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 requests (greedy and sampled) refuse UNLESS `MEMRA_SPEC_PENALTY=1`
100//! (lane/spec-exclusions-20260902, doc on [`glm5_spec_penalty_on`]): under the arm every
101//! verify round penalizes a copy of the target rows over the per-row evolving session
102//! window with the host sampler's exact arithmetic and the accept reads only that copy; the
103//! worker's admission keeps penalized requests plain while the arm is dark.
104//!
105//! FR-SPEC VOCAB MASKING (owner addition, 2026-08-30 — the house spec recipe, the q38 way):
106//! the loop consumes the existing `MEMRA_FRSPEC_TRIM` contract, no new flag. The loader
107//! already reaches this head: the trim match in `hybrid.rs` consumes the embedded head
108//! regardless of arch, `frspec_trim_own_head_name(n_trunk)` misses (glm5_next ships no
109//! private MTP lm_head) and the gather falls back to the trunk `output.weight` /
110//! `token_embd.weight` — which for glm5_next is EXACT BY CONTRACT, not merely by tying:
111//! the MTP block projects through the trunk lm_head (LANE.md; the draft gate pins
112//! `shared_head_head.is_none()` untrimmed). With a ranks artifact loaded,
113//! `mtp_head_forward_mla_cached` already projects through the gathered rows (its head is
114//! `shared_head_head.unwrap_or(trunk)`), so the DRAFT logits arrive `[n_ranks]`; this
115//! loop's seam is the remap: every draft argmax is a RANK id and maps through `d2t` back
116//! to the true vocab BEFORE it is drafted, chained (e_tok drives the embedding gather),
117//! or verified. THE VERIFY WALK STAYS FULL-VOCAB AND UNTOUCHED — a trimmed draft can only
118//! change WHICH tokens get drafted, never how they verify; that invariant is the whole
119//! design (q38's measured skipped-remap defect was 0/248 acceptance with every exactness
120//! gate green — silent, which is why the gate below makes it loud). Rank artifacts for
121//! glm5 are an INPUT DEPENDENCY: the corpus mint (SXC pools through GLM's tokenizer,
122//! per traffic class, the q38 plain-text format) is the owner's CPU-only lane; the
123//! self-trim d2t arm needs no external artifact and lands here first.
124//!
125//! PPN (lane/glm5-ppn-verify, 2026-08-30): the verify walk owns its stage split exactly as
126//! the batched decode walk does — `glm5_verify_rows_ppn` mirrors
127//! `decode_step_batch_hyper_ppn` (per-stage engine, per-stage pos_rows, ONE
128//! `[t, streams, n_embd]` boundary payload per cut; row chaining is per-LAYER through the
129//! one cache, so a straight layer-range split preserves it exactly). Rollback restores each
130//! stage's layers through that stage's engine on its stream; the MTP block, its latent
131//! plane and every draft-chain/accept-side op ride the LAST stage's engine
132//! (`glm5_head_engine` — where the loader puts the lm head and `pp::new_cache*` puts the
133//! trailing MTP plane), so the h_seed carrier never bounces devices. Gate:
134//! `glm5-spec-ppn-gate` (the tparallel battery under the split, stages=2 and 3, red-proven;
135//! the cross-device twin is the box arm). Worker admission bounds sharded placements to the
136//! GATED stage set (`glm5_sharded_placement_admits`, worker.rs) — everything else stays
137//! fail-closed by name.
138//!
139//! DRAFT SOURCE SEAM (lane/glm5-dflash-draft-src, 2026-08-30): the session's drafts come
140//! from ONE of two sources, selected at load and pinned for the session — everything from
141//! the verify walk on (accept, rollback, commit, receipts, K policy) is SHARED and
142//! source-blind, which is where the exactness invariant lives (a draft source can only
143//! move acceptance, never output):
144//!
145//!   * `NativeMtp` (existing): the embedded NextN head chains K drafts through
146//!     `mtp_head_forward_mla_cached`; requires `MEMRA_GLM5_MTP=1`.
147//!   * `Dflash2` (`MEMRA_GLM5_DFLASH=<dir-or-hf-spec>`): the pinned
148//!     incoai/GLM-5.3-Flash-DFlash2 block-diffusion drafter (owner holds WRITTEN APPROVAL
149//!     from the DFlash2 owners, 2026-08-30, for use beyond probe/eval). It REUSES the
150//!     shipped q38 DFlash2 machinery verbatim (`DflashDraft`: `ctx_features` ->
151//!     `ingest_ctx` -> `forward_round` -> `dflash2_propose_*`, mask-fill harvest, selector
152//!     walk); the ONE glm5-specific input is the drafter's measured feature contract — the
153//!     STREAM-MEAN (`hc_contract`) of the COMPLETED trunk layer output at the drafter
154//!     config's `target_layer_ids` (plan layers 5,14,24,33,42 on the real artifact), the
155//!     exact definition the probe's `MEMRA_TRACE_LAYER_ROWS` capture seam banked 0.73
156//!     acc@1 / 3.06 tokens-per-cycle against
157//!     (research/glm53-flash-bringup-20260827/dflash2-probe-20260829/RECEIPTS.md). The
158//!     features flow through [`crate::cache::HcTapSink`], a HOST sink filled by the hc
159//!     prime walks and this file's verify walk (host-staged so a ppN split needs no
160//!     cross-device tap plumbing; the drafter itself runs on the HEAD engine, where the
161//!     trunk lm_head it projects through lives). THE NATIVE MTP HEAD IS NOT LOADED for
162//!     this source (the q38 pattern — a full MoE trunk layer of VRAM back); the plan's
163//!     trailing MTP cache plane still allocates (plan-structural, ~`ctx * latent_width`
164//!     f32 per declared block — named cost, not forked). Sampled route: the drafter's
165//!     selector proposal records its true q (`DsparkDraftSample::Selector`) and the accept
166//!     rides the SAME `dspark_accept_sampled` rejection walk the q38 serve route ships,
167//!     with this session's Philox counters (`uctr` selector/accept draws, `sctr`
168//!     bonus/residual) so randomness never repeats across bursts. K is bounded by the
169//!     drafter block (K <= block_size-1 = 7): the worker clamps, the burst refuses loudly.
170//!     Selection receipt: boot logs `[glm5-spec] draft source = native-mtp` or
171//!     `[glm5-spec] draft source = dflash2 @ <sha8>`; both flags off = plain serving
172//!     (fail-closed warn). `MEMRA_GLM5_DFLASH_GATE_RED=tap-shift` is a GATE INSTRUMENT
173//!     (never a serving flag): it shifts every tap layer +1 to red-prove that a wrong
174//!     feature input collapses acceptance while the tape stays byte-identical.
175//!
176//! SERVING EXPOSURE (lane/glm5-spec-routing, 2026-08-30): `MEMRA_GLM5_SPEC` (default OFF,
177//! FLAGS.md row) is the ONE master flag — it routes `generate_spec` here for hc trunks
178//! with a loaded MTP head AND arms the worker route (`glm5_spec_capable` +
179//! `step_glm5_spec` driving [`Glm5SpecSession`] bursts). OFF = the named `refuse_hyper`
180//! refusal and zero `[glm5-spec]` log lines, byte-identical serving. The MTP_SPEC
181//! capability manifest remains deliberately UNEXTENDED — worker `mtp_spec_capable` stays
182//! false for glm5_next plans; the serving capability lives in its OWN manifest
183//! (`GLM5_SPEC`, execution_manifest.rs) whose table names exactly the glm5_next class, and
184//! a SEALED production bundle still fails closed until it banks a `glm5-spec.v1` rewrite
185//! receipt (the real-artifact qualification lane's job).
186
187use crate::Engine;
188use crate::cache::{Cache, HcTapSink};
189use crate::dflash::{DflashDraft, DflashKv, DsparkDraftSample};
190use crate::forward::argmax;
191use crate::hybrid::{HybridModel, Mixer};
192use crate::spec::SpecSampling;
193use crate::spec_phase::{
194    ProfClock, SPEC_PROF_ROUNDS, SpecFirstTokenProf, SpecPhaseNs, SpecRoundProf, SpecRoundsLog,
195    V_SEQ_ROWS, spec_prof_on,
196};
197use cudarc::driver::CudaSlice;
198
199type Res<T> = Result<T, Box<dyn std::error::Error>>;
200/// Round-cadence commit hook of `glm5_spec_session_burst_streamed` (lane/b200-spec-ttft-
201/// 20260902): called with every newly committed slice of a burst, in order, disjoint.
202pub type CommitHook<'a> = &'a mut dyn FnMut(&[u32]);
203
204/// `MEMRA_GLM5_SPEC=1` routes `generate_spec` to the glm5 T-parallel loop. Default OFF:
205/// unset/0 keeps the standing `refuse_hyper` refusal, so serving is byte-identical to the
206/// pre-lane binary. Read once (worker chunk policies read their flags the same way).
207pub fn glm5_spec_on() -> bool {
208    use std::sync::OnceLock;
209    static ON: OnceLock<bool> = OnceLock::new();
210    *ON.get_or_init(|| std::env::var("MEMRA_GLM5_SPEC").as_deref() == Ok("1"))
211}
212
213// Per-burst phase attribution moved to `crate::spec_phase` (lane/glm5-extract-general):
214// the draft/verify/accept/roll/maint split is spec-family-generic. This loop consumes
215// `MEMRA_SPEC_TRACE` (glm5 alias `MEMRA_GLM5_SPEC_TRACE` stays honored) and passes its
216// own `[glm5-phase]` / `[glm5-phase-v]` tags so every banked receipt keeps its shape.
217
218/// `MEMRA_GLM5_VERIFY_BATCH` (default ON, lane/glm5-verify-batch): the per-LAYER batched
219/// mixer walk — one t=K+1 KDA call per layer (projections/conv/gates batched through the
220/// decode-exact classes, the recurrence sequential INSIDE one scan launch) and one
221/// t=K+1 rows-exact MLA call per layer, replacing the per-row mixer loop. `0` restores
222/// the per-row walk byte-for-byte — the rollback seam. Deliberate default (new-flags
223/// law): the walk only exists behind `MEMRA_GLM5_SPEC` (default OFF in prod), per-row
224/// byte identity is bit-gated on the rig (`glm5_tparallel_verify_gpu` +
225/// `glm5_verify_batch_gpu`), and the box re-battery A/Bs this seam in one build. Read
226/// PER CALL (the `MEMRA_KDA_FUSED_PROJ` per-call precedent) so gates drive both arms in
227/// one process; one env read per verify walk.
228pub fn glm5_verify_batch_on() -> bool {
229    std::env::var("MEMRA_GLM5_VERIFY_BATCH").as_deref() != Ok("0")
230}
231
232/// `MEMRA_GLM5_DRAFT_PRIME_V2` (default OFF, lane/spec-route-depth-20260902): the DFlash2
233/// drafter's prompt prime rides the target's own chunk schedule. OFF (the pre-lane
234/// literal): the target prime fills a whole-prompt HOST tap sink (`[prompt, 5 x 4096]`
235/// f32, 21 GB at 256k) through synchronous pageable DtoHs, and round 1 re-uploads it in
236/// 256-row pageable HtoD chunks before the fc + k/v ingest. ON: each prime chunk
237/// (`hyper_prime_ranges`, 4096 rows) arms a DEVICE-staged chunk sink, and right after that
238/// chunk's trunk walk the drafter ingests it: slot DtoH into a pinned cacheable buffer,
239/// host interleave, ASYNC HtoD from pinned, fc + k/v ingest at the chunk width. The trunk
240/// program is byte-identical (the same per-range `prime_cache` calls with the same
241/// `queued_after`); the drafter's KV rows are the same GEMM class at a different M (rig
242/// gate 15 pins bit-identity at equal chunking and tape identity always). Read per session
243/// creation so a gate can flip it in-process. Rollback seam: unset.
244pub fn glm5_draft_prime_v2_on() -> bool {
245    std::env::var("MEMRA_GLM5_DRAFT_PRIME_V2").as_deref() == Ok("1")
246}
247
248/// `MEMRA_GLM5_DRAFT_TAPS_DEVICE` (DEFAULT OFF; `=1` arms it; lane/spec-route-depth-20260902):
249/// the DEVICE-RESIDENT drafter
250/// prime. Boot D on the 2x B200 pair (this arm) against boot C (host taps), TTFT s at
251/// 4k / 42k / 128k / 256k: 4.89 / 16.84 / 56.62 / 129.42 vs 5.32 / 21.23 / 72.90 /
252/// 162.26, drafter prime ms 17.9 / 59.1 / 180.7 / 362.4 vs 91.4 / 561.6 / 4430.8 /
253/// 8868.4, steady rounds and decode-after unchanged (FLAGS row, receipt path there). The trunk prime stays the ONE whole-prompt call (the
254/// chunked arm's per-range calls re-entered the PP-2 microchunk geometry and made the trunk
255/// prime itself far slower — boot B); the tap rows never leave the device: the walk stages
256/// each range's five contracted tap planes on the writing stage's device (a chunk-sized
257/// ring, `HcTapSink::new_device_staged_at`), and at the range boundary the prime's own loop
258/// hands the range to `glm5_taps_range_done`, which interleaves the planes into the fc
259/// layout with 2D device copies (a peer copy first for planes written on another stage),
260/// runs `ctx_features` + `ingest_ctx` at the range width (4096 rows: the batched GEMM
261/// class), and appends to the drafter KV. No DtoH in the prime, no HtoD in the drafter
262/// prime; the eager arm's 7.5 s of pageable HtoD at 256k (boot B's split: h2d 7505.8 of
263/// 8870 ms) goes away outright, and its `tap_dtoh` with it. Read per session creation
264/// (a gate flips it in-process). The host-tap arms (`MEMRA_GLM5_DRAFT_PRIME_V2`,
265/// `MEMRA_GLM5_DRAFT_PRIME_LAZY`) are the default and are reachable whenever this is unset.
266///
267/// DEFAULT OFF DELIBERATELY, and the flip is blocked on two named defects that a review of
268/// the ON default found. Both are properties of the device-staged ring, so neither can bite
269/// while this is unset. FIRST, the `MEMRA_B200_PRIME_V2` arm-2 branch
270/// (`prime_cache_hyper_pp2_pipelined`) never calls
271/// `glm5_taps_range_begin`/`glm5_taps_range_done`, and hands each stage-chunk an ABSOLUTE
272/// base, so from the second range on `glm5_hc_tap_slot` slices past the `sink.t * h` slot
273/// buffer: `slice_mut` panics and takes the worker with it. SECOND, the same-`ordinal`
274/// ingest reads the ring `buf` from the head engine's stream while the next range rewrites
275/// it on the stage engine's stream, with cudarc implicit event tracking disabled in this
276/// repo, so the tap planes can tear into silently wrong drafter ctx KV; verify still
277/// arbitrates the served tape, so the tape-identity gates stay green and would NOT catch
278/// it. Flipping the default needs both fixed AND a gate that fails on each before it
279/// passes.
280pub fn glm5_draft_taps_device_on() -> bool {
281    std::env::var("MEMRA_GLM5_DRAFT_TAPS_DEVICE").as_deref() == Ok("1")
282}
283
284/// The device-resident drafter prime's in-flight state (doc on
285/// [`glm5_draft_taps_device_on`]), carried type-erased on the prime's tap sink
286/// (`HcTapSink::ingest_state`) so the prime's range loop can hand it each range.
287struct Glm5DraftPrimeInflight {
288    kv: DflashKv,
289    taps: Vec<usize>,
290    n_embd: usize,
291    /// Rows the sink ring covers (the largest range the schedule emits).
292    ring: usize,
293    /// Head-device staging for planes written on another stage (lazily sized `ring x h`).
294    stage: Vec<Option<CudaSlice<f32>>>,
295    /// The interleaved `[ring, n_taps * h]` fc input on the head device (lazy).
296    rows_dev: Option<CudaSlice<f32>>,
297    prof_on: bool,
298    copy_ms: f64,
299    feat_ms: f64,
300    kv_ms: f64,
301    chunks: usize,
302}
303
304/// `MEMRA_GLM5_DRAFT_PRIME_LAZY` (default OFF, lane/spec-route-depth-20260902): `=1`
305/// restores the pre-lane placement of the EAGER arm's drafter ingest — inside round 1 of
306/// the first burst, i.e. AFTER the prime's anchor token has been emitted under the
307/// round-cadence door. Default (unset): the ingest runs at session creation, before the
308/// session is returned and before any token is emitted. WHY THE FLIP: the 2x B200 pair's
309/// boot A (`MEMRA_SPEC_PROF=1`, main + PR #101) measured the round-0 wall at 0.63 / 4.5 /
310/// 8.9 s for 42k / 128k / 256k prompts (about 35 us per prompt token, the eager ingest)
311/// while every later round sits at 55-64 ms; with the anchor already streamed, that one
312/// round lands INSIDE decode, which is the bimodal "15.4 vs 32.1 tok/s" the pair saw
313/// (256 tokens over 8.9 s + 255 rounds/37 tok/s = 16 tok/s; without the stall, 37). The
314/// work is identical either way (same rows, same GEMMs, same KV bytes: the KV is
315/// position-addressed and nothing touches it between creation and round 1); only WHEN it
316/// runs moves, from the second token's latency to TTFT. Read per session creation.
317pub fn glm5_draft_prime_lazy_on() -> bool {
318    std::env::var("MEMRA_GLM5_DRAFT_PRIME_LAZY").as_deref() == Ok("1")
319}
320
321/// Drafter ctx KV bytes at `cap` rows (the `DflashKv::new` geometry), for the profile.
322fn dflash_kv_bytes(cfg: &crate::dflash::DflashCfg, cap: usize) -> usize {
323    2 * cfg.n_layer * (cap + cfg.block_size) * cfg.n_kv * cfg.head_dim * std::mem::size_of::<f32>()
324}
325
326/// Pinned host buffer viewed as `n` f32s (page-aligned by construction).
327fn pinned_f32(buf: &crate::PinnedHostBuf, n: usize) -> &[f32] {
328    debug_assert!(n * std::mem::size_of::<f32>() <= buf.len());
329    // SAFETY: malloc_host allocations are page-aligned and the length is checked above.
330    unsafe { std::slice::from_raw_parts(buf.as_slice().as_ptr() as *const f32, n) }
331}
332
333fn pinned_f32_mut(buf: &mut crate::PinnedHostBuf, n: usize) -> &mut [f32] {
334    debug_assert!(n * std::mem::size_of::<f32>() <= buf.len());
335    // SAFETY: as above; exclusive borrow of the buffer.
336    unsafe { std::slice::from_raw_parts_mut(buf.as_mut_slice().as_mut_ptr() as *mut f32, n) }
337}
338
339/// `MEMRA_GLM5_SPEC_PREFIX` (default OFF, lane/glm5-prefix-latent2 2026-09-01): the glm5
340/// spec x prefix-cache interplay — BOTH sides, one flag (the `MEMRA_DSPARK_PREFIX_RESTORE`
341/// precedent): (capture) DFlash2-source sessions take a prompt-boundary capture at creation
342/// for the worker's deferred prefix publication, and (restore) the worker may convert a
343/// prefix hit into `glm5_spec_session_from_restored` instead of demoting it to the plain
344/// route. Requires `MEMRA_PREFIX_LATENT=1` too — a capture the worker's publisher would
345/// refuse (latent entries need the plane flag) is pure waste, so this predicate ANDs both
346/// env reads. DEFAULT OFF BY DESIGN (new-flags law): the restored-session program is
347/// unmeasured until the box battery banks restored-vs-cold byte identity on the
348/// continuation; unset restores the pre-lane posture exactly (spec sessions never capture,
349/// hits demote to plain). Read once per process.
350pub fn glm5_spec_prefix_on() -> bool {
351    use std::sync::OnceLock;
352    static ON: OnceLock<bool> = OnceLock::new();
353    *ON.get_or_init(|| {
354        let sp = std::env::var("MEMRA_GLM5_SPEC_PREFIX").as_deref() == Ok("1");
355        let pl = std::env::var("MEMRA_PREFIX_LATENT").as_deref() == Ok("1");
356        if sp && !pl {
357            // A mis-built recipe would otherwise only surface through the battery's
358            // receipt gates (PR #96 review round 2, minor) — say it at first read.
359            eprintln!(
360                "[glm5-spec] MEMRA_GLM5_SPEC_PREFIX=1 is INERT: it requires \
361                 MEMRA_PREFIX_LATENT=1 (latent entries could never publish without it)"
362            );
363        }
364        sp && pl
365    })
366}
367
368/// `MEMRA_GLM5_SPEC_FULLCOVER` (default OFF, lane/glm5-fullcover-spec-route 2026-09-02):
369/// admit the glm5 spec route on a FULL-COVER prefix hit, i.e. a hit whose restored prefix
370/// already covers the whole prompt so there is no suffix left to prime.
371///
372/// WHY IT EXISTS (memra#74): the parent lane refused an empty suffix on the premise that
373/// "the plain boundary-logits resume is faster than any prime". That is true of PREFILL and
374/// says nothing about DECODE: with the drafter left un-armed the whole generation then ran
375/// at plain speed. Measured on the live glm5 box 2026-09-02, same minute and vantage,
376/// vendor-default sampled, 67-token prompt, 512 max_tokens: repeated (full-cover hit)
377/// 29.46 s wall median / 30.8 tok/s decode vs fresh-nonce (cold, route=spec) 57.75 / 69.7.
378/// A cache hit cost the customer half the decode speed on that shape.
379///
380/// WHY IT IS WELL-FORMED: with an empty suffix the restored session's state is exactly a
381/// cold session's at the same boundary: trunk cache at `fed.len()`, drafter ctx KV at
382/// `fed.len()` (`DflashKv::from_tail`, the same rows a cold prime's tap ingest would have
383/// produced), no pending tap rows, and the anchor drawn from the ENTRY's boundary logits by
384/// the same rule the cold burst applies to its own first token. The Dflash2 round invariant
385/// `kv.len == cache.pos` therefore holds at round 1 with `pending` empty. This is the same
386/// full-cover shape the MTP restore has served since lane/spec-cache
387/// (`spec_restore_refusal`: a full-cover hit is admitted whenever the entry carries its
388/// boundary hidden + logits).
389///
390/// DEFAULT OFF BY DESIGN (new-flags law): the arm has NO GPU receipt yet. Byte identity of
391/// the full-cover restored tape against plain decode is GATE 13 of
392/// `glm5_dflash_session_gpu` (`MEMRA_GLM5_SPEC_FULLCOVER=1`), and the serving win needs the
393/// repeated-prompt cell on the glm5 box. Unset restores the pre-lane posture exactly
394/// (full-cover hits serve plain, now with `reason=full-cover-hit` on the route line).
395/// Read PER CALL (the `MEMRA_KDA_FUSED_PROJ` / `MEMRA_GLM5_VERIFY_BATCH` precedent) so one
396/// gate process can drive both arms. Rollback seam: unset.
397pub fn glm5_spec_fullcover_on() -> bool {
398    std::env::var("MEMRA_GLM5_SPEC_FULLCOVER").as_deref() == Ok("1")
399}
400
401/// `MEMRA_GLM5_SPEC_TP` (default OFF, lane/glm5-composition 2026-09-01): admit glm5 spec
402/// SESSIONS on a `MEMRA_GLM5_TP`-armed model. DEFAULT OFF BY DESIGN (new-flags law): the
403/// composition's verify/rollback wiring is rig-gated for correctness (per-rank KDA
404/// snapshot/replay, per-replica MLA latent truncation — `glm5-tp-gate` arms S*), but it has
405/// ZERO real-artifact receipts and the TP serving wiring is still the named box increment;
406/// an unmeasured composition does not default ON. `=1` lifts ONLY the session co-refusal —
407/// every other admission law holds (draft source required, batched verify walk required:
408/// the per-row rollback seam carries no TP arm and refuses by name). Read per session
409/// creation. Rollback seam: unset (the co-refusal is restored verbatim).
410pub fn glm5_spec_tp_on() -> bool {
411    std::env::var("MEMRA_GLM5_SPEC_TP").as_deref() == Ok("1")
412}
413
414/// `MEMRA_SPEC_PENALTY` (default OFF, lane/spec-exclusions-20260902): admit PENALIZED
415/// requests (repetition / frequency / presence, greedy AND sampled) to the glm5 spec route.
416/// Pre-lane every penalized request served plain (`reason=penalized`), which on agent
417/// harnesses that ship a vendor-default `presence_penalty` was the whole route's 1.48x lost
418/// on the most common request shape there is. THE ARM: every verify round penalizes a COPY of
419/// the target's verify rows on device over the per-row evolving window (`penalize_logits_rows_inc`
420/// — row r sees the session window ++ the drafts before r, exactly the plain sampler's history
421/// at that position), and the accept reads only that copy: greedy argmaxes it, the sampled
422/// walk takes its p from it, the bonus and the residual draw from it; the anchor is drawn the
423/// same way from the prime's boundary row. The DRAFT stays unpenalized (it only proposes;
424/// rejection sampling is unbiased for any proposal). NUMERIC CLASS: the device pass is the
425/// host sampler's arithmetic bit for bit (`keskar_penalize_rn`, gated by
426/// `gpu_device_penalties_are_bit_identical_to_the_host_sampler`), so greedy+penalties keeps
427/// the spec-vs-plain tape identity (gate 15) and sampled+penalties keeps the route's
428/// distribution-exact claim: the accepted stream equals the plain penalized sampler's
429/// distribution. DEFAULT OFF BY DESIGN (new-flags law): the arm has rig receipts only; the
430/// box battery (sampled vendor-default + penalties, 8-turn twin) flips it. Unset = the
431/// pre-lane refusal verbatim (the worker keeps penalized requests plain, by name). Read per
432/// session creation (the `MEMRA_GLM5_SPEC_FULLCOVER` precedent) so one gate process drives
433/// both arms.
434pub fn glm5_spec_penalty_on() -> bool {
435    std::env::var("MEMRA_SPEC_PENALTY").as_deref() == Ok("1")
436}
437
438/// `MEMRA_SPEC_WARM` (default OFF, lane/spec-exclusions-20260902): admit a prefix-hit
439/// CARRIER whose entry carries NO usable drafter tail (`reason=no-drafter-tail`: the entry
440/// was published by a plain session — a K-shed turn, a constrained or vision turn — or its
441/// tail did not cover the drafter window) by re-arming the drafter COLD at the restored
442/// boundary (`DflashKv::new_cold_at`): the trunk restores, the drafter starts with an empty
443/// context at the right absolute position and fills it from the suffix prime's taps and every
444/// committed round. Rows below the floor are never attended (the clipped round attention's
445/// floor arm), so the drafter sees exactly the program a shorter prompt runs. Can only move
446/// ACCEPTANCE, never output (verify arbitrates; gate 18 pins the restored tape against plain
447/// decode). Lives UNDER `MEMRA_GLM5_SPEC_PREFIX` (x `MEMRA_PREFIX_LATENT`): without the restore
448/// door there is no carrier to warm. DEFAULT OFF BY DESIGN: the acceptance cost of the empty
449/// context on the first rounds is unmeasured on the real artifact. Read per request.
450pub fn glm5_spec_warm_on() -> bool {
451    std::env::var("MEMRA_SPEC_WARM").as_deref() == Ok("1")
452}
453
454/// The penalty arm's per-session config (doc on [`glm5_spec_penalty_on`]): the request's
455/// Keskar coefficients and window, INDEPENDENT of the sampling regime — a greedy request
456/// with penalties admits too, which is why this is not folded into `Option<SpecSampling>`
457/// (whose `None` means greedy everywhere else on the route).
458#[derive(Clone, Copy, Debug)]
459pub struct Glm5Penalty {
460    last_n: usize,
461    rep: f32,
462    freq: f32,
463    present: f32,
464}
465
466impl Glm5Penalty {
467    /// `Some` iff the request carries a non-identity penalty with an armed window — THE
468    /// `pen_on` predicate (one definition, `SpecSampling::pen_on`).
469    pub fn of(sp: &SpecSampling) -> Option<Self> {
470        sp.pen_on().then_some(Self {
471            last_n: sp.penalty_last_n,
472            rep: sp.penalty_repeat,
473            freq: sp.penalty_freq,
474            present: sp.penalty_present,
475        })
476    }
477    /// The device window: `penalty_last_n` under the `PEN_WINDOW_MAX` cost bound — the same
478    /// trim the dspark accept walk and the boundary draw apply.
479    fn win(&self) -> usize {
480        self.last_n.min(crate::spec::PEN_WINDOW_MAX)
481    }
482}
483
484/// Print-once latch for the penalty arm's engagement line (one per process).
485static PENALTY_ARM_ANNOUNCED: std::sync::atomic::AtomicBool =
486    std::sync::atomic::AtomicBool::new(false);
487
488/// The penalty admission law, shared by the cold and restored session constructors: a
489/// penalized request opens a session ONLY under `MEMRA_SPEC_PENALTY=1`; otherwise it refuses
490/// loudly (worker admission owns the exclusion and keeps such requests plain — silently
491/// dropping a request's penalties is the failure class the refusal prevents). Logs the
492/// engagement receipt once per process on the first penalized session that opens.
493fn glm5_penalty_admit(sampling: Option<&SpecSampling>) -> Res<Option<Glm5Penalty>> {
494    let Some(pen) = sampling.and_then(Glm5Penalty::of) else {
495        return Ok(None);
496    };
497    if !glm5_spec_penalty_on() {
498        return Err(
499            "glm5 spec penalty arm is DARK (MEMRA_SPEC_PENALTY unset): penalized requests serve \
500             on the plain path (worker admission owns the exclusion; silently dropping the \
501             request's penalties is the failure class this refusal prevents)"
502                .into(),
503        );
504    }
505    if !PENALTY_ARM_ANNOUNCED.swap(true, std::sync::atomic::Ordering::AcqRel) {
506        eprintln!(
507            "[glm5-spec] penalty arm ENGAGED (MEMRA_SPEC_PENALTY=1): verify rows penalized on \
508             device over the session window (rep={} freq={} present={} last_n={}); printed \
509             once per process, every penalized request carries penalized=1 on its route line",
510            pen.rep, pen.freq, pen.present, pen.last_n,
511        );
512    }
513    Ok(Some(pen))
514}
515
516/// Seed the session's penalty window from its committed prompt (`spec::pen_window_seed`, one
517/// definition of "the window" across every spec route). Empty with penalties off.
518fn glm5_pen_window_seed(pen: Option<&Glm5Penalty>, committed: &[u32]) -> Vec<u32> {
519    match pen {
520        Some(p) => crate::spec::pen_window_seed(&[], committed, p.last_n),
521        None => Vec::new(),
522    }
523}
524
525/// The session's FIRST token from a host boundary-logits row (the prime's, or the entry's on
526/// a full-cover restore) — the plain route's own first-token rule, regime by regime:
527///   * sampled: `sample_boundary_token` (penalize over the window, filter, Gumbel at the
528///     session's Philox counter) — with the window it was always handed empty before;
529///   * greedy + penalties: the plain sampler's `sample` is penalize-then-argmax, so the row
530///     goes to the device, takes the one penalty pass (`penalize_logits`, the same bits the
531///     host produces), and the device argmax (host tie-break contract) picks;
532///   * greedy: the host argmax, byte for byte the pre-lane literal.
533#[allow(clippy::too_many_arguments)]
534// allow: the three regimes' inputs, listed rather than bundled, so the OFF arm reads as the
535// pre-lane literal at the call site
536fn glm5_anchor(
537    eh: &Engine,
538    logits: &[f32],
539    sampling: Option<&SpecSampling>,
540    pen: Option<&Glm5Penalty>,
541    pen_hist: &[u32],
542    sctr: &mut u32,
543    site: &str,
544) -> Res<u32> {
545    match (sampling, pen) {
546        (Some(sp), _) => crate::spec::sample_boundary_token(eh, logits, sp, pen_hist, sctr, site),
547        (None, Some(p)) => {
548            let n = logits.len();
549            let mut col = eh.htod(logits)?;
550            let w0 = pen_hist.len().saturating_sub(p.win());
551            let hist = &pen_hist[w0..];
552            if !hist.is_empty() {
553                let hd = eh.htod_u32_v(hist)?;
554                eh.penalize_logits(&mut col, &hd, hist.len(), p.rep, p.freq, p.present, n)?;
555            }
556            let td = eh.argmax_token_device(&col, n)?;
557            crate::spec::guard_vocab_token(
558                eh.dtoh_u32_one(&td)?,
559                n,
560                &format!("glm5 penalized greedy anchor ({site})"),
561            )
562        }
563        (None, None) => Ok(argmax(logits) as u32),
564    }
565}
566
567/// `MEMRA_SPEC_PMIN`, honored by the glm5 loop (loop-port 2 — the step37 shipping family,
568/// `MEMRA_SPEC_PMIN=0.5 MEMRA_SPEC_PMIN0=1` is what step37 serves; NO new flag): stop the
569/// draft chain early when the drafter's confidence in its own pick drops below p_min.
570/// Native chain: p = the head's softmax confidence in its pick (the spec.rs `g_p`
571/// statistic, `prob_of_token_device`). DFlash2: q = the selector's recorded per-slot
572/// candidate-set confidence (`q_chosen`; T=1 twin on the greedy walk) — the owner's
573/// "take only high confidence offers" tau-slot form, truncated PRE-verify.
574/// Unset/0 = OFF (today's rounds, byte-identical). The VALUE is a per-model measurement
575/// (spec.rs bank: q27 PMIN=0.3 was -1.9% on one pack; step37 ships 0.5) — the box-B tau
576/// ladder prices glm5's.
577pub(crate) fn glm5_pmin() -> f32 {
578    use std::sync::OnceLock;
579    static P: OnceLock<f32> = OnceLock::new();
580    *P.get_or_init(|| {
581        std::env::var("MEMRA_SPEC_PMIN")
582            .ok()
583            .and_then(|v| v.parse().ok())
584            .unwrap_or(0.0)
585    })
586}
587
588/// `MEMRA_SPEC_PMIN0=1` (llama.cpp's draft gating, vendored via spec.rs): the p-min gate
589/// applies at slot 0 too, so a low-confidence round drafts NOTHING and the verify batch is
590/// just the anchor row — m=1 = a plain decode step. Always legal for glm5 (the anchor row
591/// exists every round). "llama's 35B win rides exactly this — draft acceptance 76% at mean
592/// len 2.5 because unpredictable stretches never pay draft+verify overhead" (spec.rs).
593pub(crate) fn glm5_pmin0() -> bool {
594    use std::sync::OnceLock;
595    static P: OnceLock<bool> = OnceLock::new();
596    *P.get_or_init(|| std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1"))
597}
598
599/// MEMRA_SPEC_PMIN break semantics — hoisted to the shared K-policy surface
600/// ([`crate::spec::spec_conf_keep`], lane/glm5-extract-general); re-exported here so the
601/// glm5 gates and call sites keep their name.
602pub use crate::spec::spec_conf_keep as glm5_conf_keep;
603
604/// The loaded DFlash2 drafter (module doc, DRAFT SOURCE SEAM): the model-level half of the
605/// `Dflash2` draft source — weights loaded ONCE per model on the head engine (`hybrid.rs`,
606/// `MEMRA_GLM5_DFLASH`); per-session state lives in [`Glm5DraftState`].
607///
608/// HOISTED to the general seam ([`crate::dflash::DflashDrafter`], lane/glm5-extract2): the
609/// holder is `{ drafter weights, byte-identity pin }` with nothing glm5 in it. Re-exported
610/// here under its old name so glm5's call sites and gates keep the name they were written
611/// against, exactly as `glm5_conf_keep` does above.
612pub use crate::dflash::DflashDrafter as Glm5DflashDrafter;
613
614/// Process-wide count of verify rounds whose drafts came through a RANK-TRIMMED draft head
615/// (either source; lane/frspec-dflash2-20260902). Gates read it beside the per-session
616/// [`Glm5SpecSession::rank_trimmed_rounds`]; the worker's `[glm5-acc]` line carries the
617/// session's own count.
618static GLM5_RANK_TRIMMED_DRAFT_ROUNDS: std::sync::atomic::AtomicU64 =
619    std::sync::atomic::AtomicU64::new(0);
620
621/// Read [`GLM5_RANK_TRIMMED_DRAFT_ROUNDS`] (doc on the static).
622pub fn glm5_rank_trimmed_draft_rounds() -> u64 {
623    GLM5_RANK_TRIMMED_DRAFT_ROUNDS.load(std::sync::atomic::Ordering::Relaxed)
624}
625
626/// Per-session draft-source state (module doc, DRAFT SOURCE SEAM). Selected at session
627/// creation from the model's loaded sources and pinned for the session's lifetime.
628pub(crate) enum Glm5DraftState {
629    /// Embedded NextN head: state = the MTP latent plane + `Glm5SpecSession::pending`
630    /// (token, h_seed) pairs — the pre-seam program, byte-identical.
631    NativeMtp,
632    /// DFlash2 block-diffusion drafter: state = the drafter's own ctx-feature KV cache
633    /// plus host feature rows not yet ingested. Invariant at every round boundary:
634    /// `kv.len + pending.len()/(taps.len()*n_embd) == committed.len()` — the drafter's
635    /// context is exactly the committed tokens (the probe's `F_feat[new_lo:start]` walk).
636    Dflash2 {
637        kv: DflashKv,
638        /// Committed-position feature rows awaiting ingest, `[n, n_taps*n_embd]` host
639        /// (the prompt's prime taps at session start; each round's kept verify taps after).
640        pending: Vec<f32>,
641        /// Resolved tap layers (drafter config `target_layer_ids`, red-arm shift applied).
642        taps: Vec<usize>,
643    },
644}
645
646/// The retained q side of one round's draft chain — what the sampled accept walk consumes.
647/// Greedy rounds carry `None` (the accept is the byte-deterministic prefix walk).
648enum Glm5DraftQ {
649    None,
650    /// Native MTP chain: per-slot retained draft logits (rank space under a trim) + the
651    /// filtered stats of the distribution each draft was drawn from.
652    Mtp {
653        draft_idx: Vec<u32>,
654        draft_logits: Vec<CudaSlice<f32>>,
655        draft_stats: Vec<(f32, f32, f32)>,
656    },
657    /// DFlash2 selector proposal (the recorded candidate-set q) + the retained draft-logit
658    /// rows `dl` (`dspark_accept_sampled`'s buffer contract; unread on the Selector q path).
659    Selector {
660        prop: DsparkDraftSample,
661        dl: CudaSlice<f32>,
662    },
663}
664
665/// Resolve the drafter's tap layers against the trunk: the drafter config's
666/// `target_layer_ids` are memra PLAN layer indices whose COMPLETED output feeds the fc
667/// (the probe's capture convention: `MEMRA_TRACE_LAYER_ROWS_LAYERS=5,14,24,33,42` == the
668/// drafter's own `target_layer_ids`, asserted 1:1 in `score_dflash2.py`).
669/// `MEMRA_GLM5_DFLASH_GATE_RED=tap-shift` is the RED-ARM INSTRUMENT: every tap moves +1
670/// layer — deliberately wrong features whose acceptance collapse the gate asserts while
671/// the output tape stays byte-identical. Unknown values refuse loudly.
672///
673/// The resolution itself is the general seam ([`crate::dflash::resolve_tap_layers`],
674/// lane/glm5-extract2); what stays here is glm5's OWN red arm — the gate instrument reads its
675/// env, prints its `[glm5-spec]` tag, and hands the shift in as a parameter. Error bytes are
676/// unchanged ("glm5 DFlash2" is the `what` label).
677fn glm5_dflash_tap_layers(draft: &DflashDraft, n_trunk: usize) -> Res<Vec<usize>> {
678    let shift = match std::env::var("MEMRA_GLM5_DFLASH_GATE_RED").ok().as_deref() {
679        Some("tap-shift") => {
680            eprintln!(
681                "[glm5-spec] RED-ARM tap-shift: drafter tap layers shifted +1 (gate \
682                 instrument, never a serving flag)"
683            );
684            1
685        }
686        Some("") | None => 0,
687        Some(other) => {
688            return Err(format!(
689                "MEMRA_GLM5_DFLASH_GATE_RED={other:?}: unknown red arm (want tap-shift)"
690            )
691            .into());
692        }
693    };
694    Ok(crate::dflash::resolve_tap_layers(
695        &draft.cfg.target_layer_ids,
696        n_trunk,
697        shift,
698        "glm5 DFlash2",
699    )?)
700}
701
702/// Pre-round state checkpoint for one glm5 verify round. Captured by `glm5_verify_rows`
703/// BEFORE any row runs; consumed by `glm5_verify_rollback`.
704///
705/// Covers exactly the state planes a glm5_next trunk mutates in a verify round:
706/// - `latent_len`: per-layer MLA latent length at round start (rollback = truncate).
707/// - KDA state (loop-port 3, the module doc's GdnStash/ReplaySSM diet LANDED): the old
708///   per-row (conv, ssm) column clones — 4 MiB x 34 layers per COLUMN, ~0.95 GiB of
709///   transient at K=7 — are replaced by
710///   * `kda_ssm_snap[il]`: ONE recurrent-state clone per layer per round (the state
711///     BEFORE row 0),
712///   * `kda_scan_stash[il][r]`: row r's scan-input buffers, STOLEN from the step (zero
713///     copies, ~160 KB/row/layer — `kda::KdaScanInputs`), rows `0..t-1` except the last
714///     (`keep == t` needs no restore, so row `t-1` is never a replay target),
715///   * `kda_conv_cols[il][r]`: the conv ring stays PER-ROW CLONED (288 KiB, 1.4% of the
716///     ssm plane it rode beside — not worth a replay arm).
717///
718///   Partial-accept rollback REPLAYS rows `0..keep` from the snapshot
719///   (`kda::kda_scan_replay`): each replay is the original t=1 scan launch re-issued over
720///   the very buffers that row consumed, so the rebuilt state is byte-identical to the
721///   clone it replaces by construction. Under a ppN split every clone/stash lives on its
722///   layer's OWNING stage engine; rollback restores through the same per-stage seam.
723/// - `pos`: `cache.pos` at round start.
724///
725/// glm5_next has no Full/Linear trunk mixers (the walk refuses them by name), so `kv`,
726/// `tp_kv` and GDN stashes have no arm here — growing one is a deliberate extension with
727/// its own gate, not a silent default.
728pub struct Glm5VerifyCkpt {
729    pos: usize,
730    latent_len: Vec<Option<usize>>,
731    /// Per-row conv-ring clones, rows `0..t-1` except the last (doc above). PER-ROW walk
732    /// only (`MEMRA_GLM5_VERIFY_BATCH=0`); the batched walk fills `kda_rows` instead.
733    kda_conv_cols: Vec<Option<Vec<CudaSlice<f32>>>>,
734    /// The recurrent state BEFORE row 0, one clone per KDA layer per round (doc above).
735    /// BOTH walks fill this — it is the batched replay's scan base too.
736    kda_ssm_snap: Vec<Option<CudaSlice<f32>>>,
737    /// Stolen per-row scan inputs, rows `0..t-1` except the last (doc above). PER-ROW
738    /// walk only.
739    kda_scan_stash: Vec<Option<Vec<crate::kda::KdaScanInputs>>>,
740    /// BATCHED walk (lane/glm5-verify-batch): one [`crate::kda::KdaRowsStash`] per KDA
741    /// layer per round — ring snapshot + stolen raw conv rows + stolen batched scan
742    /// inputs; rollback re-rolls the ring and replays the scan ONCE at T=keep.
743    kda_rows: Vec<Option<crate::kda::KdaRowsStash>>,
744    /// glm5 TP composition (lane/glm5-composition): per-rank rollback material of each
745    /// SHARDED KDA layer's batched verify call — the rank-indexed twin of
746    /// (`kda_ssm_snap`, `kda_rows`), restored through each rank's own engine. `None` on
747    /// every unsharded layer.
748    kda_tp: Vec<Option<crate::glm5_tp::Glm5TpKdaVerifyStash>>,
749    /// Row count of the walk that filled this ckpt; rollback validates `keep` against it.
750    rows: usize,
751}
752
753impl Glm5VerifyCkpt {
754    /// GATE RECEIPT (wiring anchor, not a serving surface): how many KDA layers filled
755    /// the BATCHED rows stash vs the PER-ROW column stash — the flag A/B gate asserts
756    /// the arm it set actually ran (wiring-assertions-match-prose law: anchor on the
757    /// invocation's artifact, never the log prose).
758    pub fn kda_stash_kinds(&self) -> (usize, usize) {
759        (
760            self.kda_rows.iter().filter(|s| s.is_some()).count(),
761            self.kda_conv_cols.iter().filter(|s| s.is_some()).count(),
762        )
763    }
764}
765
766/// Position buffers for one verify walk range (per stage engine under a split — the
767/// per-stage pos_d law): `all` = the `[t]` vector the BATCHED per-layer mixer calls
768/// consume; `rows` = the per-row single-position buffers of the per-row arm, built only
769/// when that arm can run (flag off) — the batched arm never reads them.
770struct Glm5VerifyPos {
771    pos0: usize,
772    t: usize,
773    all: CudaSlice<i32>,
774    rows: Vec<CudaSlice<i32>>,
775}
776
777impl Glm5VerifyPos {
778    fn new(e: &Engine, pos0: usize, t: usize) -> Res<Self> {
779        let v: Vec<i32> = (0..t as i32).map(|r| pos0 as i32 + r).collect();
780        let all = e.htod_i32(&v)?;
781        let rows = if glm5_verify_batch_on() && t > 1 {
782            Vec::new()
783        } else {
784            (0..t)
785                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
786                .collect::<Result<_, _>>()?
787        };
788        Ok(Self { pos0, t, all, rows })
789    }
790}
791
792impl HybridModel {
793    /// THE T-PARALLEL VERIFY WALK: score `tokens` (row 0 = the last committed token, rows
794    /// 1..t = the K drafted tokens) in ONE forward over the hc trunk at positions
795    /// `cache.pos .. cache.pos + t`, in the batched-decode kernel classes (module doc).
796    ///
797    /// Returns `(logits [t, n_vocab] device, collapsed [t, n_embd] device, ckpt)`:
798    /// `logits` row r is bit-identical to the plain `decode_step_hyper` logits after
799    /// consuming `tokens[r]` at that position (the gate's bar); `collapsed` row r is the
800    /// pre-output_norm hidden — the MTP `h_seed` for position `cache.pos + r`.
801    ///
802    /// State effects: every trunk MLA plane appends `t` rows; every trunk KDA state
803    /// advances `t` steps (per-step columns stashed in the ckpt); `cache.pos` is NOT moved
804    /// (rollback owns it). The MTP block's plane (il = n_trunk) is untouched.
805    pub fn glm5_verify_rows(
806        &self,
807        e: &Engine,
808        tokens: &[u32],
809        cache: &mut Cache,
810    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
811        let topology = *self
812            .hyper
813            .as_ref()
814            .ok_or("glm5_verify_rows on a model with no HyperConnections topology")?;
815        let t = tokens.len();
816        let cap = Self::hyper_batch_cap();
817        if t == 0 {
818            return Err("glm5_verify_rows: empty verify row set".into());
819        }
820        if t > cap {
821            return Err(format!(
822                "glm5_verify_rows: t={t} > cap {cap} — at t >= PRIME_MIN_T (16) the MoE \
823                 shared-expert trio crosses off the decode-exact class (the batched-decode \
824                 gate's measured B=16 knee), so per-row bit-identity vs plain decode breaks. \
825                 K <= cap-1 drafts per round"
826            )
827            .into());
828        }
829        let mut any_sharded = false;
830        for (il, layer) in self.layers.iter().enumerate() {
831            match &layer.mixer {
832                Mixer::Kda(la) => any_sharded |= la.tp.is_some(),
833                Mixer::Mla(mla) => any_sharded |= mla.tp.is_some(),
834                _ => {
835                    return Err(format!(
836                        "glm5_verify_rows: trunk layer {il} is not a KDA or MLA mixer — the \
837                         rollback contract below is built and gated for glm5_next's two state \
838                         classes only; a Full/Linear arm needs its own ckpt plane and gate"
839                    )
840                    .into());
841                }
842            }
843        }
844        // spec x TP composition (lane/glm5-composition): the per-row walk carries no TP
845        // rollback arm — a sharded trunk demands the BATCHED walk at t > 1 (t = 1 rounds
846        // ride the TP decode walk below; full accept is the only legal outcome there).
847        if any_sharded && t > 1 && !glm5_verify_batch_on() {
848            return Err(
849                "glm5_verify_rows: the trunk is glm5-TP-SHARDED and MEMRA_GLM5_VERIFY_BATCH=0 — \
850                 the per-row rollback seam carries no TP arm; the spec x TP composition \
851                 requires the batched verify walk (unset MEMRA_GLM5_VERIFY_BATCH or run \
852                 without the TP door)"
853                    .into(),
854            );
855        }
856
857        let n_embd = self.cfg.n_embd as usize;
858        let pos0 = cache.pos;
859
860        // Ckpt BEFORE any state moves.
861        let mut ckpt = Glm5VerifyCkpt {
862            pos: pos0,
863            latent_len: cache
864                .latent
865                .iter()
866                .take(self.layers.len())
867                .map(|plane| plane.as_ref().map(|plane| plane.len))
868                .collect(),
869            kda_conv_cols: (0..self.layers.len()).map(|_| None).collect(),
870            kda_ssm_snap: (0..self.layers.len()).map(|_| None).collect(),
871            kda_scan_stash: (0..self.layers.len()).map(|_| None).collect(),
872            kda_rows: (0..self.layers.len()).map(|_| None).collect(),
873            kda_tp: (0..self.layers.len()).map(|_| None).collect(),
874            rows: t,
875        };
876
877        // ppN door — the verify walk owns its stage split exactly as the batched decode
878        // walk does (`decode_step_batch_hyper_ppn`, decode_batch.rs). Loud refusal on an
879        // unqualified pipeline rewrite, never a single-engine walk over stage-sharded
880        // weights.
881        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
882            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
883                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
884            }
885            return self.glm5_verify_rows_ppn(e, tokens, cache, ckpt, &topology, &fence);
886        }
887
888        let pos = Glm5VerifyPos::new(e, pos0, t)?;
889        let embedded = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
890        let x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
891        let x = self.glm5_verify_range(
892            e,
893            &topology,
894            x,
895            0,
896            self.layers.len(),
897            &pos,
898            cache,
899            &mut ckpt,
900        )?;
901        let (logits, collapsed) = self.glm5_verify_head(e, &topology, &x, t)?;
902        Ok((logits, collapsed, ckpt))
903    }
904
905    /// One hc layer RANGE `[lo, hi)` of the verify walk — the body `glm5_verify_rows` ran
906    /// inline before the ppN twin landed, extracted so the unsplit walk and every pipeline
907    /// stage run the SAME code over their own range (the `hyper_range_decode` /
908    /// `decode_batch_layers` precedent: bit-identity between the arms is then structural,
909    /// not a coincidence of two maintained copies). At `lo=0, hi=n_layers` the launch
910    /// sequence is identical to the pre-extraction walk.
911    ///
912    /// KDA ckpt columns are cloned THROUGH `e` — under a split that is the owning stage's
913    /// engine, so each column lives on the device (and is ordered on the stream) that owns
914    /// its layer's state; `glm5_verify_rollback` restores through the same per-stage seam.
915    #[allow(clippy::too_many_arguments)]
916    // allow: the parameter list mirrors the range-walk call contract its siblings share
917    fn glm5_verify_range(
918        &self,
919        e: &Engine,
920        topology: &crate::hyper::HyperTopology,
921        mut x: CudaSlice<f32>,
922        lo: usize,
923        hi: usize,
924        pos: &Glm5VerifyPos,
925        cache: &mut Cache,
926        ckpt: &mut Glm5VerifyCkpt,
927    ) -> Res<CudaSlice<f32>> {
928        let t = pos.t;
929        let n_embd = self.cfg.n_embd as usize;
930        let eps = self.cfg.rms_eps;
931        // THE BATCHED MIXER ARM (lane/glm5-verify-batch, default ON): one t=K+1 call per
932        // layer per class instead of the per-row loop — KDA batches projections/conv/
933        // gates through the decode-exact classes with the recurrence sequential INSIDE
934        // one scan launch; MLA runs the SAME cached core at t rows on the rows-exact
935        // matmul classes (per-query causal selection + attention by construction).
936        // `0` = the per-row walk below, byte-for-byte (the rollback seam). Engagement is
937        // a receipt, announced once per process.
938        let batch = glm5_verify_batch_on() && t > 1;
939        {
940            static SAID: std::sync::Once = std::sync::Once::new();
941            SAID.call_once(|| {
942                if batch {
943                    eprintln!(
944                        "[glm5-spec] verify walk BATCHED per layer: kda=one t-call (scan \
945                         sequential in-kernel), mla=rows-exact t-call, head=rows-exact, \
946                         moe=pairs rows-call where qualified \
947                         (MEMRA_GLM5_VERIFY_BATCH default ON)"
948                    );
949                } else {
950                    eprintln!("[glm5-spec] verify walk PER-ROW (MEMRA_GLM5_VERIFY_BATCH=0 or t=1)");
951                }
952            });
953        }
954        let trace_v = crate::spec_phase::spec_trace_level() >= 2;
955        // Sub-phase clock (trace level 2 only): drain the walking stream so the elapsed
956        // ns lands in the mixer-class bucket — shares, never walls.
957        let vclock = |on: bool| -> Option<std::time::Instant> {
958            on.then(|| {
959                let _ = e.stream().synchronize();
960                std::time::Instant::now()
961            })
962        };
963        for il in lo..hi {
964            let layer = &self.layers[il];
965            let hyper = layer.hyper.as_ref().ok_or_else(|| {
966                format!("layer {il} carries no hyper-connection weights under an hc plan")
967            })?;
968
969            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, t, n_embd)?;
970            let mut h = e.uninit(t * n_embd)?;
971            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
972            // The batched arm refuses per-layer only for an MLA layer WITHOUT the DSA
973            // indexer: the absorbed t>1 attention arm has no per-row bit-identity claim
974            // at this seam, so it stays on the per-row loop by name (glm5_next always
975            // carries the indexer, so this is a foreign-geometry guard, not a live path).
976            let layer_batched = batch
977                && match &layer.mixer {
978                    Mixer::Kda(_) => true,
979                    Mixer::Mla(mla) => mla.index.is_some(),
980                    Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
981                };
982            let mixed = if layer_batched {
983                match &layer.mixer {
984                    // spec x TP composition: sharded mixers ride the TP verify walks —
985                    // per-rank batched rows calls, column-parallel-over-gather joins on
986                    // the rows-exact classes, per-rank rollback stash into the ckpt.
987                    Mixer::Kda(la) if la.tp.is_some() => {
988                        let t0 = vclock(trace_v);
989                        let mut scan_ns = 0u64;
990                        let (out, stash) = crate::glm5_tp::kda_tp_verify_rows(
991                            e,
992                            la,
993                            &h,
994                            t,
995                            eps,
996                            cache,
997                            il,
998                            trace_v.then_some(&mut scan_ns),
999                        )?;
1000                        ckpt.kda_tp[il] = Some(stash);
1001                        if let Some(t0) = t0 {
1002                            let _ = e.stream().synchronize();
1003                            use std::sync::atomic::Ordering;
1004                            crate::spec_phase::V_KDA_NS
1005                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
1006                            crate::spec_phase::V_KDA_SCAN_NS.fetch_add(scan_ns, Ordering::Relaxed);
1007                        }
1008                        out
1009                    }
1010                    Mixer::Mla(mla) if mla.tp.is_some() => {
1011                        let t0 = vclock(trace_v);
1012                        let out =
1013                            self.mla_tp_attn_cached(e, mla, &h, &pos.all, t, il, cache, true)?;
1014                        if let Some(t0) = t0 {
1015                            let _ = e.stream().synchronize();
1016                            use std::sync::atomic::Ordering;
1017                            crate::spec_phase::V_MLA_NS
1018                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
1019                        }
1020                        out
1021                    }
1022                    Mixer::Kda(la) => {
1023                        // Pre-round snapshot: ONE ssm clone per layer per round, BEFORE
1024                        // the batched call advances the resident state (ckpt doc; also
1025                        // the batched rollback's scan-replay base).
1026                        {
1027                            let rl = cache.recur[il]
1028                                .as_ref()
1029                                .ok_or("glm5 verify KDA layer has no recurrent state")?;
1030                            ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
1031                        }
1032                        let t0 = vclock(trace_v);
1033                        let mut scan_ns = 0u64;
1034                        let (out, stash) = crate::kda::kda_verify_rows_cached(
1035                            e,
1036                            la,
1037                            &h,
1038                            t,
1039                            eps,
1040                            cache,
1041                            il,
1042                            trace_v.then_some(&mut scan_ns),
1043                        )?;
1044                        ckpt.kda_rows[il] = Some(stash);
1045                        if let Some(t0) = t0 {
1046                            let _ = e.stream().synchronize();
1047                            use std::sync::atomic::Ordering;
1048                            crate::spec_phase::V_KDA_NS
1049                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
1050                            crate::spec_phase::V_KDA_SCAN_NS.fetch_add(scan_ns, Ordering::Relaxed);
1051                        }
1052                        out
1053                    }
1054                    Mixer::Mla(mla) => {
1055                        let t0 = vclock(trace_v);
1056                        let out =
1057                            self.mla_attn_cached_rows_exact(e, mla, &h, &pos.all, t, il, cache)?;
1058                        if let Some(t0) = t0 {
1059                            let _ = e.stream().synchronize();
1060                            use std::sync::atomic::Ordering;
1061                            crate::spec_phase::V_MLA_NS
1062                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
1063                        }
1064                        out
1065                    }
1066                    Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
1067                }
1068            } else {
1069                // ---- PER-ROW mixer walk (the rollback seam; also t=1 rounds and the
1070                // no-indexer MLA guard): row r's state input is row r-1's state output
1071                // (KDA) / rows 0..pos0+r (MLA latent) — each row the SAME t=1 call its
1072                // plain decode step makes.
1073                // h_row is hoisted out of the row loop (loop-port 3): the mixer consumes
1074                // it in stream order before the next row's overwrite, so ONE buffer per
1075                // layer replaces t allocations (stream-ordered pool churn is the dsv4
1076                // lesson).
1077                // Depth attribution (lane/spec-route-depth-20260902): every row that walks
1078                // this arm is counted; the per-round log samples the counter.
1079                V_SEQ_ROWS.fetch_add(t as u64, std::sync::atomic::Ordering::Relaxed);
1080                let mut mixed = e.uninit(t * n_embd)?;
1081                let mut h_row = e.uninit(n_embd)?;
1082                #[allow(clippy::needless_range_loop)]
1083                // allow: r is the sequential row cursor (slices h, offsets pos); iterating pos buffers would hide the row-chaining contract
1084                for r in 0..t {
1085                    e.dtod_copy_view(&h.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
1086                    // The per-row position buffer: prebuilt when the per-row arm owns the
1087                    // walk; built on demand for the rare per-layer refusal under batch.
1088                    let pos_row: CudaSlice<i32>;
1089                    let pos_r = if let Some(p) = pos.rows.get(r) {
1090                        p
1091                    } else {
1092                        pos_row = e.htod_i32(&[(pos.pos0 + r) as i32])?;
1093                        &pos_row
1094                    };
1095                    let out_row = match &layer.mixer {
1096                        // spec x TP composition on the per-row arm. KDA shards reach
1097                        // here at t == 1 ONLY, asserted locally: the walk-entry guard is
1098                        // a FLAG check two frames up (MEMRA_GLM5_VERIFY_BATCH=0 at t>1
1099                        // refuses), and under the batched flag layer_batched is
1100                        // unconditionally true for KDA — but neither is a structural
1101                        // invariant of THIS arm (#80 review's latent-trap finding). A
1102                        // sharded NO-INDEXER MLA layer legally lands here at any t
1103                        // (append + truncate rollback covers every keep; foreign
1104                        // geometry, never a live glm5_next path).
1105                        Mixer::Kda(la) if la.tp.is_some() => {
1106                            if t > 1 {
1107                                return Err(format!(
1108                                    "glm5 verify per-row arm reached a sharded KDA \
1109                                     layer {il} at t={t}: no per-rank rollback stash \
1110                                     exists on this arm (walk-entry guard bypassed?)"
1111                                )
1112                                .into());
1113                            }
1114                            crate::glm5_tp::kda_tp_cached(
1115                                e,
1116                                la,
1117                                &h_row,
1118                                1,
1119                                eps,
1120                                cache,
1121                                il,
1122                                crate::kda::ConvArm::Decode,
1123                            )?
1124                        }
1125                        Mixer::Mla(mla) if mla.tp.is_some() => {
1126                            self.mla_tp_attn_cached(e, mla, &h_row, pos_r, 1, il, cache, false)?
1127                        }
1128                        Mixer::Kda(la) => {
1129                            // Pre-round snapshot: ONE ssm clone per layer per round taken
1130                            // before row 0 mutates the resident state (loop-port 3).
1131                            if r == 0 && t > 1 {
1132                                let rl = cache.recur[il]
1133                                    .as_ref()
1134                                    .ok_or("glm5 verify KDA layer has no recurrent state")?;
1135                                ckpt.kda_ssm_snap[il] = Some(e.clone_dtod(&rl.ssm_state)?);
1136                            }
1137                            if r + 1 < t {
1138                                // Steal the row's scan inputs for the replay stash (zero
1139                                // copies); clone only the small conv ring per row.
1140                                let (out, inputs) = crate::kda::kda_decode_cached_stash(
1141                                    e, la, &h_row, eps, cache, il,
1142                                )?;
1143                                let rl = cache.recur[il]
1144                                    .as_ref()
1145                                    .ok_or("glm5 verify KDA layer has no recurrent state")?;
1146                                ckpt.kda_conv_cols[il]
1147                                    .get_or_insert_with(Vec::new)
1148                                    .push(e.clone_dtod(&rl.conv_state)?);
1149                                ckpt.kda_scan_stash[il]
1150                                    .get_or_insert_with(Vec::new)
1151                                    .push(inputs);
1152                                out
1153                            } else {
1154                                crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?
1155                            }
1156                        }
1157                        Mixer::Mla(mla) => {
1158                            self.mla_attn_cached(e, mla, &h_row, pos_r, 1, il, cache)?
1159                        }
1160                        // Refused at entry; unreachable keeps the match total without a silent arm.
1161                        Mixer::Full(_) | Mixer::Linear(_) => unreachable!("refused at walk entry"),
1162                    };
1163                    e.copy_into(&mut mixed, r * n_embd, &out_row, n_embd)?;
1164                }
1165                mixed
1166            };
1167            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
1168
1169            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, t, n_embd)?;
1170            let mut z = e.uninit(t * n_embd)?;
1171            e.rms_norm(
1172                &y,
1173                layer.post_attn_norm.float_data(),
1174                &mut z,
1175                n_embd,
1176                t,
1177                eps,
1178            )?;
1179            // FFN branch: `batch` arms the pairs-shaped batched MoE across the t rows
1180            // (lane/glm5-vrest — fail-closed inside to the byte-identical sequential
1181            // loop); the =0 arm keeps the pre-lane per-(token,expert) class. Clocked
1182            // into the vffn sub-bucket at trace level 2 (batched arm only, like vkda).
1183            let t0 = vclock(trace_v && batch);
1184            let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, t, il, batch)?;
1185            if let Some(t0) = t0 {
1186                let _ = e.stream().synchronize();
1187                use std::sync::atomic::Ordering;
1188                crate::spec_phase::V_FFN_NS
1189                    .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
1190            }
1191            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
1192            // glm5 DFlash2 feature tap (module doc, DRAFT SOURCE SEAM): the verify rows'
1193            // contracted completed-layer outputs are next round's drafter context.
1194            self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
1195        }
1196        Ok(x)
1197    }
1198
1199    /// Write one tapped layer's CONTRACTED completed output into the armed
1200    /// [`HcTapSink`] — the glm5 DFlash2 drafter's measured feature contract (stream-mean
1201    /// over the hyper streams, the probe's `hc_contract` capture definition). Staged
1202    /// through the WALKING engine `e` (the owning stage engine under a ppN split), so the
1203    /// sink is placement-invariant. One Option check when unarmed; nothing else pays.
1204    ///
1205    /// Two staging arms (loop-port 1):
1206    ///   * `device_stage` (the verify-round sink): ONE async D2D into the slot's device
1207    ///     buffer — the walk never blocks; the round drains all slots post-walk in its
1208    ///     single sync point (`glm5_tap_drain`). Kills the five in-walk DtoHs the 3way
1209    ///     window priced into the 31.6 ms fixed round cost (map row #17).
1210    ///   * host-staged (prime sinks): the pre-port behavior — per-chunk DtoH, amortized
1211    ///     over the prime's >= 256-row chunks.
1212    pub(crate) fn glm5_hc_tap(
1213        &self,
1214        e: &Engine,
1215        cache: &mut Cache,
1216        topology: &crate::hyper::HyperTopology,
1217        il: usize,
1218        x: &CudaSlice<f32>,
1219        t: usize,
1220    ) -> Res<()> {
1221        let Some(sink) = cache.hc_taps.as_mut() else {
1222            return Ok(());
1223        };
1224        let base = sink.base;
1225        self.glm5_hc_tap_into(e, sink, base, topology, il, x, t)
1226    }
1227
1228    /// [`Self::glm5_hc_tap`] with the sink and its row base passed EXPLICITLY instead of read
1229    /// from the cache.
1230    ///
1231    /// The pipelined mHC prime needs this seam. Its two stage threads run different CHUNKS at
1232    /// the same moment, so a single `sink.base` field cannot describe both, and
1233    /// [`PrimeCacheStages`] hands each stage a cache shell with `hc_taps: None` — which is why
1234    /// arm 2 used to refuse outright when the DFlash2 drafter had armed a sink. Passing the
1235    /// base per call is what makes one shared sink correct for two concurrent walks.
1236    #[allow(clippy::too_many_arguments)] // allow: the list is the tap contract, base included
1237    pub(crate) fn glm5_hc_tap_into(
1238        &self,
1239        e: &Engine,
1240        sink: &mut HcTapSink,
1241        base_abs: usize,
1242        topology: &crate::hyper::HyperTopology,
1243        il: usize,
1244        x: &CudaSlice<f32>,
1245        t: usize,
1246    ) -> Res<()> {
1247        let Some(slot) = sink.layer_ids.iter().position(|&l| l == il) else {
1248            return Ok(());
1249        };
1250        let saved = sink.base;
1251        sink.base = base_abs;
1252        let out = self.glm5_hc_tap_slot(e, sink, slot, topology, x, t);
1253        sink.base = saved;
1254        out
1255    }
1256
1257    fn glm5_hc_tap_slot(
1258        &self,
1259        e: &Engine,
1260        sink: &mut HcTapSink,
1261        slot: usize,
1262        topology: &crate::hyper::HyperTopology,
1263        x: &CudaSlice<f32>,
1264        t: usize,
1265    ) -> Res<()> {
1266        let h = sink.hidden;
1267        let n_taps = sink.layer_ids.len();
1268        // Sink-relative row of this walk's row 0 (doc on `HcTapSink::origin`): fresh-prompt
1269        // sinks have origin 0 and this is exactly the pre-field arithmetic; a suffix-prime
1270        // sink is anchored at the restored boundary. A base below the origin is a caller
1271        // bug (a walk over rows the sink does not cover) — refuse, never wrap.
1272        let base = sink.base.checked_sub(sink.origin).ok_or_else(|| {
1273            format!(
1274                "hc tap base {} below sink origin {} (walk outside the sink's window)",
1275                sink.base, sink.origin,
1276            )
1277        })?;
1278        debug_assert!(
1279            base + t <= sink.t,
1280            "hc tap window {base}+{t} exceeds sink {}",
1281            sink.t
1282        );
1283        let contracted = crate::hyper::contract_mean(e, topology, x, t, h)?;
1284        if sink.device_stage {
1285            // Lazy slot buffer on the WRITING engine (this layer always walks on one
1286            // stage, so the buffer's device is stable for the sink's lifetime). Every
1287            // walk row writes every tapped layer, so the buffer is fully covered by the
1288            // walk that armed the sink.
1289            if sink.dev[slot].is_none() {
1290                sink.dev[slot] = Some(e.uninit(sink.t * h)?);
1291            }
1292            let buf = sink.dev[slot].as_mut().expect("just filled");
1293            e.copy_into(buf, base * h, &contracted, t * h)?;
1294            return Ok(());
1295        }
1296        let t_dtoh = std::time::Instant::now();
1297        let host = e.dtoh(&contracted)?;
1298        for r in 0..t {
1299            let dst = (base + r) * n_taps * h + slot * h;
1300            sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
1301        }
1302        sink.dtoh_ns += t_dtoh.elapsed().as_nanos() as u64;
1303        Ok(())
1304    }
1305
1306    /// Drain a device-staged tap sink into its host `rows` — the round's ONE post-walk
1307    /// sync point for tap features (loop-port 1). Each slot reads back through its
1308    /// layer's OWNING engine (the stage engine under a live split, the caller's engine
1309    /// otherwise); the verify walk's terminal drain has already retired every stage's
1310    /// writes (stream program order: the slot copy precedes its stage's TX, and the
1311    /// TX-wait chain covers it transitively — the pp.rs multi-stream law).
1312    fn glm5_tap_drain(&self, e: &Engine, sink: &mut HcTapSink) -> Res<()> {
1313        if !sink.device_stage {
1314            return Ok(());
1315        }
1316        let h = sink.hidden;
1317        let n_taps = sink.layer_ids.len();
1318        let split = match crate::pp::pp_cuts(self.layers.len()) {
1319            Some(fence) if !crate::pp::pp2_streams_off() => {
1320                Some((crate::pp::PpNRt::get(e)?, fence))
1321            }
1322            _ => None,
1323        };
1324        for slot in 0..n_taps {
1325            let Some(buf) = sink.dev[slot].take() else {
1326                continue;
1327            };
1328            let il = sink.layer_ids[slot];
1329            let es = match split.as_ref() {
1330                Some((rt, fence)) => {
1331                    let stage = fence
1332                        .windows(2)
1333                        .position(|w| il >= w[0] && il < w[1])
1334                        .ok_or_else(|| format!("tap layer {il} outside every stage range"))?;
1335                    rt.engine(stage, e)
1336                }
1337                None => e,
1338            };
1339            let host = es.dtoh(&buf)?;
1340            for r in 0..sink.t {
1341                let dst = r * n_taps * h + slot * h;
1342                sink.rows[dst..dst + h].copy_from_slice(&host[r * h..(r + 1) * h]);
1343            }
1344        }
1345        Ok(())
1346    }
1347
1348    /// Trunk exit of the verify walk, the batched head's decode-exact form
1349    /// (`hyper_batch_head_logits`), with the collapsed pre-output_norm rows kept — they are
1350    /// the h_seeds the MTP head re-seeds from (LANE.md §A). Under a split this runs on the
1351    /// LAST stage's engine, where the loader put `output_norm` + the lm head.
1352    fn glm5_verify_head(
1353        &self,
1354        e: &Engine,
1355        topology: &crate::hyper::HyperTopology,
1356        x: &CudaSlice<f32>,
1357        t: usize,
1358    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>)> {
1359        let n_embd = self.cfg.n_embd as usize;
1360        let eps = self.cfg.rms_eps;
1361        let collapsed =
1362            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
1363        let mut hn = e.uninit(t * n_embd)?;
1364        e.rms_norm(
1365            &collapsed,
1366            self.output_norm.float_data(),
1367            &mut hn,
1368            n_embd,
1369            t,
1370            eps,
1371        )?;
1372        // Under the batched walk the lm head rides the rows-exact classes too (the bf16
1373        // tcols twin reads the 1.27 GB head ONCE per round instead of once per row);
1374        // per-row bits unchanged by contract, the tcols bit-gate holds it.
1375        let logits = if glm5_verify_batch_on() && t > 1 {
1376            e.matmul_rows_exact(&self.output, &hn, t)?
1377        } else {
1378            e.matmul_decode_exact(&self.output, &hn, t)?
1379        };
1380        Ok((logits, collapsed))
1381    }
1382
1383    /// ppN twin of the verify walk (lane/glm5-ppn-verify, 2026-08-30), mirroring
1384    /// `decode_step_batch_hyper_ppn` (decode_batch.rs): the t=K+1 rows walk as N stage
1385    /// subgraphs — per-stage engine, per-stage pos_rows uploads, ONE `[t, streams, n_embd]`
1386    /// boundary payload per fence cut. Row chaining is per-LAYER through the one cache
1387    /// (row r+1 at layer il depends only on row r at layer il), so a straight layer-range
1388    /// split preserves it exactly; no row ever crosses a boundary individually. Head +
1389    /// collapsed rows land on the LAST stage's engine — where the loader put the lm head
1390    /// and where `pp::new_cache*` places the MTP plane the re-seed feeds.
1391    ///
1392    /// DRAIN CONTRACT: this walk returns DEVICE buffers with no terminal dtoh (unlike its
1393    /// decode twins, whose epilogue reads back on the last stage's stream), so it owns the
1394    /// settle — the per-stage arm synchronizes the LAST stage's stream before returning.
1395    /// The TX-wait chain transitively covers every earlier stage (pp.rs multi-stream law),
1396    /// so the logits, the collapsed rows AND the ckpt's per-stage KDA columns are all safe
1397    /// for consumption from the caller's streams after this returns.
1398    #[allow(clippy::too_many_arguments)]
1399    // allow: the parameter list mirrors its decode twin's stage-walk contract
1400    fn glm5_verify_rows_ppn(
1401        &self,
1402        e: &Engine,
1403        tokens: &[u32],
1404        cache: &mut Cache,
1405        mut ckpt: Glm5VerifyCkpt,
1406        topology: &crate::hyper::HyperTopology,
1407        fence: &[usize],
1408    ) -> Res<(CudaSlice<f32>, CudaSlice<f32>, Glm5VerifyCkpt)> {
1409        let t = tokens.len();
1410        let n_embd = self.cfg.n_embd as usize;
1411        let pos0 = ckpt.pos;
1412        let payload = t * topology.streams * n_embd;
1413        // Position buffers through THIS stage's engine (the per-stage pos_d law:
1414        // allocated, consumed and freed on one stage's stream).
1415        let pos_on = |eng: &Engine| -> Res<Glm5VerifyPos> { Glm5VerifyPos::new(eng, pos0, t) };
1416
1417        // Same-stream seam (MEMRA_PP_STREAMS=0): one engine, boundary copies between
1418        // ranges — the shape every hc ppN walk uses for this knob.
1419        if crate::pp::pp2_streams_off() {
1420            let pos = pos_on(e)?;
1421            let embedded = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
1422            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
1423            x =
1424                self.glm5_verify_range(e, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1425            for s in 1..fence.len() - 1 {
1426                let boundary_tx = e.clone_dtod(&x)?;
1427                let boundary_rx = e.clone_dtod(&boundary_tx)?;
1428                x = self.glm5_verify_range(
1429                    e,
1430                    topology,
1431                    boundary_rx,
1432                    fence[s],
1433                    fence[s + 1],
1434                    &pos,
1435                    cache,
1436                    &mut ckpt,
1437                )?;
1438            }
1439            let (logits, collapsed) = self.glm5_verify_head(e, topology, &x, t)?;
1440            return Ok((logits, collapsed, ckpt));
1441        }
1442
1443        let rt = crate::pp::PpNRt::get(e)?;
1444        let n_st = fence.len() - 1;
1445        assert_eq!(
1446            rt.n_stages(),
1447            n_st,
1448            "PpNRt stage count {} != fence stages {n_st}",
1449            rt.n_stages()
1450        );
1451        // #87 reverse publication (see decode_step_batch_ppn): order every stage stream
1452        // behind the caller before this body's first stage allocation.
1453        rt.fence_stages_behind(&e.stream())?;
1454
1455        // ---- STAGE 0: embed + expand (no weights) + layers [0, fence[1]) + TX ----
1456        let mut slot = {
1457            let _st0 = rt.enter(0);
1458            let e0 = rt.engine(0, e);
1459            let pos = pos_on(e0)?;
1460            let embedded = e0.htod(&self.embd.try_gather(n_embd, tokens)?)?;
1461            let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
1462            let x = self
1463                .glm5_verify_range(e0, topology, x, fence[0], fence[1], &pos, cache, &mut ckpt)?;
1464            rt.tx(0, &x, payload)?
1465        };
1466
1467        // ---- MIDDLE STAGES: RX -> range -> TX ----
1468        for s in 1..n_st - 1 {
1469            let _st = rt.enter(s);
1470            let es = rt.engine(s, e);
1471            let pos = pos_on(es)?;
1472            let x = rt.rx(s - 1, slot, payload)?;
1473            let x = self.glm5_verify_range(
1474                es,
1475                topology,
1476                x,
1477                fence[s],
1478                fence[s + 1],
1479                &pos,
1480                cache,
1481                &mut ckpt,
1482            )?;
1483            slot = rt.tx(s, &x, payload)?;
1484        }
1485
1486        // ---- LAST STAGE: RX + final range + collapse/head + the drain (doc above) ----
1487        let _stl = rt.enter(n_st - 1);
1488        let el = rt.engine(n_st - 1, e);
1489        let pos = pos_on(el)?;
1490        let x = rt.rx(n_st - 2, slot, payload)?;
1491        let x = self.glm5_verify_range(
1492            el,
1493            topology,
1494            x,
1495            fence[n_st - 1],
1496            fence[n_st],
1497            &pos,
1498            cache,
1499            &mut ckpt,
1500        )?;
1501        let (logits, collapsed) = self.glm5_verify_head(el, topology, &x, t)?;
1502        // el.stream() under the enter-guard IS the stage stream (memra_runtime ambient
1503        // override) — this drain settles the whole walk transitively.
1504        el.stream().synchronize()?;
1505        drop(_stl);
1506        // EXIT PUBLICATION (lane/glm5-accrace): the drain above settles the LAST stage, and
1507        // the TX-wait chain covers every earlier stage's work only UP TO its `ev_tx`. Each
1508        // earlier stage's stream still holds the stage-scope tail its locals enqueue when
1509        // they drop under the override (`pos`, the boundary residual, the per-layer
1510        // transients, this round's ckpt clones). The caller resumes and allocates for the
1511        // accept walk and the MTP re-seed, so it must be ordered behind ALL stages.
1512        self.glm5_publish_stages(e)?;
1513        Ok((logits, collapsed, ckpt))
1514    }
1515
1516    /// The engine that owns the trunk exit (collapse + output_norm + lm head), the MTP
1517    /// block's weights AND its latent plane under a ppN split: the LAST stage's engine —
1518    /// `hybrid.rs` uploads the head there (`pp::layer_engine(e, n_trunk, n_trunk - 1)`)
1519    /// and `pp::new_cache*` maps trailing MTP/NextN planes to the last stage. Door shut or
1520    /// the same-stream seam: the caller's engine, unchanged (single-device callers pay
1521    /// nothing — `e` is returned by identity).
1522    fn glm5_head_engine<'e>(&self, e: &'e Engine) -> Res<&'e Engine> {
1523        match crate::pp::pp_cuts(self.layers.len()) {
1524            Some(fence) if !crate::pp::pp2_streams_off() => {
1525                let rt = crate::pp::PpNRt::get(e)?;
1526                Ok(rt.engine(fence.len() - 2, e))
1527            }
1528            _ => Ok(e),
1529        }
1530    }
1531
1532    /// Roll the trunk back to exactly `keep` accepted verify rows (1 <= keep <= t; keep =
1533    /// j+1: the always-committed anchor row plus j accepted drafts).
1534    ///
1535    /// - MLA latent planes: `len = snapshot + keep` (truncate; rows are position-addressed
1536    ///   and append-only, so the kept rows ARE what a plain decode chain would have
1537    ///   written — the decode-exact contract), device `len_d` in lock-step, and
1538    ///   `truncate_index_pool_keys(pool)` clamps pool-key finality to what the shortened
1539    ///   `len` still justifies (the tail-ring residency tripwire fires on the next call if
1540    ///   this clamp is ever skipped).
1541    /// - KDA state: restore column keep-1 (state after the last kept row); full accept
1542    ///   (keep == t) keeps the resident state — the columns are clones OF it.
1543    /// - `cache.pos = snapshot + keep`.
1544    ///
1545    /// Under a live ppN split each stage's layers restore THROUGH that stage's engine ON
1546    /// its stream: the state planes and the ckpt columns live on the owning stage's device
1547    /// (per-stage `KvDev` allocation; per-stage clones in the walk), and enqueuing the
1548    /// restores on the same stage streams the walk writes on orders them relative to the
1549    /// walk without any extra fence.
1550    ///
1551    /// THE EXIT PUBLICATION IS NOT OPTIONAL (lane/glm5-accrace 2026-09-01). This body used
1552    /// to return with the restores merely ENQUEUED on the stage streams, on the reasoning
1553    /// that "the next walk's own entry fence covers the primary-stream seam". It does not:
1554    /// `fence_stages_behind` orders the STAGE streams behind the CALLER, and everything the
1555    /// round does after this point — the MTP plane reset, the h_seed rows, the next round's
1556    /// whole draft chain, the next SESSION's cache allocation and prime — runs on the
1557    /// CALLER's stream and ALLOCATES. cudarc's drops carry no read guard, so the pool could
1558    /// hand the caller a block whose stage-stream lifetime had not retired and the caller's
1559    /// writes landed under queued rollback work.
1560    ///
1561    /// MEASURED CONSEQUENCE, and why a "rollback ordering" bug showed up as a PRIME bug:
1562    /// with per-stage streams on one device the hc ppN prime over a fixed 24-token prompt
1563    /// returned three distinct logit fingerprints inside one process (a third of all primes
1564    /// non-canonical); downstream, one glm5 spec round lost an acceptance silently
1565    /// (14/42 -> 13/42) and the e2e tape diverged. Publishing here took non-canonical primes
1566    /// from 20/110 to 2/110 in an interleaved A/B, and the walk's own exit publication
1567    /// closed the remainder. Receipts:
1568    /// `research/glm53-flash-bringup-20260827/accrace-20260901/LANE.md`.
1569    pub fn glm5_verify_rollback(
1570        &self,
1571        e: &Engine,
1572        cache: &mut Cache,
1573        ckpt: &Glm5VerifyCkpt,
1574        keep: usize,
1575    ) -> Res<()> {
1576        if keep == 0 || keep > ckpt.rows {
1577            return Err(format!(
1578                "glm5_verify_rollback: keep={keep} outside 1..={} (the anchor row is always \
1579                 committed; keep = accepted drafts + 1)",
1580                ckpt.rows
1581            )
1582            .into());
1583        }
1584        match crate::pp::pp_cuts(self.layers.len()) {
1585            Some(fence) if !crate::pp::pp2_streams_off() => {
1586                let rt = crate::pp::PpNRt::get(e)?;
1587                for s in 0..fence.len() - 1 {
1588                    let _st = rt.enter(s);
1589                    let es = rt.engine(s, e);
1590                    for il in fence[s]..fence[s + 1] {
1591                        self.glm5_rollback_layer(es, cache, ckpt, keep, il)?;
1592                    }
1593                }
1594                // EXIT PUBLICATION (doc above): every stage stream, to the caller's.
1595                self.glm5_publish_stages(e)?;
1596            }
1597            _ => {
1598                for il in 0..self.layers.len() {
1599                    self.glm5_rollback_layer(e, cache, ckpt, keep, il)?;
1600                }
1601            }
1602        }
1603        cache.pos = ckpt.pos + keep;
1604        Ok(())
1605    }
1606
1607    /// Restore ONE trunk layer to the ckpt's `keep`-row state (the per-plane contract in
1608    /// [`Self::glm5_verify_rollback`]'s doc). `e` is the layer's OWNING engine — the stage
1609    /// engine under a split, the caller's engine otherwise.
1610    fn glm5_rollback_layer(
1611        &self,
1612        e: &Engine,
1613        cache: &mut Cache,
1614        ckpt: &Glm5VerifyCkpt,
1615        keep: usize,
1616        il: usize,
1617    ) -> Res<()> {
1618        match &self.layers[il].mixer {
1619            Mixer::Mla(mla) => {
1620                if keep == ckpt.rows {
1621                    // Full accept: the walk already advanced len AND the len_d device
1622                    // mirror to saved + rows on the canonical plane and every replica
1623                    // (append-time stores), so the restore below would rewrite unchanged
1624                    // values — ~11 synchronizing pageable 4-byte copies per round on the
1625                    // HOT outcome (the KDA arms' early-out twin; #82 review).
1626                    return Ok(());
1627                }
1628                let saved = ckpt.latent_len[il].ok_or_else(|| {
1629                    format!("glm5_verify_rollback: MLA layer {il} missing from the ckpt")
1630                })?;
1631                let plane = cache.latent[il].as_mut().ok_or_else(|| {
1632                    format!("glm5_verify_rollback: MLA layer {il} has no latent plane")
1633                })?;
1634                plane.len = saved + keep;
1635                let len_i32 =
1636                    i32::try_from(plane.len).map_err(|_| "latent length exceeds i32 mirror")?;
1637                // Door H (`MEMRA_GLM5_HTOD_DIET`): async `i32_set_k` instead of the synchronizing
1638                // pageable 4-byte copy — 11 of these per round, and unconditional (unlike the
1639                // KDA arm, which short-circuits when `keep == rows`).
1640                e.i32_mirror_store(&mut plane.len_d, len_i32)?;
1641                if let Some(indexer) = mla.index.as_ref() {
1642                    plane.truncate_index_pool_keys(indexer.geom.pool);
1643                }
1644                // spec x TP composition: the PEER latent replicas append in lock-step with
1645                // the canonical plane (the TP walk's construction), so the same truncation
1646                // restores each of them — through its own rank's engine for the device
1647                // `len_d` mirror. Full accept skips the loop (lens already read
1648                // saved + rows — the KDA arm's early-out twin; each skipped store is a
1649                // synchronizing pageable copy per rank per layer on the HOT outcome).
1650                // Missing replicas after a verify walk are a wiring bug and refuse by
1651                // name, never a silent canonical-only restore (#80 review hardening).
1652                if let Some(tp) = mla.tp.as_ref() {
1653                    let replicas = cache.glm5_tp_latent_peer[il].as_mut().ok_or_else(|| {
1654                        format!(
1655                            "glm5_verify_rollback: sharded MLA layer {il} has no peer \
1656                             latent replicas (the TP verify walk hydrates them; a \
1657                             rollback without them would silently restore the canonical \
1658                             plane only)"
1659                        )
1660                    })?;
1661                    for (i, replica) in replicas.iter_mut().enumerate() {
1662                        replica.len = saved + keep;
1663                        tp.rt.peers[i].i32_mirror_store(&mut replica.len_d, len_i32)?;
1664                        if let Some(indexer) = mla.index.as_ref() {
1665                            replica.truncate_index_pool_keys(indexer.geom.pool);
1666                        }
1667                    }
1668                }
1669            }
1670            Mixer::Kda(la) if la.tp.is_some() => {
1671                if keep == ckpt.rows {
1672                    return Ok(()); // resident per-rank states ARE the post-keep states
1673                }
1674                let stash = ckpt.kda_tp[il].as_ref().ok_or_else(|| {
1675                    format!(
1676                        "glm5_verify_rollback: sharded KDA layer {il} has no per-rank stash \
1677                         (the batched TP verify walk fills it; the per-row arm is refused \
1678                         at walk entry)"
1679                    )
1680                })?;
1681                crate::glm5_tp::kda_tp_verify_rollback(e, la, stash, keep, cache, il)?;
1682            }
1683            Mixer::Kda(la) => {
1684                if keep == ckpt.rows {
1685                    return Ok(()); // resident state IS the state after the last kept row
1686                }
1687                // BATCHED-walk stash (lane/glm5-verify-batch): ring restore + re-roll,
1688                // then ONE scan replay at T=keep from the pre-round snapshot.
1689                if let Some(stash) = ckpt.kda_rows[il].as_ref() {
1690                    let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1691                        format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1692                    })?;
1693                    return crate::kda::kda_verify_rollback_rows(
1694                        e, la, snap, stash, keep, cache, il,
1695                    );
1696                }
1697                // Conv ring: restore the cloned column (unchanged — 288 KiB).
1698                let conv_cols = ckpt.kda_conv_cols[il].as_ref().ok_or_else(|| {
1699                    format!("glm5_verify_rollback: KDA layer {il} has no conv columns")
1700                })?;
1701                let conv = &conv_cols[keep - 1];
1702                {
1703                    let rl = cache.recur[il].as_mut().ok_or_else(|| {
1704                        format!("glm5_verify_rollback: KDA layer {il} has no recurrent state")
1705                    })?;
1706                    e.copy_into(&mut rl.conv_state, 0, conv, conv.len())?;
1707                }
1708                // Recurrent state: REPLAY rows 0..keep from the pre-round snapshot
1709                // (loop-port 3; ckpt doc) — each replay re-issues that row's original
1710                // t=1 scan over its stolen inputs, so the rebuilt state is byte-identical
1711                // to the per-row clone this retires.
1712                let snap = ckpt.kda_ssm_snap[il].as_ref().ok_or_else(|| {
1713                    format!("glm5_verify_rollback: KDA layer {il} has no ssm snapshot")
1714                })?;
1715                let stash = ckpt.kda_scan_stash[il].as_ref().ok_or_else(|| {
1716                    format!("glm5_verify_rollback: KDA layer {il} has no scan stash")
1717                })?;
1718                crate::kda::kda_scan_replay(e, la, snap, &stash[..keep], cache, il)?;
1719            }
1720            Mixer::Full(_) | Mixer::Linear(_) => {
1721                return Err(format!(
1722                    "glm5_verify_rollback: layer {il} mixer class was refused at walk \
1723                     entry and cannot appear in a ckpt"
1724                )
1725                .into());
1726            }
1727        }
1728        Ok(())
1729    }
1730
1731    /// Reset the MTP draft plane (il = n_trunk) to `len` rows — the LANE.md rollback
1732    /// contract ("one row per step, rollback = plane len reset") plus the same pool-key
1733    /// clamp every len-shortening path owes the tail ring. The plane lives on the LAST
1734    /// stage under a split (`pp::new_cache*` maps trailing MTP planes there), so the
1735    /// device mirror writes through the head engine.
1736    fn glm5_mtp_plane_reset(&self, e: &Engine, cache: &mut Cache, len: usize) -> Res<()> {
1737        let e = self.glm5_head_engine(e)?;
1738        let mtp = self
1739            .mtp
1740            .as_ref()
1741            .ok_or("glm5_mtp_plane_reset with no MTP head loaded")?;
1742        let il = self
1743            .plan
1744            .mtp_blocks
1745            .first()
1746            .ok_or("ModelPlan declares no MTP block")?
1747            .layer
1748            .index as usize;
1749        let plane = cache
1750            .latent
1751            .get_mut(il)
1752            .and_then(|plane| plane.as_mut())
1753            .ok_or_else(|| format!("MTP block layer {il} has no latent cache plane"))?;
1754        if len > plane.len {
1755            return Err(format!(
1756                "glm5_mtp_plane_reset: target {len} is past the plane's {} rows — a reset \
1757                 only ever shortens",
1758                plane.len
1759            )
1760            .into());
1761        }
1762        plane.len = len;
1763        let len_i32 = i32::try_from(len).map_err(|_| "latent length exceeds i32 mirror")?;
1764        e.stream().memcpy_htod(&[len_i32], &mut plane.len_d)?;
1765        if let Mixer::Mla(mla) = &mtp.mixer
1766            && let Some(indexer) = mla.index.as_ref()
1767        {
1768            plane.truncate_index_pool_keys(indexer.geom.pool);
1769        }
1770        Ok(())
1771    }
1772
1773    /// Single-shot glm5 speculative generation: draft (MTP head, K steps) -> verify (one
1774    /// t=K+1 walk) -> accept j (greedy longest matching prefix) -> rollback -> re-seed.
1775    /// Returns `(tokens, drafted, accepted)` — `generate_spec`'s contract. One-shot form:
1776    /// builds a [`Glm5SpecSession`] over a fresh cache and drives it to `max_new` — the
1777    /// SAME round machinery the serve worker bursts, so the tparallel gate's byte-identity
1778    /// pins cover the served path's rounds too.
1779    pub fn generate_spec_glm5(
1780        &self,
1781        e: &Engine,
1782        prompt: &[u32],
1783        max_new: usize,
1784        k: usize,
1785    ) -> Res<(Vec<u32>, usize, usize)> {
1786        self.generate_spec_glm5_gated(e, prompt, max_new, k, Glm5SpecKnobs::default())
1787    }
1788
1789    /// `generate_spec_glm5` with GATE INSTRUMENTS (never a serving surface): a draft
1790    /// override for deterministic forced-accept / forced-reject rounds, and a
1791    /// rollback-disable arm that red-proves the end-to-end byte-identity gate.
1792    pub fn generate_spec_glm5_gated(
1793        &self,
1794        e: &Engine,
1795        prompt: &[u32],
1796        max_new: usize,
1797        k: usize,
1798        mut knobs: Glm5SpecKnobs<'_>,
1799    ) -> Res<(Vec<u32>, usize, usize)> {
1800        let cap = Self::hyper_batch_cap();
1801        if k == 0 || k + 1 > cap {
1802            return Err(format!(
1803                "generate_spec_glm5: k={k} outside 1..={} (verify rows = k+1 must stay \
1804                 inside the decode-exact knee, cap {cap})",
1805                cap - 1
1806            )
1807            .into());
1808        }
1809        if max_new == 0 {
1810            return Ok((Vec::new(), 0, 0));
1811        }
1812        let max_ctx = prompt.len() + max_new + k + 8;
1813        let mut sess = self.glm5_spec_session_new(e, prompt, max_ctx, None)?;
1814        let mut out: Vec<u32> = Vec::with_capacity(max_new + k);
1815        let mut drafted = 0usize;
1816        let mut accepted = 0usize;
1817        while out.len() < max_new && !sess.finished() {
1818            let (burst, d, a) = self.glm5_spec_session_burst_gated(
1819                e,
1820                &mut sess,
1821                max_new - out.len(),
1822                k,
1823                &[],
1824                &mut knobs,
1825            )?;
1826            if burst.is_empty() {
1827                break; // ctx guard tripped with nothing new — never spin
1828            }
1829            out.extend(burst);
1830            drafted += d;
1831            accepted += a;
1832        }
1833        out.truncate(max_new);
1834        Ok((out, drafted, accepted))
1835    }
1836
1837    /// BATCHED MTP-PLANE WARM (loop-port fold-in; doc at the call site in
1838    /// `glm5_spec_session_new`): fill the NextN block's latent plane with rows for pairs
1839    /// `(tokens_next[i], hiddens row i)`, i in `0..t`, in chunked t-parallel passes —
1840    /// ops 1-7 of `mtp_head_forward_mla_cached` batched over the chunk (embed gather,
1841    /// enorm/hnorm, the eh_proj concat via `place_rows_strided`, attn_norm), then ONE
1842    /// `mla_attn_cached` append per chunk (the prime-class t>1 arm the trunk's own MLA
1843    /// layers warm through; its attention output is discarded — the plane rows are the
1844    /// product). The MoE FFN, final norm and lm-head of the per-token chain are never
1845    /// run: they fed nothing but the (discarded) draft logits of prompt positions.
1846    fn glm5_mtp_plane_fill(
1847        &self,
1848        e: &Engine,
1849        tokens_next: &[u32],
1850        hiddens: &CudaSlice<f32>,
1851        t: usize,
1852        cache: &mut Cache,
1853    ) -> Res<()> {
1854        let mtp = self
1855            .mtp
1856            .as_ref()
1857            .ok_or("glm5_mtp_plane_fill with no MTP head loaded")?;
1858        let il = self
1859            .plan
1860            .mtp_blocks
1861            .first()
1862            .ok_or("ModelPlan declares no MTP block")?
1863            .layer
1864            .index as usize;
1865        let Mixer::Mla(mla) = &mtp.mixer else {
1866            return Err("glm5_mtp_plane_fill serves MLA-mixer MTP blocks only".into());
1867        };
1868        if tokens_next.len() < t {
1869            return Err(format!(
1870                "glm5_mtp_plane_fill: {t} rows requested over {} successor tokens",
1871                tokens_next.len()
1872            )
1873            .into());
1874        }
1875        let n_embd = self.cfg.n_embd as usize;
1876        let eps = self.cfg.rms_eps;
1877        // Chunk bound: the trunk prime's workspace discipline — bounds the t>1 attention
1878        // workspace and the transient buffers below without changing the append semantics
1879        // (`mla_attn_cached` appends at the plane's running length either way).
1880        const CHUNK: usize = 512;
1881        let mut done = 0usize;
1882        while done < t {
1883            let tc = (t - done).min(CHUNK);
1884            let e_emb = e.htod(
1885                &self
1886                    .embd
1887                    .try_gather(n_embd, &tokens_next[done..done + tc])?,
1888            )?;
1889            let mut e_norm = e.uninit(tc * n_embd)?;
1890            e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, tc, eps)?;
1891            // hnorm over the chunk's hidden rows (one contiguous view copy — rms_norm
1892            // takes an owned-slice operand).
1893            let hv = e.view(hiddens, (done + tc) * n_embd);
1894            let mut h_rows = e.uninit(tc * n_embd)?;
1895            e.copy_view_into(
1896                &mut h_rows,
1897                0,
1898                &hv.slice(done * n_embd..(done + tc) * n_embd),
1899                tc * n_embd,
1900            )?;
1901            let mut h_norm = e.uninit(tc * n_embd)?;
1902            e.rms_norm(
1903                &h_rows,
1904                mtp.hnorm.float_data(),
1905                &mut h_norm,
1906                n_embd,
1907                tc,
1908                eps,
1909            )?;
1910            // concat rows [tc, 2*n_embd] = [enorm ; hnorm] — two strided placements.
1911            let mut concat = e.uninit(tc * 2 * n_embd)?;
1912            e.place_rows_strided(&e_norm, &mut concat, n_embd, tc, 2 * n_embd, 0)?;
1913            e.place_rows_strided(&h_norm, &mut concat, n_embd, tc, 2 * n_embd, n_embd)?;
1914            let inp_sa = e.matmul(&mtp.eh_proj, &concat, tc)?;
1915            let mut a_norm = e.uninit(tc * n_embd)?;
1916            e.rms_norm(
1917                &inp_sa,
1918                mtp.attn_norm.float_data(),
1919                &mut a_norm,
1920                n_embd,
1921                tc,
1922                eps,
1923            )?;
1924            let pos: Vec<i32> = (done as i32..(done + tc) as i32).collect();
1925            let pos_d = e.htod_i32(&pos)?;
1926            let _ = self.mla_attn_cached(e, mla, &a_norm, &pos_d, tc, il, cache)?;
1927            done += tc;
1928        }
1929        Ok(())
1930    }
1931
1932    /// Row `row` of a `[rows, n_embd]` device stack, copied into its own `[n_embd]` buffer
1933    /// (the MTP `h_seed` handoff shape).
1934    fn glm5_seed_row(
1935        &self,
1936        e: &Engine,
1937        src: &CudaSlice<f32>,
1938        rows: usize,
1939        row: usize,
1940    ) -> Res<CudaSlice<f32>> {
1941        let n_embd = self.cfg.n_embd as usize;
1942        let stack = e.view(src, rows * n_embd);
1943        let view = stack.slice(row * n_embd..(row + 1) * n_embd);
1944        let mut seed = e.uninit(n_embd)?;
1945        e.copy_view_into(&mut seed, 0, &view, n_embd)?;
1946        Ok(seed)
1947    }
1948
1949    /// SERVED-SESSION ENTRY (lane/glm5-spec-routing): prime the prompt, warm the MTP plane,
1950    /// draw the boundary token, and hand back a [`Glm5SpecSession`] the worker bursts.
1951    ///
1952    /// `sampling`: `None` / `temp <= 0` = the greedy byte-contract route (the instrument);
1953    /// `Some` with `temp > 0` = the sampled route — the boundary token, the draft chain and
1954    /// the accept walk all draw through the session's own Philox counters (`sctr` device
1955    /// events, `uctr` host accept-test uniforms via `spec::host_u01`, tag 0xFFFF_FFFE), so a
1956    /// session's randomness never repeats across bursts (the session-continuity law).
1957    /// PENALIZED requests (greedy or sampled) open only under `MEMRA_SPEC_PENALTY=1`
1958    /// (`glm5_penalty_admit`; doc on [`glm5_spec_penalty_on`]) and refuse loudly otherwise —
1959    /// worker admission keeps them on the plain path while the arm is dark.
1960    pub fn glm5_spec_session_new(
1961        &self,
1962        e: &Engine,
1963        prompt: &[u32],
1964        ctx_cap: usize,
1965        sampling: Option<SpecSampling>,
1966    ) -> Res<Glm5SpecSession> {
1967        if self.hyper.is_none() {
1968            return Err("generate_spec_glm5 requires a HyperConnections trunk".into());
1969        }
1970        // Two parallel/spec programs on one model never silently coexist unless the
1971        // composition is EXPLICITLY armed: the spec x TP verify/rollback wiring
1972        // (lane/glm5-composition) exists and is rig-gated, but it has zero real-artifact
1973        // receipts, so sessions on a SHARDED model stay co-refused unless
1974        // MEMRA_GLM5_SPEC_TP=1 lifts the refusal (default OFF by design — the FLAGS row).
1975        // The predicate is the MODEL's own sharding (the same per-layer truth the verify
1976        // walk keys on), never the MEMRA_GLM5_TP env: sharding is a load-time property,
1977        // and an env read here is bypassable after load (set/load/unset) and spuriously
1978        // refuses an UNSHARDED model in a process that still carries the env (the #80
1979        // review's confirmed finding).
1980        let tp_sharded = self.layers.iter().any(|l| match &l.mixer {
1981            Mixer::Kda(la) => la.tp.is_some(),
1982            Mixer::Mla(mla) => mla.tp.is_some(),
1983            _ => false,
1984        });
1985        if tp_sharded {
1986            if !glm5_spec_tp_on() {
1987                return Err(
1988                    "glm5 spec is co-refused on a MEMRA_GLM5_TP-sharded model: set \
1989                     MEMRA_GLM5_SPEC_TP=1 to run the gated spec x TP composition \
1990                     (default OFF — zero real-artifact receipts; every other admission \
1991                     law still holds)"
1992                        .into(),
1993                );
1994            }
1995            if !glm5_verify_batch_on() {
1996                return Err("MEMRA_GLM5_SPEC_TP=1 requires the BATCHED verify walk \
1997                     (MEMRA_GLM5_VERIFY_BATCH must not be 0): the per-row rollback seam \
1998                     carries no TP arm"
1999                    .into());
2000            }
2001            // The ARMED announce prints AFTER the last admission law below — a session
2002            // refused later (draft source, ppN qualification, penalties, ctx) must never
2003            // log the engagement receipt (the fleet's prints-ARMED-then-serves-plain
2004            // trap class; rig-gates/03 caught SF3 logging it on a refused session).
2005        }
2006        // DRAFT-SOURCE SELECTION — through the GENERAL law
2007        // ([`crate::spec::resolve_draft_source_kind`], lane/glm5-extract2): DFlash2 when the
2008        // drafter is loaded (MEMRA_GLM5_DFLASH — it wins over a co-loaded MTP head, the boot
2009        // receipt states the selection), native MTP otherwise; neither = the loud refusal
2010        // below, whose bytes name the glm5 flags because the FAMILY owns the how-to-arm text.
2011        // The native MTP head is NOT required for the DFlash2 source (the q38 pattern).
2012        let dflash_src = self.glm5_dflash.as_ref();
2013        let source_kind = crate::spec::resolve_draft_source_kind(
2014            self.plan.draft_source,
2015            self.mtp.is_some(),
2016            dflash_src.is_some(),
2017        )
2018        .map_err(|why| {
2019            // The general law says WHY there is no usable source; the family owns the
2020            // how-to-arm text. Composed so the sentence is true on BOTH of the law's refusal
2021            // branches (nothing loaded / a head loaded under a plan that does not claim it),
2022            // rather than asserting "requires a draft source" at an operator whom the bracket
2023            // then tells a head IS loaded.
2024            format!(
2025                "generate_spec_glm5 cannot select a draft source ({why}). Arm one: the \
2026                 embedded MTP head (MEMRA_GLM5_MTP=1; a full MoE layer, unloaded by default) \
2027                 or the DFlash2 drafter (MEMRA_GLM5_DFLASH=<dir-or-hf-spec>)"
2028            )
2029        })?;
2030        if prompt.len() < 2 {
2031            return Err(
2032                "generate_spec_glm5 needs a prompt of >= 2 tokens (the MTP plane warms on \
2033                 (token[i+1], hidden[i]) pairs)"
2034                    .into(),
2035            );
2036        }
2037        if dflash_src.is_none() && crate::spec::spec_hpost() {
2038            // MTP-carrier-specific refusal: the DFlash2 source consumes tapped trunk
2039            // features, not the h_seed carrier, so the flag has nothing to flip there.
2040            return Err(
2041                "generate_spec_glm5 has no MEMRA_SPEC_HPOST arm: the flag flips the MTP \
2042                 carrier to the post-norm hidden, but this loop seeds every committed pair \
2043                 from the trunk's PRE-output_norm collapsed rows (LANE.md §A). Mixing the \
2044                 two silently degrades drafts; the HPOST twin needs its own gate before it \
2045                 may run"
2046                    .into(),
2047            );
2048        }
2049        // ppN split (lane/glm5-ppn-verify): the verify walk, the rollback and the MTP
2050        // chain all run under the split now — but an UNQUALIFIED pipeline rewrite still
2051        // refuses loudly at the session seam, before any cache is allocated over
2052        // stage-sharded weights (worker admission additionally bounds the stage count to
2053        // the gated set; see glm5_spec_capable).
2054        if crate::pp::pp_cuts(self.layers.len()).is_some()
2055            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
2056        {
2057            return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2058        }
2059        // PENALTY ARM admission (lane/spec-exclusions-20260902): a penalized request opens
2060        // only under MEMRA_SPEC_PENALTY=1; `pen` rides the session for every round, and
2061        // `sampling` keeps its route meaning (`None` = greedy verify) — a greedy request
2062        // with penalties is `sampling: None, pen: Some`.
2063        let pen = glm5_penalty_admit(sampling.as_ref())?;
2064        let sampling = sampling.filter(|sp| sp.temp > 0.0);
2065        // Room for the prompt, the anchor row and at least one verify round.
2066        if prompt.len() + 4 > ctx_cap {
2067            return Err(format!(
2068                "glm5 spec session needs ctx for prompt {} + anchor + one verify round, \
2069                 cap {ctx_cap}",
2070                prompt.len()
2071            )
2072            .into());
2073        }
2074        let n_vocab = self.output.out_features();
2075        // FR-SPEC TRIM (module doc): a loaded `MEMRA_FRSPEC_TRIM` artifact means the draft
2076        // head projects over gathered top-N rows and every draft pick is a RANK id that
2077        // must remap through d2t to the true vocabulary BEFORE it is chained or verified.
2078        // The verify walk stays full-vocab regardless — the invariant under gate.
2079        if let Some(map) = self.glm5_d2t() {
2080            if map.iter().any(|&t| t as usize >= n_vocab) {
2081                return Err(format!(
2082                    "glm5 FR-Spec d2t carries a token id >= n_vocab {n_vocab} — the ranks \
2083                     artifact was minted for a different vocabulary"
2084                )
2085                .into());
2086            }
2087            // Engagement receipt (the dspark trim receipt's shape): the server-log line
2088            // the trim arm's per-session engagement is verified by.
2089            eprintln!("{}", self.glm5_trim_engagement_line(map));
2090        }
2091        // Confidence-gate engagement receipt (loop-port 2; the deploy-gate greps this —
2092        // never-serve-greedy law's receipt discipline): armed iff MEMRA_SPEC_PMIN > 0.
2093        if glm5_pmin() > 0.0 {
2094            eprintln!(
2095                "[glm5-spec] draft confidence gate armed: PMIN={:.3} PMIN0={} (native \
2096                 chain p-of-pick; DFlash2 selector-q tau-slot truncation)",
2097                glm5_pmin(),
2098                glm5_pmin0() as u8,
2099            );
2100        }
2101        // The MTP plane index is a NATIVE-arm need; the DFlash2 source never touches the
2102        // plane (it still allocates below — plan-structural, the named cost in the module
2103        // doc — but nothing reads or resets it).
2104        let mtp_il = match dflash_src {
2105            Some(_) => None,
2106            None => Some(
2107                self.plan
2108                    .mtp_blocks
2109                    .first()
2110                    .ok_or("ModelPlan declares no MTP block")?
2111                    .layer
2112                    .index as usize,
2113            ),
2114        };
2115        // FIRST-TOKEN PROFILE (lane/b200-spec-ttft-20260902, `MEMRA_SPEC_PROF=1`): the
2116        // head engine resolves before any device work (a pure placement lookup), so the
2117        // clock can bound every creation phase on both the primary and head streams.
2118        let eh = self.glm5_head_engine(e)?;
2119        let mut prof = spec_prof_on().then(SpecFirstTokenProf::default);
2120        let mut pclk = prof.as_ref().map(|_| ProfClock::start(e, eh));
2121        if let Some(pf) = prof.as_mut() {
2122            pf.free_mb_before = self.glm5_free_mb(e);
2123        }
2124        // Stage-owned allocation under a split (each layer's planes on its stage's device,
2125        // trailing MTP plane on the last stage); door shut = plain `Cache::new_planned`.
2126        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, ctx_cap)?;
2127        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2128            pf.cache_alloc_ms = ck.lap(e, eh);
2129        }
2130
2131        // ---- prime, boundary token, draft-source warm over the prompt ----
2132        // `prime_cache` routes to its own ppN twin under the split; `hiddens` is owned by
2133        // the LAST stage's engine (its published contract) — exactly where the MTP chain
2134        // below runs, so the warm consumes it with no device bounce. DFlash2 source: the
2135        // prime walk fills the armed HcTapSink with every prompt row's contracted tap
2136        // features (the drafter's context; round 1 ingests them into its own KV).
2137        let plen = prompt.len();
2138        let n_embd = self.cfg.n_embd as usize;
2139        let tap_layers = match dflash_src {
2140            Some(dr) => Some(glm5_dflash_tap_layers(&dr.draft, self.layers.len())?),
2141            None => None,
2142        };
2143        // DRAFTER PRIME ARM (lane/spec-route-depth-20260902): the chunked arm primes the
2144        // trunk per engine chunk and ingests each chunk's taps into the drafter KV right
2145        // after that chunk's walk (`glm5_draft_prime_chunked`); the eager arm (the pre-lane
2146        // literal) primes once over a whole-prompt host sink and leaves the ingest to
2147        // round 1. `hidden_rows` is the row count of the returned `hiddens` stack (the whole
2148        // prompt on the eager arm, the LAST chunk on the chunked arm) — the boundary
2149        // capture indexes its last row through it.
2150        let mut v2_kv: Option<DflashKv> = None;
2151        let (logits0, hiddens, hidden_rows) = match (dflash_src, tap_layers.as_ref()) {
2152            (Some(dr), Some(taps)) if glm5_draft_taps_device_on() => {
2153                // DEVICE-RESIDENT ARM (doc on `glm5_draft_taps_device_on`): ONE whole-prompt
2154                // prime; the ingest happens inside it at every range boundary.
2155                let kv = DflashKv::new(eh, &dr.draft.cfg, ctx_cap)?;
2156                if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2157                    pf.draft_alloc_ms = ck.lap(e, eh);
2158                    pf.draft_kv_mb = dflash_kv_bytes(&dr.draft.cfg, ctx_cap) as f64 / 1e6;
2159                }
2160                let ring = crate::hybrid_forward::hyper_prime_call_rows(
2161                    plen,
2162                    self.layers.len(),
2163                    self.gdn_prime_grid_on(),
2164                );
2165                let mut sink = HcTapSink::new_device_staged_at(taps.clone(), n_embd, ring, 0);
2166                sink.ingest_state = Some(Box::new(Glm5DraftPrimeInflight {
2167                    kv,
2168                    taps: taps.clone(),
2169                    n_embd,
2170                    ring,
2171                    stage: (0..taps.len()).map(|_| None).collect(),
2172                    rows_dev: None,
2173                    prof_on: prof.is_some(),
2174                    copy_ms: 0.0,
2175                    feat_ms: 0.0,
2176                    kv_ms: 0.0,
2177                    chunks: 0,
2178                }));
2179                cache.hc_taps = Some(sink);
2180                let (l, _seed, h) = self.prime_cache(e, prompt, &mut cache, 0)?;
2181                let walk_ms = pclk.as_mut().map(|ck| ck.lap(e, eh));
2182                let mut sink = cache
2183                    .hc_taps
2184                    .take()
2185                    .ok_or("device-resident drafter prime: tap sink vanished")?;
2186                let st = sink
2187                    .ingest_state
2188                    .take()
2189                    .ok_or("device-resident drafter prime: ingest state vanished")?
2190                    .downcast::<Glm5DraftPrimeInflight>()
2191                    .map_err(|_| "device-resident drafter prime: ingest state of the wrong type")?;
2192                let st = *st;
2193                if st.kv.len != plen {
2194                    return Err(format!(
2195                        "device-resident drafter prime covered {} of {plen} prompt rows \
2196                         (the prime's range loop must hand every range to the ingest)",
2197                        st.kv.len
2198                    )
2199                    .into());
2200                }
2201                if let (Some(pf), Some(walk)) = (prof.as_mut(), walk_ms) {
2202                    // The ingest ran INSIDE the prime walk: split it back out so `prime`
2203                    // stays the trunk's share and `draft_prime` the drafter's.
2204                    let ingest = st.copy_ms + st.feat_ms + st.kv_ms;
2205                    pf.prime_ms = (walk - ingest).max(0.0);
2206                    pf.draft_prime_ms = ingest;
2207                    pf.draft_prime_h2d_ms = st.copy_ms;
2208                    pf.draft_prime_feat_ms = st.feat_ms;
2209                    pf.draft_prime_kv_ms = st.kv_ms;
2210                    pf.draft_prime_rows = plen;
2211                    pf.draft_prime_chunks = st.chunks;
2212                    pf.draft_prime_arm = "device";
2213                }
2214                v2_kv = Some(st.kv);
2215                (l, h, plen)
2216            }
2217            (Some(dr), Some(taps)) if glm5_draft_prime_v2_on() => {
2218                let mut kv = DflashKv::new(eh, &dr.draft.cfg, ctx_cap)?;
2219                if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2220                    pf.draft_alloc_ms = ck.lap(e, eh);
2221                    pf.draft_kv_mb = dflash_kv_bytes(&dr.draft.cfg, ctx_cap) as f64 / 1e6;
2222                }
2223                let out = self.glm5_draft_prime_chunked(
2224                    e,
2225                    eh,
2226                    &dr.draft,
2227                    prompt,
2228                    &mut cache,
2229                    &mut kv,
2230                    taps,
2231                    prof.as_mut(),
2232                    pclk.as_mut(),
2233                )?;
2234                v2_kv = Some(kv);
2235                out
2236            }
2237            _ => {
2238                if let Some(taps) = tap_layers.as_ref() {
2239                    let t_sink = std::time::Instant::now();
2240                    cache.hc_taps = Some(HcTapSink::new(taps.clone(), n_embd, plen));
2241                    if let Some(pf) = prof.as_mut() {
2242                        pf.sink_alloc_ms = t_sink.elapsed().as_secs_f64() * 1e3;
2243                    }
2244                }
2245                let (l, _seed, h) = self.prime_cache(e, prompt, &mut cache, 0)?;
2246                if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2247                    pf.prime_ms = ck.lap(e, eh);
2248                    pf.prime_tap_dtoh_ms = cache
2249                        .hc_taps
2250                        .as_ref()
2251                        .map(|sk| sk.dtoh_ns as f64 / 1e6)
2252                        .unwrap_or(0.0);
2253                }
2254                (l, h, plen)
2255            }
2256        };
2257        // Prompt-boundary capture (lane/glm5-prefix-latent2): taken NOW — after the prime
2258        // filled every plane to the boundary, before the anchor/draft machinery below and
2259        // before any burst mutates the conv/ssm state or laps the tail ring. DFlash2-only:
2260        // the native arm's plane fill moves the MTP latent layer past the boundary before a
2261        // capture could be taken, and restore refuses the native source anyway. A refusal
2262        // drops the capture loudly and the session serves regardless.
2263        let prefix_capture = if glm5_spec_prefix_on() && dflash_src.is_some() {
2264            self.glm5_prefix_boundary_capture(e, eh, &cache, &logits0, &hiddens, plen, hidden_rows)
2265        } else {
2266            None
2267        };
2268        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2269            pf.capture_ms = ck.lap(e, eh);
2270        }
2271        // The penalty window spans the SESSION (`pen_window_seed`): the prompt now, every
2272        // committed token from here. Empty with penalties off — the anchor rule below is
2273        // then the pre-lane literal in both regimes.
2274        let pen_hist = glm5_pen_window_seed(pen.as_ref(), prompt);
2275        let mut sctr = 0u32;
2276        let anchor = glm5_anchor(
2277            eh,
2278            &logits0,
2279            sampling.as_ref(),
2280            pen.as_ref(),
2281            &pen_hist,
2282            &mut sctr,
2283            "glm5-prime",
2284        )?;
2285        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2286            pf.anchor_ms = ck.lap(e, eh);
2287        }
2288
2289        // Keyed on the KIND the general law returned, not on a second local re-derivation:
2290        // a seam whose answer is recomputed by its consumer is decoration. (`tap_layers` is
2291        // `Some` exactly when `dflash_src` is, and the law returns `Dflash2` exactly then, so
2292        // this is the same program the pre-seam code ran — the `_` arm's refusal is the
2293        // never-taken proof of that rather than a silent fallback.)
2294        let (mut draft, pending) = match (source_kind, dflash_src, tap_layers) {
2295            (crate::spec::DraftSourceKind::Dflash2, Some(dr), Some(taps)) => {
2296                if let Some(kv) = v2_kv.take() {
2297                    // Chunked arm: the KV already holds every prompt row (kv.len == plen);
2298                    // round 1 finds nothing pending and walks straight to its block forward.
2299                    debug_assert_eq!(kv.len, plen, "chunked drafter prime must cover the prompt");
2300                    (
2301                        Glm5DraftState::Dflash2 {
2302                            kv,
2303                            pending: Vec::new(),
2304                            taps,
2305                        },
2306                        Vec::new(),
2307                    )
2308                } else {
2309                    let sink = cache
2310                        .hc_taps
2311                        .take()
2312                        .ok_or("glm5 dflash prime tap sink vanished")?;
2313                    // Drafter ctx KV on the HEAD engine (where the drafter weights loaded and
2314                    // every round's chain runs); prompt feature rows ride `pending` so round 1
2315                    // ingests them through the one chunked path.
2316                    let kv = DflashKv::new(eh, &dr.draft.cfg, ctx_cap)?;
2317                    if let Some(pf) = prof.as_mut() {
2318                        pf.draft_kv_mb = dflash_kv_bytes(&dr.draft.cfg, ctx_cap) as f64 / 1e6;
2319                    }
2320                    (
2321                        Glm5DraftState::Dflash2 {
2322                            kv,
2323                            pending: sink.rows,
2324                            taps,
2325                        },
2326                        Vec::new(),
2327                    )
2328                }
2329            }
2330            (crate::spec::DraftSourceKind::Dflash2, _, _) => {
2331                return Err(
2332                    "the draft-source law selected DFlash2 but this session resolved no tap \
2333                     layers — a load-path bug, refused instead of silently drafting from the \
2334                     MTP plane (a VANISHED tap sink is a different failure, caught by name in \
2335                     the Dflash2 arm itself)"
2336                        .into(),
2337                );
2338            }
2339            (crate::spec::DraftSourceKind::NativeMtp, _, _) => {
2340                // BATCHED PLANE WARM (loop-port fold-in — the map's #4, the spec.rs
2341                // `mtp_kv_fill_all` pattern re-aimed at the MLA plane): pairs
2342                // (prompt[i+1], h_i) at plane pos i, i in 0..P-1, filled in CHUNKED
2343                // t-parallel passes instead of P-1 sequential full-block forwards. The
2344                // sequential warm ran ~400 tok/s — the measured +2.5 s TTFT per 1k
2345                // prompt tokens, spec-battery flip condition 1 by name. MTP rows are
2346                // INDEPENDENT given the trunk hiddens (no row-to-row recurrence — the
2347                // plane is the only carrier), so the fill is exact in structure; the
2348                // t>1 attention takes the prime-class program, which can only move
2349                // DRAFTS, never output (verify arbitrates; the byte-identity batteries
2350                // stay the proof).
2351                self.glm5_mtp_plane_fill(eh, &prompt[1..], &hiddens, plen - 1, &mut cache)?;
2352                // pending = committed (token, h_seed) pairs not yet fed to the MTP plane.
2353                // The LAST pair's logits are the next round's first draft — the re-warm
2354                // doubles as draft 1.
2355                let pending = vec![(anchor, self.glm5_seed_row(eh, &hiddens, plen, plen - 1)?)];
2356                (Glm5DraftState::NativeMtp, pending)
2357            }
2358        };
2359        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
2360            pf.draft_alloc_ms += ck.lap(e, eh);
2361        }
2362        // EAGER-ARM INGEST AT CREATION (doc on `glm5_draft_prime_lazy_on`): the prompt's tap
2363        // rows go into the drafter KV NOW, before the session (and its anchor) is handed to
2364        // the worker, unless the lazy seam asks for the round-1 placement. The chunked arm
2365        // arrives here with nothing pending.
2366        if !glm5_draft_prime_lazy_on()
2367            && let (Glm5DraftState::Dflash2 { kv, pending, taps }, Some(dr)) =
2368                (&mut draft, dflash_src)
2369            && !pending.is_empty()
2370        {
2371            let rows = std::mem::take(pending);
2372            let stats = self.glm5_dflash_ingest_rows(
2373                eh,
2374                &dr.draft,
2375                kv,
2376                &rows,
2377                taps.len() * n_embd,
2378                pclk.as_mut(),
2379            )?;
2380            if let Some(pf) = prof.as_mut() {
2381                stats.write(pf, "eager");
2382            }
2383        }
2384        if let Some(pf) = prof.as_mut() {
2385            pf.free_mb_after = self.glm5_free_mb(e);
2386        }
2387        // Composition engagement receipt — printed immediately before the session is
2388        // RETURNED (after every admission law, the d2t vocabulary check, the cache
2389        // allocation and the prompt prime), so a grep for this line counts sessions that
2390        // actually opened; a refusal or a failure anywhere above never logs it (the #82
2391        // review moved it here after finding four fallible steps below its first home).
2392        if tp_sharded {
2393            eprintln!(
2394                "[glm5-spec] spec x TP composition ARMED (MEMRA_GLM5_SPEC_TP=1): verify \
2395                 rows ride the TP shards; rollback restores per-rank planes \
2396                 performance_claim=false"
2397            );
2398        }
2399        Ok(Glm5SpecSession {
2400            cache,
2401            committed: prompt.to_vec(),
2402            anchor,
2403            anchor_emitted: false,
2404            pending,
2405            draft,
2406            sampling,
2407            pen,
2408            pen_hist,
2409            sctr,
2410            uctr: 0,
2411            rounds: 0,
2412            rank_trimmed_rounds: 0,
2413            done: false,
2414            max_ctx: ctx_cap,
2415            mtp_il,
2416            prefix_capture,
2417            prof_rounds: prof.as_ref().map(|_| SpecRoundsLog::default()),
2418            prof,
2419        })
2420    }
2421
2422    /// Range-begin hook of the device-resident drafter prime (called by the prime's range
2423    /// loop, `hybrid_forward.rs`): anchor the chunk ring at the range's first row so the
2424    /// walk's `base - origin` lands in `[0, ring)`. No-op on every other sink.
2425    pub(crate) fn glm5_taps_range_begin(&self, cache: &mut Cache, start: usize) {
2426        if let Some(sink) = cache.hc_taps.as_mut()
2427            && sink.ingest_state.is_some()
2428        {
2429            sink.origin = start;
2430        }
2431    }
2432
2433    /// Range-done hook of the device-resident drafter prime: the range's tap planes are
2434    /// complete on their writing devices (the range call returned host logits, so the last
2435    /// stage drained, and every earlier stage's rows were consumed by it), so ingest them
2436    /// now. No-op on every other sink; a state of another type is put back untouched.
2437    pub(crate) fn glm5_taps_range_done(
2438        &self,
2439        e: &Engine,
2440        cache: &mut Cache,
2441        start: usize,
2442        end: usize,
2443    ) -> Res<()> {
2444        let Some(sink) = cache.hc_taps.as_mut() else {
2445            return Ok(());
2446        };
2447        let Some(state) = sink.ingest_state.take() else {
2448            return Ok(());
2449        };
2450        let mut st = match state.downcast::<Glm5DraftPrimeInflight>() {
2451            Ok(st) => st,
2452            Err(other) => {
2453                sink.ingest_state = Some(other);
2454                return Ok(());
2455            }
2456        };
2457        let res = self.glm5_taps_ingest_range(e, sink, &mut st, start, end);
2458        sink.ingest_state = Some(st);
2459        res
2460    }
2461
2462    /// One range of the device-resident drafter prime: per tap slot, the plane on its
2463    /// writing device is interleaved into the head-device fc input with a 2D copy (after a
2464    /// peer copy into head-device staging when the slot's stage is another device; the
2465    /// source stream is drained before the plane is read), then `ctx_features` +
2466    /// `ingest_ctx` at the range width append the rows to the drafter KV.
2467    fn glm5_taps_ingest_range(
2468        &self,
2469        e: &Engine,
2470        sink: &mut HcTapSink,
2471        st: &mut Glm5DraftPrimeInflight,
2472        start: usize,
2473        end: usize,
2474    ) -> Res<()> {
2475        let t = end - start;
2476        if t == 0 {
2477            return Ok(());
2478        }
2479        if t > st.ring {
2480            return Err(format!(
2481                "device-resident tap ingest: range {start}..{end} ({t} rows) exceeds the \
2482                 sink ring of {} rows",
2483                st.ring
2484            )
2485            .into());
2486        }
2487        if st.kv.len != start {
2488            return Err(format!(
2489                "device-resident tap ingest: drafter KV holds {} rows but the range starts \
2490                 at {start} (a range was skipped or handed twice)",
2491                st.kv.len
2492            )
2493            .into());
2494        }
2495        let eh = self.glm5_head_engine(e)?;
2496        let dr = self
2497            .glm5_dflash
2498            .as_ref()
2499            .ok_or("device-resident tap ingest without a loaded drafter")?;
2500        let h = st.n_embd;
2501        let n_taps = st.taps.len();
2502        let mut pclk = st.prof_on.then(|| ProfClock::start(e, eh));
2503        if st.rows_dev.is_none() {
2504            st.rows_dev = Some(eh.uninit(st.ring * n_taps * h)?);
2505        }
2506        let rows_dev = st.rows_dev.as_mut().expect("just filled");
2507        for (slot, &il) in st.taps.iter().enumerate() {
2508            let buf = sink.dev[slot].as_ref().ok_or_else(|| {
2509                format!(
2510                    "device-resident tap ingest: slot {slot} (layer {il}) was never written \
2511                     over rows {start}..{end}"
2512                )
2513            })?;
2514            let es = self.glm5_tap_slot_engine(e, il)?;
2515            if es.ctx().ordinal() == eh.ctx().ordinal() {
2516                eh.copy_2d_dtod_async(rows_dev, slot * h, n_taps * h, buf, h, h, t)?;
2517            } else {
2518                if st.stage[slot].is_none() {
2519                    st.stage[slot] = Some(eh.uninit(st.ring * h)?);
2520                }
2521                let staging = st.stage[slot].as_mut().expect("just filled");
2522                eh.copy_peer_from_async(staging, es, buf, t * h)?;
2523                es.stream().synchronize()?;
2524                eh.copy_2d_dtod_async(rows_dev, slot * h, n_taps * h, staging, h, h, t)?;
2525            }
2526        }
2527        if let Some(ck) = pclk.as_mut() {
2528            st.copy_ms += ck.lap(e, eh);
2529        }
2530        let feats = dr.draft.ctx_features(eh, rows_dev, t)?;
2531        if let Some(ck) = pclk.as_mut() {
2532            st.feat_ms += ck.lap(e, eh);
2533        }
2534        let pos: Vec<i32> = ((st.kv.len as i32)..(st.kv.len + t) as i32).collect();
2535        dr.draft.ingest_ctx(eh, &mut st.kv, &feats, &pos, t)?;
2536        if let Some(ck) = pclk.as_mut() {
2537            st.kv_ms += ck.lap(e, eh);
2538        }
2539        st.chunks += 1;
2540        Ok(())
2541    }
2542
2543    /// The EAGER drafter ingest (one implementation for session creation and round 1):
2544    /// host feature rows `[n, row_w]` -> 256-row pageable HtoD chunks -> `ctx_features` ->
2545    /// `ingest_ctx`, appended at `kv.len`. Returns the profile buckets (all 0 with the
2546    /// clock absent).
2547    fn glm5_dflash_ingest_rows(
2548        &self,
2549        eh: &Engine,
2550        draft: &DflashDraft,
2551        kv: &mut DflashKv,
2552        rows: &[f32],
2553        row_w: usize,
2554        mut pclk: Option<&mut ProfClock>,
2555    ) -> Res<DraftIngestStats> {
2556        debug_assert_eq!(rows.len() % row_w, 0, "ragged feature rows");
2557        let n_new = rows.len() / row_w;
2558        let mut st = DraftIngestStats {
2559            rows: n_new,
2560            ..Default::default()
2561        };
2562        let mut r0 = 0usize;
2563        while r0 < n_new {
2564            let t_c = (n_new - r0).min(256);
2565            let chunk = eh.htod(&rows[r0 * row_w..(r0 + t_c) * row_w])?;
2566            if let Some(ck) = pclk.as_deref_mut() {
2567                st.h2d_ms += ck.lap(eh, eh);
2568            }
2569            let feats = draft.ctx_features(eh, &chunk, t_c)?;
2570            if let Some(ck) = pclk.as_deref_mut() {
2571                st.feat_ms += ck.lap(eh, eh);
2572            }
2573            let pos_c: Vec<i32> = ((kv.len as i32)..(kv.len + t_c) as i32).collect();
2574            draft.ingest_ctx(eh, kv, &feats, &pos_c, t_c)?;
2575            if let Some(ck) = pclk.as_deref_mut() {
2576                st.kv_ms += ck.lap(eh, eh);
2577            }
2578            r0 += t_c;
2579            st.chunks += 1;
2580        }
2581        Ok(st)
2582    }
2583
2584    /// Free device memory per device this model can hold state on (primary + ppN stages),
2585    /// for the depth profile's before/after samples.
2586    fn glm5_free_mb(&self, e: &Engine) -> Vec<(usize, u64)> {
2587        let mut out = vec![(e.ctx().ordinal(), e.free_mem_mb())];
2588        if let Ok(rt) = crate::pp::PpNRt::get(e) {
2589            for stage in 0..rt.n_stages() {
2590                let se = rt.engine(stage, e);
2591                let ord = se.ctx().ordinal();
2592                if !out.iter().any(|(o, _)| *o == ord) {
2593                    out.push((ord, se.free_mem_mb()));
2594                }
2595            }
2596        }
2597        out
2598    }
2599
2600    /// The engine that OWNS tap layer `il`'s device (the stage engine under a live split,
2601    /// the caller's engine otherwise) — the `glm5_tap_drain` placement rule, shared.
2602    fn glm5_tap_slot_engine<'e>(&self, e: &'e Engine, il: usize) -> Res<&'e Engine> {
2603        match crate::pp::pp_cuts(self.layers.len()) {
2604            Some(fence) if !crate::pp::pp2_streams_off() => {
2605                let rt = crate::pp::PpNRt::get(e)?;
2606                let stage = fence
2607                    .windows(2)
2608                    .position(|w| il >= w[0] && il < w[1])
2609                    .ok_or_else(|| format!("tap layer {il} outside every stage range"))?;
2610                Ok(rt.engine(stage, e))
2611            }
2612            _ => Ok(e),
2613        }
2614    }
2615
2616    /// CHUNKED DRAFTER PRIME (lane/spec-route-depth-20260902, `MEMRA_GLM5_DRAFT_PRIME_V2`):
2617    /// prime the trunk over the engine's own chunk schedule and ingest each chunk's tap
2618    /// rows into the drafter KV right after that chunk's walk. Returns the LAST chunk's
2619    /// (boundary logits, hidden stack, stack rows).
2620    ///
2621    /// TRUNK PROGRAM, unchanged: `prime_cache` over `[start, end)` with
2622    /// `queued_after = plen - end` is exactly the call the whole-prompt entry makes per
2623    /// range (`prime_cache_hyper`'s loop: `queued_after + (t - end)`, `seq_end` invariant),
2624    /// so the trunk cache after the loop is byte-identical to the eager arm's — the
2625    /// restored-suffix gates (11/12) already pin "continue == cold" for this call shape.
2626    ///
2627    /// DATA MOVEMENT, the whole point: per chunk, five device slot buffers (one per tap
2628    /// layer, on the writing stage's device) are read back into ONE pinned cacheable slot
2629    /// buffer (sync, `dtoh_f32_into_pinned`), CPU-interleaved into the pinned
2630    /// `[t, n_taps * hidden]` rows buffer (the fc input layout), uploaded ASYNC from pinned
2631    /// (`htod_f32_from_pinned_async`), then `ctx_features` + `ingest_ctx` at the chunk
2632    /// width. The host transient is two chunk-sized pinned buffers (335 MB + 67 MB at 4096
2633    /// rows) instead of the eager arm's `[prompt, 5 x hidden]` pageable Vec; the eager
2634    /// arm's 21 GB of pageable DtoH + 21 GB of pageable HtoD at 256k become pinned DMA.
2635    /// The ingest kernels are stream-ordered behind the upload; the chunk ends with one
2636    /// stream drain so the rows buffer can be refilled (the upload helper's contract).
2637    #[allow(clippy::too_many_arguments)]
2638    // allow: the parameter list is the prime contract (engines, drafter, prompt, cache,
2639    // KV, taps) plus the two profile instruments; bundling would hide which is which
2640    fn glm5_draft_prime_chunked(
2641        &self,
2642        e: &Engine,
2643        eh: &Engine,
2644        draft: &DflashDraft,
2645        prompt: &[u32],
2646        cache: &mut Cache,
2647        kv: &mut DflashKv,
2648        taps: &[usize],
2649        mut prof: Option<&mut SpecFirstTokenProf>,
2650        mut pclk: Option<&mut ProfClock>,
2651    ) -> Res<(Vec<f32>, CudaSlice<f32>, usize)> {
2652        let plen = prompt.len();
2653        let n_embd = self.cfg.n_embd as usize;
2654        let n_taps = taps.len();
2655        let ranges = crate::hybrid_forward::hyper_prime_ranges(
2656            plen,
2657            self.layers.len(),
2658            self.gdn_prime_grid_on(),
2659        );
2660        let t_max = ranges.iter().map(|&(a, b)| b - a).max().unwrap_or(0);
2661        let f32b = std::mem::size_of::<f32>();
2662        let mut slot_buf = crate::PinnedHostBuf::new(t_max * n_embd * f32b)?;
2663        let mut rows_buf = crate::PinnedHostBuf::new(t_max * n_taps * n_embd * f32b)?;
2664        let (mut prime_ms, mut h2d_ms, mut feat_ms, mut kv_ms) = (0f64, 0f64, 0f64, 0f64);
2665        let mut last: Option<(Vec<f32>, CudaSlice<f32>, usize)> = None;
2666        for &(start, end) in &ranges {
2667            let t = end - start;
2668            cache.hc_taps = Some(HcTapSink::new_device_staged_at(
2669                taps.to_vec(),
2670                n_embd,
2671                t,
2672                start,
2673            ));
2674            let (l, _seed, h) = self.prime_cache(e, &prompt[start..end], cache, plen - end)?;
2675            if let Some(ck) = pclk.as_deref_mut() {
2676                prime_ms += ck.lap(e, eh);
2677            }
2678            let mut sink = cache
2679                .hc_taps
2680                .take()
2681                .ok_or("chunked drafter prime: chunk tap sink vanished")?;
2682            // ---- drain: device slots -> pinned slot -> interleaved pinned rows ----
2683            for (slot, &il) in taps.iter().enumerate() {
2684                let buf = sink.dev[slot].take().ok_or_else(|| {
2685                    format!(
2686                        "chunked drafter prime: tap slot {slot} (layer {il}) was never \
2687                         written by the prime walk over rows {start}..{end}"
2688                    )
2689                })?;
2690                let es = self.glm5_tap_slot_engine(e, il)?;
2691                es.dtoh_f32_into_pinned(&buf, &mut slot_buf, t * n_embd)?;
2692                let src = pinned_f32(&slot_buf, t * n_embd);
2693                let dst = pinned_f32_mut(&mut rows_buf, t * n_taps * n_embd);
2694                for r in 0..t {
2695                    let d0 = (r * n_taps + slot) * n_embd;
2696                    dst[d0..d0 + n_embd].copy_from_slice(&src[r * n_embd..(r + 1) * n_embd]);
2697                }
2698            }
2699            let feats_in = eh.htod_f32_from_pinned_async(&rows_buf, t * n_taps * n_embd)?;
2700            if let Some(ck) = pclk.as_deref_mut() {
2701                h2d_ms += ck.lap(e, eh);
2702            }
2703            let feats = draft.ctx_features(eh, &feats_in, t)?;
2704            if let Some(ck) = pclk.as_deref_mut() {
2705                feat_ms += ck.lap(e, eh);
2706            }
2707            let pos: Vec<i32> = ((kv.len as i32)..(kv.len + t) as i32).collect();
2708            draft.ingest_ctx(eh, kv, &feats, &pos, t)?;
2709            // The rows buffer is refilled by the next chunk: drain the stream first (the
2710            // async-upload contract). Also bounds the kv bucket.
2711            eh.stream().synchronize()?;
2712            if let Some(ck) = pclk.as_deref_mut() {
2713                kv_ms += ck.lap(e, eh);
2714            }
2715            last = Some((l, h, t));
2716        }
2717        debug_assert_eq!(kv.len, plen, "chunked drafter prime must cover the prompt");
2718        if let Some(pf) = prof.as_mut() {
2719            pf.prime_ms = prime_ms;
2720            pf.draft_prime_ms = h2d_ms + feat_ms + kv_ms;
2721            pf.draft_prime_h2d_ms = h2d_ms;
2722            pf.draft_prime_feat_ms = feat_ms;
2723            pf.draft_prime_kv_ms = kv_ms;
2724            pf.draft_prime_rows = plen;
2725            pf.draft_prime_chunks = ranges.len();
2726            pf.draft_prime_arm = "chunked";
2727        }
2728        last.ok_or_else(|| "chunked drafter prime: empty prime schedule".into())
2729    }
2730
2731    /// Boundary capture for the DEFERRED prefix publication (lane/glm5-prefix-latent2,
2732    /// 2026-09-01): the generation-destroyed state at `pos == plen` — conv/ssm via
2733    /// `Cache::snapshot` (D2D copies), per-layer latent tails via `snapshot_tail`, the
2734    /// prime's boundary logits, the pre-output_norm boundary hidden. `None` (loud) on any
2735    /// refusal — a capture is an optimization the session must never fail on. The
2736    /// append-only planes (latent rows, final pool keys, full-attn KV) are NOT copied here:
2737    /// the worker slices them from the live cache at publish (`snapshot_plane_at`), which is
2738    /// legal because the glm5 verify rollback never truncates below the prime boundary.
2739    #[allow(clippy::too_many_arguments)]
2740    // allow: the parameter list is the capture contract (two engines, the cache, the
2741    // boundary logits and hidden stack, the boundary and the stack's own row count);
2742    // bundling would hide which row index is which
2743    fn glm5_prefix_boundary_capture(
2744        &self,
2745        e: &Engine,
2746        eh: &Engine,
2747        cache: &Cache,
2748        logits0: &[f32],
2749        hiddens: &CudaSlice<f32>,
2750        plen: usize,
2751        // Rows in `hiddens` (its last row is the boundary hidden): the whole prompt on the
2752        // eager cold prime, the last chunk on the chunked arm, the suffix on a restore.
2753        hidden_rows: usize,
2754    ) -> Option<crate::spec::SpecBoundaryCapture> {
2755        debug_assert_eq!(
2756            cache.pos, plen,
2757            "boundary capture must sit at the prime boundary"
2758        );
2759        let snap = match cache.snapshot(e) {
2760            Ok(s) => s,
2761            Err(err) => {
2762                eprintln!("[glm5-spec] prefix boundary capture SKIPPED (cache snapshot: {err})");
2763                return None;
2764            }
2765        };
2766        // The ONLY latent layer that may legitimately sit empty at the boundary is the
2767        // MTP/NextN plane (allocated, never executed by the trunk on the DFlash2 arm).
2768        // Identity-keyed, not length-keyed (PR #96 review round 2, finding 1): a TRUNK
2769        // plane empty at the boundary is a regression that must refuse the capture, or
2770        // three length-keyed layers downstream would each wave it through and reproduce
2771        // the parent lane's fabrication shape per-layer.
2772        let mtp_plane_il = self.plan.mtp_blocks.first().map(|b| b.layer.index as usize);
2773        let mut latent_tails = Vec::with_capacity(cache.latent.len());
2774        for (il, l) in cache.latent.iter().enumerate() {
2775            match l {
2776                Some(l) if l.len == 0 && cache.pos > 0 => {
2777                    if Some(il) != mtp_plane_il {
2778                        eprintln!(
2779                            "[glm5-spec] prefix boundary capture SKIPPED (trunk latent \
2780                             layer {il} is EMPTY at the boundary — not the MTP plane; a \
2781                             capture would publish an absent history for a live layer)"
2782                        );
2783                        return None;
2784                    }
2785                    latent_tails.push(None)
2786                }
2787                Some(l) => {
2788                    if l.len != plen {
2789                        eprintln!(
2790                            "[glm5-spec] prefix boundary capture SKIPPED (latent layer {il} \
2791                             len {} != boundary {plen})",
2792                            l.len,
2793                        );
2794                        return None;
2795                    }
2796                    match l.snapshot_tail(e) {
2797                        Ok(t) => latent_tails.push(Some(t)),
2798                        Err(err) => {
2799                            eprintln!(
2800                                "[glm5-spec] prefix boundary capture SKIPPED (latent layer \
2801                                 {il}: {err})"
2802                            );
2803                            return None;
2804                        }
2805                    }
2806                }
2807                None => latent_tails.push(None),
2808            }
2809        }
2810        let last_h = crate::spec::capture_boundary_hidden(
2811            eh,
2812            hiddens,
2813            hidden_rows,
2814            self.cfg.n_embd as usize,
2815        );
2816        Some(crate::spec::SpecBoundaryCapture {
2817            snap,
2818            pos: plen,
2819            logits: logits0.to_vec(),
2820            last_h,
2821            latent_tails,
2822        })
2823    }
2824
2825    /// Re-arm a glm5 spec session from a RESTORED trunk cache plus a published DFlash2
2826    /// drafter tail — the glm5 twin of `dspark_spec_session_from_restored`, EXTENDED with
2827    /// the suffix prime the multi-turn shape needs (lane/glm5-prefix-latent2, 2026-09-01).
2828    ///
2829    /// WHY IT IS EQUIVALENT TO A COLD PRIME, field by field:
2830    /// * `cache` — the caller's whole-entry restored trunk cache at `fed.len()` (KDA
2831    ///   conv/ssm + MLA latent rows + kpool keys + tail ring, the parent lane's restore),
2832    ///   and the SUFFIX primes onto it through `prime_cache` — the same continuation
2833    ///   program every chunk after the first of a cold prime runs.
2834    /// * drafter — `dkv` is rebuilt from the entry's tail into the SAME absolute rows
2835    ///   (`DflashKv::from_tail`, caller-side while the prefix cache is borrowable), and the
2836    ///   suffix's tap rows ride `pending` exactly as a cold session's prompt rows do — the
2837    ///   drafter's context is the committed tokens either way (a truncated tail below the
2838    ///   window can only move ACCEPTANCE, never output: verify arbitrates).
2839    /// * anchor — drawn from the SUFFIX prime's boundary logits with the request's own
2840    ///   sampler, exactly the cold composition; Philox counters fresh (a restore is a NEW
2841    ///   session — the frspec continuity law, the dspark restore's own convention).
2842    /// * republish — with `glm5_spec_prefix_on()` the session takes a NEW boundary capture
2843    ///   at `fed + suffix`, so the next turn hits a DEEPER prefix (the
2844    ///   MEMRA_SPEC_RESTORE_REPUBLISH posture; the worker's has_key dedupe drops equals).
2845    ///
2846    /// Refuses (never asserts) whenever the restored halves disagree — a caller that gets
2847    /// `Err` serves the plain hit (correct, slower).
2848    #[allow(clippy::too_many_arguments)]
2849    pub fn glm5_spec_session_from_restored(
2850        &self,
2851        e: &Engine,
2852        mut cache: Cache,
2853        fed: &[u32],
2854        suffix: &[u32],
2855        // The ENTRY's boundary logits (`ReuseEntry::last_logits`), read ONLY on the
2856        // full-cover arm (`suffix.is_empty()`), where there is no suffix prime to draw the
2857        // anchor from. A suffix-bearing restore ignores it and uses the prime's own row.
2858        boundary_logits: &[f32],
2859        dkv: DflashKv,
2860        ctx_cap: usize,
2861        sampling: Option<SpecSampling>,
2862    ) -> Res<Glm5SpecSession> {
2863        if self.hyper.is_none() {
2864            return Err("glm5_spec_session_from_restored requires a HyperConnections trunk".into());
2865        }
2866        // The composition laws a cold session enforces hold here too, fail-closed and by
2867        // name — a restored session must never be the door around an admission law.
2868        let tp_sharded = self.layers.iter().any(|l| match &l.mixer {
2869            Mixer::Kda(la) => la.tp.is_some(),
2870            Mixer::Mla(mla) => mla.tp.is_some(),
2871            _ => false,
2872        });
2873        if tp_sharded {
2874            return Err(
2875                "restored glm5 spec sessions carry no TP arm (the spec x TP composition is \
2876                 cold-session gated only); the plain hit serves"
2877                    .into(),
2878            );
2879        }
2880        if crate::pp::pp_cuts(self.layers.len()).is_some()
2881            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
2882        {
2883            return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2884        }
2885        // DFlash2 source ONLY: the native MTP plane fill consumes trunk hiddens the restored
2886        // range does not have (re-running the trunk over it would be a second prime — the
2887        // whole cost this restore exists to avoid).
2888        let dr = self.glm5_dflash.as_ref().ok_or(
2889            "restored glm5 spec sessions require the DFlash2 drafter (MEMRA_GLM5_DFLASH): \
2890             the native MTP plane cannot be re-warmed from restored KV",
2891        )?;
2892        let source = crate::spec::resolve_draft_source_kind(
2893            self.plan.draft_source,
2894            self.mtp.is_some(),
2895            true,
2896        )
2897        .map_err(|why| format!("restored glm5 spec session has no draft source ({why})"))?;
2898        if !matches!(source, crate::spec::DraftSourceKind::Dflash2) {
2899            return Err(
2900                "restored glm5 spec sessions require the DFlash2 draft source; the plan \
2901                 selected another"
2902                    .into(),
2903            );
2904        }
2905        // PENALTY ARM admission: the cold constructor's law, verbatim (doc there).
2906        let pen = glm5_penalty_admit(sampling.as_ref())?;
2907        let sampling = sampling.filter(|sp| sp.temp > 0.0);
2908        if fed.is_empty() {
2909            return Err("restored glm5 spec session needs a non-empty restored prefix".into());
2910        }
2911        // FULL-COVER ARM (memra#74, lane/glm5-fullcover-spec-route): an empty suffix is the
2912        // repeated-prompt shape. It refuses unless the arm is armed AND the entry carried
2913        // its boundary logits: without them there is no anchor row and no way to start a
2914        // round, which is exactly the `spec_restore_refusal` full-cover rule for MTP.
2915        if suffix.is_empty() {
2916            if !glm5_spec_fullcover_on() {
2917                return Err(
2918                    "restored glm5 spec session needs a non-empty suffix unless \
2919                     MEMRA_GLM5_SPEC_FULLCOVER=1 (empty-suffix full-cover hits otherwise \
2920                     keep the plain boundary-logits resume)"
2921                        .into(),
2922                );
2923            }
2924            if boundary_logits.len() != self.output.out_features() {
2925                return Err(format!(
2926                    "full-cover glm5 spec restore needs the entry's boundary logits \
2927                     ({} rows, got {})",
2928                    self.output.out_features(),
2929                    boundary_logits.len(),
2930                )
2931                .into());
2932            }
2933        }
2934        if cache.pos != fed.len() {
2935            return Err(format!(
2936                "restored glm5 spec session needs a whole-entry trunk cache: cache.pos {} \
2937                 != restored prefix {}",
2938                cache.pos,
2939                fed.len(),
2940            )
2941            .into());
2942        }
2943        if dkv.len != fed.len() {
2944            return Err(format!(
2945                "restored draft KV len {} != restored prefix {}",
2946                dkv.len,
2947                fed.len(),
2948            )
2949            .into());
2950        }
2951        if dkv.cap != ctx_cap {
2952            return Err(
2953                format!("restored draft KV cap {} != session ctx {ctx_cap}", dkv.cap,).into(),
2954            );
2955        }
2956        // Room for prefix + suffix, the anchor row and at least one verify round.
2957        if fed.len() + suffix.len() + 4 > ctx_cap {
2958            return Err(format!(
2959                "restored glm5 spec session needs ctx for prefix {} + suffix {} + anchor + \
2960                 one verify round, cap {ctx_cap}",
2961                fed.len(),
2962                suffix.len(),
2963            )
2964            .into());
2965        }
2966        let n_vocab = self.output.out_features();
2967        if let Some(map) = self.glm5_d2t() {
2968            if map.iter().any(|&t| t as usize >= n_vocab) {
2969                return Err(format!(
2970                    "glm5 FR-Spec d2t carries a token id >= n_vocab {n_vocab} — the ranks \
2971                     artifact was minted for a different vocabulary"
2972                )
2973                .into());
2974            }
2975            eprintln!("{}", self.glm5_trim_engagement_line(map));
2976        }
2977        if glm5_pmin() > 0.0 {
2978            eprintln!(
2979                "[glm5-spec] draft confidence gate armed: PMIN={:.3} PMIN0={} (native \
2980                 chain p-of-pick; DFlash2 selector-q tau-slot truncation)",
2981                glm5_pmin(),
2982                glm5_pmin0() as u8,
2983            );
2984        }
2985        // ---- suffix prime over the restored planes (the continuation program), taps armed
2986        // for exactly the suffix rows (`HcTapSink::origin` anchors the sink at the restored
2987        // boundary; the chunked walk's absolute bases rebase through it).
2988        let n_embd = self.cfg.n_embd as usize;
2989        let taps = glm5_dflash_tap_layers(&dr.draft, self.layers.len())?;
2990        let eh = self.glm5_head_engine(e)?;
2991        // ---- STREAM ORDERING, UNCONDITIONAL AND BEFORE THE SUFFIX BRANCH (memra#95, the
2992        // fleet-fatal full-cover panic).
2993        //
2994        // Everything this session restores was written on the CALLER's stream, in the
2995        // worker's admission path: the trunk planes by `prefix_restore_at`, and the DFlash2
2996        // drafter ctx KV by `DflashKv::from_tail` (worker.rs, the "glm5 spec restore, half 1
2997        // of 2" block — a fresh allocation plus per-layer `memset_zeros_view` +
2998        // `copy_range_into`). The draft phase of round 1 then READS that KV through
2999        // `eh` = `glm5_head_engine`, which under a live ppN split is the last stage's OWN
3000        // Engine (`pp::PpNRt::build` gives every stage s>0 its own Engine even on the primary
3001        // device, for scratch-pool isolation) and is used OUTSIDE any `rt.enter` scope — so
3002        // it launches on that Engine's own stream, which is neither the caller's stream nor
3003        // the stage stream. Nothing ordered the two, and the drafter's first `forward_round`
3004        // could read the ctx rows before the import landed.
3005        //
3006        // `fence_stages_behind` is NOT the fix and was the first thing this lane got wrong:
3007        // it orders `StageRt::stream`, the enter-scope stream, and says nothing about a stage
3008        // ENGINE's own stream. `order_engine_behind` is the seam for a body that hands a
3009        // stage engine to a helper without entering the stage.
3010        //
3011        // Why only the FULL-COVER arm: the suffix arm calls `prime_cache` on `e`, whose
3012        // logits come back through `Engine::dtoh` (a `stream().synchronize()`), which drains
3013        // the caller's stream before round 1 ever runs. The full-cover arm has no prime and
3014        // no device readback at all between the import and the round — the anchor is drawn
3015        // from the entry's host-side boundary logits — so it is the only restored path with
3016        // nothing between the two. That is why production never saw this on the suffix path,
3017        // and it is why the ordering is done HERE, once per session and for both arms,
3018        // instead of relying on a host sync that belongs to the prime.
3019        //
3020        // The measured failure: a NaN drafter row -> the top-k selector's documented
3021        // exhausted-slot sentinel (`0xffffffff`, `cu/kernels.cu` `topk_rows_f32` and its
3022        // sharded twin) -> `walk: candidate 4294967295 outside codebook vocab` at
3023        // dflash.rs:221/:804, on the FIRST round of a restored session (the receipts put the
3024        // panic between the `RESTORED session` line and the round's first `[glm5-acc]`),
3025        // fleet-fatal after one respawn.
3026        crate::pp::PpNRt::order_engine_behind(e, eh)?;
3027        // FIRST-TOKEN PROFILE (`MEMRA_SPEC_PROF=1`): the restored shape pays no cache or
3028        // drafter allocation here (the caller restored both); its prime bucket is the
3029        // SUFFIX prime (0 on a full-cover hit), which is what makes the restore worth having.
3030        let mut prof = spec_prof_on().then(SpecFirstTokenProf::default);
3031        let mut pclk = prof.as_ref().map(|_| ProfClock::start(e, eh));
3032        // FULL COVER: no suffix, so no prime, no taps and no republish (the entry ALREADY
3033        // sits at this boundary, and a capture here would be the same key the worker's has_key
3034        // dedupe drops). The boundary row is the entry's own; the drafter ctx KV is already
3035        // at `cache.pos` from the tail, so `pending` is empty and round 1's
3036        // `kv.len == cache.pos` invariant holds without an ingest.
3037        let (logits_s, tap_rows, prefix_capture) = if suffix.is_empty() {
3038            (boundary_logits.to_vec(), Vec::new(), None)
3039        } else {
3040            cache.hc_taps = Some(HcTapSink::new_at(
3041                taps.clone(),
3042                n_embd,
3043                suffix.len(),
3044                fed.len(),
3045            ));
3046            let (logits_s, _seed, hiddens) = self.prime_cache(e, suffix, &mut cache, 0)?;
3047            if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
3048                pf.prime_ms = ck.lap(e, eh);
3049            }
3050            // Republish capture at the NEW (deeper) boundary — pos == fed + suffix here.
3051            let capture = if glm5_spec_prefix_on() {
3052                // `hiddens` is the SUFFIX stack: its last row is the boundary hidden
3053                // (indexed through `suffix.len()`, not the absolute boundary).
3054                self.glm5_prefix_boundary_capture(
3055                    e,
3056                    eh,
3057                    &cache,
3058                    &logits_s,
3059                    &hiddens,
3060                    cache.pos,
3061                    suffix.len(),
3062                )
3063            } else {
3064                None
3065            };
3066            let sink = cache
3067                .hc_taps
3068                .take()
3069                .ok_or("glm5 restored-session suffix tap sink vanished")?;
3070            (logits_s, sink.rows, capture)
3071        };
3072        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
3073            pf.capture_ms = ck.lap(e, eh);
3074        }
3075        let mut committed = Vec::with_capacity(fed.len() + suffix.len());
3076        committed.extend_from_slice(fed);
3077        committed.extend_from_slice(suffix);
3078        // The penalty window is the whole committed prompt (restored prefix + suffix), which
3079        // is what the plain hit's sampler replays too ("penalty history replayed over the
3080        // whole prefix"). Empty with penalties off.
3081        let pen_hist = glm5_pen_window_seed(pen.as_ref(), &committed);
3082        let mut sctr = 0u32;
3083        let anchor = glm5_anchor(
3084            eh,
3085            &logits_s,
3086            sampling.as_ref(),
3087            pen.as_ref(),
3088            &pen_hist,
3089            &mut sctr,
3090            "glm5-restore",
3091        )?;
3092        if let (Some(pf), Some(ck)) = (prof.as_mut(), pclk.as_mut()) {
3093            pf.anchor_ms = ck.lap(e, eh);
3094        }
3095        // Engagement receipt (the dspark restore's shape — the deploy gate greps this; a
3096        // cached_tokens number alone cannot distinguish a spec restore from a plain hit).
3097        eprintln!(
3098            "[glm5-spec] RESTORED session: {} prefix tokens + {} suffix from cache — no \
3099             cold prime (drafter tail rows {}, arm {})",
3100            fed.len(),
3101            suffix.len(),
3102            dkv.len,
3103            if suffix.is_empty() {
3104                "full-cover"
3105            } else {
3106                "suffix-prime"
3107            },
3108        );
3109        Ok(Glm5SpecSession {
3110            cache,
3111            committed,
3112            anchor,
3113            anchor_emitted: false,
3114            pending: Vec::new(),
3115            draft: Glm5DraftState::Dflash2 {
3116                kv: dkv,
3117                pending: tap_rows,
3118                taps,
3119            },
3120            sampling,
3121            pen,
3122            pen_hist,
3123            sctr,
3124            uctr: 0,
3125            rounds: 0,
3126            rank_trimmed_rounds: 0,
3127            done: false,
3128            max_ctx: ctx_cap,
3129            mtp_il: None,
3130            prefix_capture,
3131            prof_rounds: prof.as_ref().map(|_| SpecRoundsLog::default()),
3132            prof,
3133        })
3134    }
3135
3136    /// The loaded FR-Spec draft->target map for the session's DRAFT SOURCE (None = full-vocab
3137    /// head, rank id == token id). A loaded DFlash2 drafter wins the source selection, so its
3138    /// trim ([`Self::glm5_dflash_trim`]) is the map; otherwise the embedded head's own map.
3139    fn glm5_d2t(&self) -> Option<&[u32]> {
3140        if self.glm5_dflash.is_some() {
3141            return self.glm5_dflash_trim().map(|(_, d2t)| d2t);
3142        }
3143        self.mtp
3144            .as_ref()
3145            .and_then(|head| head.d2t.as_deref())
3146            .filter(|map| !map.is_empty())
3147    }
3148
3149    /// The DFlash2 round's trimmed draft head `(rows, d2t)`, in the round's preference order:
3150    /// the MtpHead self-trim when its rows came from the TARGET head (an external student's
3151    /// head is never borrowed), else the `dflash_trim` slab (the MEMRA_MTP_SKIP stub, or the
3152    /// glm5 DFlash2 slab built when the NextN block is not loaded, lane/frspec-dflash2-
3153    /// 20260902). None = the full trunk head, stated in the boot receipt.
3154    pub fn glm5_dflash_trim(&self) -> Option<(&crate::model::GpuTensor, &[u32])> {
3155        self.mtp
3156            .as_ref()
3157            .filter(|m| m.d2t_from_target_head)
3158            .and_then(|m| m.shared_head_head.as_ref().zip(m.d2t.as_deref()))
3159            .or_else(|| {
3160                self.dflash_trim
3161                    .as_ref()
3162                    .map(|t| (&t.head, t.d2t.as_slice()))
3163            })
3164            .filter(|(_, d2t)| !d2t.is_empty())
3165    }
3166
3167    /// The trim engagement receipt at session admission (the dspark trim receipt's shape;
3168    /// the deploy gate greps the server log for it). Source-keyed wording: the DFlash2 route
3169    /// names the slab and the ranks file (`RANK-TRIMMED n_ranks=<n> src=<sha16>`), the native
3170    /// chain keeps its line.
3171    fn glm5_trim_engagement_line(&self, map: &[u32]) -> String {
3172        if self.glm5_dflash.is_some() {
3173            format!(
3174                "[glm5-spec] draft head RANK-TRIMMED n_ranks={} src={}",
3175                map.len(),
3176                self.frspec_src_sha16.as_deref().unwrap_or("unknown")
3177            )
3178        } else {
3179            format!(
3180                "[glm5-spec] draft head TRIMMED to {} rows (FR-Spec d2t engaged)",
3181                map.len()
3182            )
3183        }
3184    }
3185
3186    /// ONE serve burst (the worker's per-tick call, `step_glm5_spec`): rounds of
3187    /// draft(K) -> `glm5_verify_rows` -> accept -> rollback/commit until `target` new
3188    /// tokens are out, EOS commits, or the context guard trips. Returns
3189    /// `(burst, drafted, accepted)`; the burst may overshoot `target` by up to K (a
3190    /// round commits j+1 tokens atomically — the engine surplus stays committed in the
3191    /// session cache and the WORKER clamps public emission to the request budget, the
3192    /// SpecSession overshoot contract).
3193    pub fn glm5_spec_session_burst(
3194        &self,
3195        e: &Engine,
3196        sess: &mut Glm5SpecSession,
3197        target: usize,
3198        k: usize,
3199        eos: &[u32],
3200    ) -> Res<(Vec<u32>, usize, usize)> {
3201        self.glm5_spec_session_burst_inner(
3202            e,
3203            sess,
3204            target,
3205            k,
3206            eos,
3207            &mut Glm5SpecKnobs::default(),
3208            None,
3209        )
3210    }
3211
3212    /// [`glm5_spec_session_burst`] with a ROUND-CADENCE commit hook (lane/b200-spec-ttft-
3213    /// 20260902, the `MEMRA_SPEC_FIRST_TOKEN_EAGER` door's engine half; the spec.rs
3214    /// `on_commit` sse-cadence pattern): `on_commit` is called with every newly committed
3215    /// slice of the burst — first the prime's anchor alone, then each round's `j` accepted
3216    /// drafts + bonus — as DISJOINT, IN-ORDER slices whose concatenation IS the returned
3217    /// burst, byte for byte. The hook returns nothing and the loop's control flow never
3218    /// reads it, so the tokens produced are exactly `glm5_spec_session_burst`'s; only WHEN
3219    /// the caller learns them moves (from burst end to round end). Without a hook the
3220    /// first token of a cold session waits for the whole `target`-token burst.
3221    pub fn glm5_spec_session_burst_streamed(
3222        &self,
3223        e: &Engine,
3224        sess: &mut Glm5SpecSession,
3225        target: usize,
3226        k: usize,
3227        eos: &[u32],
3228        on_commit: CommitHook<'_>,
3229    ) -> Res<(Vec<u32>, usize, usize)> {
3230        self.glm5_spec_session_burst_inner(
3231            e,
3232            sess,
3233            target,
3234            k,
3235            eos,
3236            &mut Glm5SpecKnobs::default(),
3237            Some(on_commit),
3238        )
3239    }
3240
3241    /// [`glm5_spec_session_burst`] with GATE INSTRUMENTS (`Glm5SpecKnobs` — never a serving
3242    /// surface; no serving path constructs a non-default value).
3243    pub fn glm5_spec_session_burst_gated(
3244        &self,
3245        e: &Engine,
3246        sess: &mut Glm5SpecSession,
3247        target: usize,
3248        k: usize,
3249        eos: &[u32],
3250        knobs: &mut Glm5SpecKnobs<'_>,
3251    ) -> Res<(Vec<u32>, usize, usize)> {
3252        self.glm5_spec_session_burst_inner(e, sess, target, k, eos, knobs, None)
3253    }
3254
3255    /// The one burst loop behind the three public entries (plain / streamed / gated).
3256    #[allow(clippy::too_many_arguments)]
3257    // allow: the parameter list is the burst contract plus the two instruments (gate
3258    // knobs, commit hook); bundling would hide which inputs are serving vs instrument
3259    fn glm5_spec_session_burst_inner(
3260        &self,
3261        e: &Engine,
3262        sess: &mut Glm5SpecSession,
3263        target: usize,
3264        k: usize,
3265        eos: &[u32],
3266        knobs: &mut Glm5SpecKnobs<'_>,
3267        mut on_commit: Option<CommitHook<'_>>,
3268    ) -> Res<(Vec<u32>, usize, usize)> {
3269        let cap = Self::hyper_batch_cap();
3270        if k == 0 || k + 1 > cap {
3271            return Err(format!(
3272                "glm5_spec_session_burst: k={k} outside 1..={} (verify rows = k+1 must stay \
3273                 inside the decode-exact knee, cap {cap})",
3274                cap - 1
3275            )
3276            .into());
3277        }
3278        if let Glm5DraftState::Dflash2 { .. } = sess.draft {
3279            let b = self
3280                .glm5_dflash
3281                .as_ref()
3282                .ok_or("dflash session on a model with no loaded drafter")?
3283                .draft
3284                .cfg
3285                .block_size;
3286            if k + 1 > b {
3287                return Err(format!(
3288                    "glm5_spec_session_burst: k={k} exceeds the DFlash2 drafter's block \
3289                     (block_size {b} = anchor + {} drafts, the trained mask pattern) — \
3290                     the worker clamps operator K pins to {} for this source; refusing \
3291                     loudly rather than drafting an untrained shape",
3292                    b - 1,
3293                    b - 1
3294                )
3295                .into());
3296            }
3297        }
3298        let d2t = self.glm5_d2t();
3299        if d2t.is_some() && knobs.skip_d2t_remap {
3300            eprintln!("[glm5-spec] d2t REMAP SKIPPED — red-arm instrument, drafts are rank ids");
3301        }
3302        let sp_on: Option<SpecSampling> = sess.sampling.filter(|sp| sp.temp > 0.0);
3303        let mut out: Vec<u32> = Vec::with_capacity(target + k);
3304        let mut drafted = 0usize;
3305        let mut accepted = 0usize;
3306        let mut phase: Option<SpecPhaseNs> =
3307            crate::spec_phase::spec_trace_on().then(SpecPhaseNs::default);
3308        // FIRST-TOKEN PROFILE: the first burst's wall (its tokens are host-visible at
3309        // return, so no drain is needed to bound it).
3310        let first_burst = (!sess.anchor_emitted && sess.prof.is_some())
3311            .then(|| (std::time::Instant::now(), sess.rounds));
3312        // The hook's own share of the first burst's wall (host-only detext + sends), so
3313        // the profile can separate engine time from emission time under the eager door.
3314        let mut hook_ns: u64 = 0;
3315        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
3316        let mut flushed = 0usize;
3317        if !sess.anchor_emitted {
3318            // The prime's boundary token: emitted exactly once, by the first burst.
3319            out.push(sess.anchor);
3320            sess.anchor_emitted = true;
3321            if eos.contains(&sess.anchor) {
3322                sess.done = true;
3323            }
3324            if let Some(cb) = on_commit.as_mut() {
3325                let t_cb = first_burst.map(|_| std::time::Instant::now());
3326                cb(&out[flushed..]);
3327                if let Some(t_cb) = t_cb {
3328                    hook_ns += t_cb.elapsed().as_nanos() as u64;
3329                }
3330                flushed = out.len();
3331            }
3332        }
3333        while out.len() < target && !sess.done {
3334            // Context guard: a round appends up to k+1 trunk rows from `cache.pos` (and the
3335            // draft plane stays <= pos + k), so the next round must fit with one row slack.
3336            if sess.cache.pos + k + 2 > sess.max_ctx {
3337                sess.done = true;
3338                break;
3339            }
3340            // DEPTH LOG (lane/spec-route-depth-20260902, `MEMRA_SPEC_PROF=1`): the first
3341            // SPEC_PROF_ROUNDS rounds of a session get their own phase accumulator (the
3342            // trace's drains), a wall clock, and the per-row-arm row count — folded into
3343            // the per-burst trace accumulator afterwards so both instruments agree.
3344            let log_round = sess.prof_rounds.as_ref().is_some_and(|l| l.wants_more());
3345            let mut rp = log_round.then(SpecPhaseNs::default);
3346            let t_round = log_round.then(std::time::Instant::now);
3347            let seq0 = V_SEQ_ROWS.load(std::sync::atomic::Ordering::Relaxed);
3348            let ctx0 = sess.cache.pos;
3349            let (round_tokens, n_drafted) = {
3350                let ph: Option<&mut SpecPhaseNs> = match (rp.as_mut(), phase.as_mut()) {
3351                    (Some(r), _) => Some(r),
3352                    (None, p) => p,
3353                };
3354                self.glm5_spec_round(e, sess, k, d2t, sp_on.as_ref(), knobs, ph)?
3355            };
3356            if let (Some(r), Some(t0)) = (rp.as_ref(), t_round) {
3357                if let Some(p) = phase.as_mut() {
3358                    p.add(r);
3359                }
3360                if let Some(log) = sess.prof_rounds.as_mut() {
3361                    let ms = |ns: u64| ns as f32 / 1e6;
3362                    log.push(SpecRoundProf {
3363                        wall_ms: t0.elapsed().as_secs_f32() * 1e3,
3364                        draft_ms: ms(r.draft),
3365                        verify_ms: ms(r.verify),
3366                        accept_ms: ms(r.accept),
3367                        rest_ms: ms(r.roll + r.maint),
3368                        k: n_drafted as u16,
3369                        j: (round_tokens.len() - 1) as u16,
3370                        ctx: ctx0 as u32,
3371                        seq_rows: (V_SEQ_ROWS.load(std::sync::atomic::Ordering::Relaxed) - seq0)
3372                            as u32,
3373                    });
3374                }
3375            }
3376            drafted += n_drafted;
3377            accepted += round_tokens.len() - 1; // j accepted drafts + the bonus row
3378            for &tok in &round_tokens {
3379                if eos.contains(&tok) {
3380                    sess.done = true;
3381                }
3382            }
3383            out.extend_from_slice(&round_tokens);
3384            sess.rounds += 1;
3385            if let Some(cb) = on_commit.as_mut() {
3386                // sse-cadence: this round's accepted drafts + bonus are committed — hand
3387                // the caller exactly the not-yet-flushed tail (disjoint, in order).
3388                let t_cb = first_burst.map(|_| std::time::Instant::now());
3389                cb(&out[flushed..]);
3390                if let Some(t_cb) = t_cb {
3391                    hook_ns += t_cb.elapsed().as_nanos() as u64;
3392                }
3393                flushed = out.len();
3394            }
3395        }
3396        debug_assert!(
3397            on_commit.is_none() || flushed == out.len(),
3398            "every committed token must have been handed to on_commit"
3399        );
3400        if let Some(ph) = phase.as_ref() {
3401            ph.emit("glm5-phase", "glm5-phase-v", k);
3402        }
3403        if let (Some((t0, rounds0)), Some(pf)) = (first_burst, sess.prof.as_mut()) {
3404            pf.first_burst_ms = t0.elapsed().as_secs_f64() * 1e3;
3405            pf.first_burst_hook_ms = hook_ns as f64 / 1e6;
3406            pf.first_burst_rounds = sess.rounds - rounds0;
3407            pf.first_burst_tokens = out.len();
3408        }
3409        Ok((out, drafted, accepted))
3410    }
3411
3412    /// One draft->verify->accept->rollback->re-seed round over the session state. Returns
3413    /// `(round_tokens, n_drafted)`: the round's committed tokens (`j` accepted drafts + the
3414    /// bonus token) and how many drafts actually entered the verify (== `k` today; the
3415    /// confidence gate may truncate it below `k`).
3416    #[allow(clippy::too_many_arguments)]
3417    // allow: the parameter list mirrors the round contract (session + policy + gate knobs +
3418    // the trace accumulator); bundling would hide which inputs are serving vs instrument
3419    fn glm5_spec_round(
3420        &self,
3421        e: &Engine,
3422        sess: &mut Glm5SpecSession,
3423        k: usize,
3424        d2t: Option<&[u32]>,
3425        sp: Option<&SpecSampling>,
3426        knobs: &mut Glm5SpecKnobs<'_>,
3427        mut phase: Option<&mut SpecPhaseNs>,
3428    ) -> Res<(Vec<u32>, usize)> {
3429        let n_vocab = self.output.out_features();
3430        let n_embd = self.cfg.n_embd as usize;
3431        // The MTP block / DFlash2 drafter, the trunk lm head and the verify walk's returned
3432        // rows all live on the LAST stage under a split — every draft-chain and accept-side
3433        // op below runs through the head engine (identity when the door is shut).
3434        let eh = self.glm5_head_engine(e)?;
3435        let mut t_mark = phase.as_ref().map(|_| SpecPhaseNs::clock(e, eh));
3436        // FIRST-TOKEN PROFILE (`MEMRA_SPEC_PROF=1`): round 1 only — the same drains as the
3437        // trace, bucketed into the session's one-shot profile instead of the per-burst
3438        // accumulator. `pclk` rides into the DFlash2 draft fn so the prompt ingest (the
3439        // drafter prime) gets its own bucket before the draft bucket starts.
3440        let mut pclk = (sess.rounds == 0 && sess.prof.is_some()).then(|| ProfClock::start(e, eh));
3441        // Phase-boundary bump: drain, bucket the elapsed ns, restart the clock. No-op with
3442        // the trace off (t_mark is None and no stream is ever synchronized).
3443        macro_rules! bump {
3444            ($field:ident) => {
3445                if let (Some(ph), Some(t0)) = (phase.as_deref_mut(), t_mark.as_mut()) {
3446                    let now = SpecPhaseNs::clock(e, eh);
3447                    ph.$field += now.duration_since(*t0).as_nanos() as u64;
3448                    *t0 = now;
3449                }
3450            };
3451        }
3452        // Profile lap into the named first-round bucket (no-op with the profile off).
3453        macro_rules! plap {
3454            ($field:ident) => {
3455                if let (Some(ck), Some(pf)) = (pclk.as_mut(), sess.prof.as_mut()) {
3456                    pf.$field = ck.lap(e, eh);
3457                }
3458            };
3459        }
3460
3461        // CONFIDENCE GATE resolution (loop-port 2): the env pair is the serving surface
3462        // (the step37 family, no new flags); the knobs override is the gate instrument.
3463        let (p_min, pmin0) = knobs
3464            .pmin_override
3465            .unwrap_or_else(|| (glm5_pmin(), glm5_pmin0()));
3466
3467        // ---- 1+2. produce the K drafts (+ the retained q side), SOURCE-KEYED. Everything
3468        // after this point is shared and source-blind — the exactness seam (module doc).
3469        let (drafts, qside, mtp_committed_len) = match sess.draft {
3470            Glm5DraftState::Dflash2 { .. } => {
3471                let (d, q) = self.glm5_dflash_round_drafts(
3472                    eh,
3473                    sess,
3474                    k,
3475                    sp,
3476                    knobs,
3477                    p_min,
3478                    pmin0,
3479                    pclk.as_mut(),
3480                )?;
3481                (d, q, 0)
3482            }
3483            Glm5DraftState::NativeMtp => {
3484                let mtp_il = sess.mtp_il.ok_or("native-mtp arm without a plane index")?;
3485                // ---- feed pending committed pairs; the last call yields draft 1 ----
3486                let mut last: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
3487                for (tok, h) in sess.pending.drain(..) {
3488                    let plane_len = sess.cache.latent[mtp_il]
3489                        .as_ref()
3490                        .ok_or("MTP plane missing")?
3491                        .len;
3492                    last = Some(self.mtp_head_forward_mla_cached(
3493                        eh,
3494                        0,
3495                        tok,
3496                        &h,
3497                        &mut sess.cache,
3498                        plane_len,
3499                    )?);
3500                }
3501                let (mut d_logits, mut carrier) =
3502                    last.ok_or("glm5 spec round started with no pending committed pair")?;
3503                let mtp_committed_len = sess.cache.latent[mtp_il]
3504                    .as_ref()
3505                    .ok_or("MTP plane missing")?
3506                    .len;
3507
3508                // ---- chain K drafts. Greedy route: argmax over the draft head. Sampled
3509                // route: filtered Gumbel draw through the session's device Philox stream
3510                // (`sctr`), with the per-step filtered stats + logits retained — they are
3511                // the q side of the accept walk. Trimmed heads yield RANK ids that remap
3512                // through d2t to true vocab before anything consumes them (chain feed,
3513                // verify, output); the q gather keeps the rank id.
3514                let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
3515                let mut drafts: Vec<u32> = Vec::with_capacity(k); // true-vocab tokens
3516                let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // draft-head rank ids
3517                let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // sampled route only
3518                let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (mx, th, z), sampled only
3519                for ki in 0..k {
3520                    let (idx, sampled_stats) = match sp {
3521                        Some(sp) => {
3522                            let (idx, stats) =
3523                                glm5_sampled_draft(eh, &d_logits, d_vocab, sp, &mut sess.sctr)?;
3524                            (idx, Some(stats))
3525                        }
3526                        None => {
3527                            // Device argmax + ONE 4-byte readback per draft (loop-port 1)
3528                            // — replaces the full d_vocab logits DtoH + host argmax the
3529                            // map names at this seam. Same tie-break contract
3530                            // (argmax_gate); drafts never decide exactness anyway, the
3531                            // verify arbitrates. The #87 sentinel guard mirrors the
3532                            // spec.rs graph chain: a device argmax may emit a sentinel
3533                            // on a NaN row — refuse loudly, never gather an OOB embed.
3534                            let td = eh.argmax_token_device(&d_logits, d_vocab)?;
3535                            let idx = crate::spec::guard_vocab_token(
3536                                eh.dtoh_u32_one(&td)?,
3537                                d_vocab,
3538                                &format!(
3539                                    "glm5 native draft argmax at round {} ki={ki}",
3540                                    sess.rounds
3541                                ),
3542                            )?;
3543                            (idx, None)
3544                        }
3545                    };
3546                    // P-MIN CONFIDENCE GATE (loop-port 2, the spec.rs chain break): p =
3547                    // the head's softmax confidence in its own pick (the `g_p` statistic,
3548                    // prob_of_token_device kernels), one 4-byte read — armed rounds only.
3549                    // Break BEFORE the pick is drafted or the next full-MoE-layer chain
3550                    // forward is paid; a discarded sampled draw's Philox advance stands
3551                    // (spec.rs eager parity: "counts the p-min-discarded token too").
3552                    if p_min > 0.0 {
3553                        let tok_d = eh.htod_u32_v(&[idx])?;
3554                        let p_d = eh.prob_of_token_device(&d_logits, &tok_d, d_vocab)?;
3555                        let p = eh.dtoh(&p_d)?[0];
3556                        if p < p_min && (ki > 0 || pmin0) {
3557                            break;
3558                        }
3559                    }
3560                    if let Some(stats) = sampled_stats {
3561                        draft_stats.push(stats);
3562                        draft_logits.push(eh.clone_dtod(&d_logits)?);
3563                    }
3564                    let mut d = match d2t {
3565                        Some(map) if !knobs.skip_d2t_remap => map[idx as usize],
3566                        _ => idx,
3567                    };
3568                    if let Some(over) = knobs.draft_override.as_mut() {
3569                        d = over(sess.rounds, ki, d);
3570                    }
3571                    drafts.push(d);
3572                    draft_idx.push(idx);
3573                    if ki + 1 < k {
3574                        let plane_len = sess.cache.latent[mtp_il]
3575                            .as_ref()
3576                            .ok_or("MTP plane missing")?
3577                            .len;
3578                        let (lg, ca) = self.mtp_head_forward_mla_cached(
3579                            eh,
3580                            0,
3581                            d,
3582                            &carrier,
3583                            &mut sess.cache,
3584                            plane_len,
3585                        )?;
3586                        d_logits = lg;
3587                        carrier = ca;
3588                    }
3589                }
3590                let q = match sp {
3591                    Some(_) => Glm5DraftQ::Mtp {
3592                        draft_idx,
3593                        draft_logits,
3594                        draft_stats,
3595                    },
3596                    None => Glm5DraftQ::None,
3597                };
3598                (drafts, q, mtp_committed_len)
3599            }
3600        };
3601        bump!(draft);
3602        plap!(first_draft_ms);
3603        // RANK-TRIM COUNTER (lane/frspec-dflash2-20260902): this round drafted through a
3604        // rank-trimmed head (`d2t` is the source-keyed map the drafts were produced under ,
3605        // the DFlash2 slab or the native trimmed chain). Counted on the head, not the remap:
3606        // the skip_d2t_remap red arm still ran over the slab.
3607        if d2t.is_some() {
3608            sess.rank_trimmed_rounds += 1;
3609            GLM5_RANK_TRIMMED_DRAFT_ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3610        }
3611
3612        // DFlash2 source: arm the verify tap — the walk's rows are next round's drafter
3613        // context features (rows 0..keep survive the accept; the sink is taken in step 7).
3614        // DEVICE-STAGED (loop-port 1): the walk D2Ds each tapped layer's contracted rows
3615        // instead of blocking on five in-walk DtoHs; step 7 drains post-walk.
3616        if let Glm5DraftState::Dflash2 { taps, .. } = &sess.draft {
3617            sess.cache.hc_taps = Some(HcTapSink::new_device_staged(
3618                taps.clone(),
3619                n_embd,
3620                drafts.len() + 1,
3621            ));
3622        }
3623
3624        // ---- 3. verify: one t=K+1 walk over the trunk ----
3625        let mut rows: Vec<u32> = Vec::with_capacity(drafts.len() + 1);
3626        rows.push(sess.anchor);
3627        rows.extend_from_slice(&drafts);
3628        let (vlogits, collapsed, ckpt) = self.glm5_verify_rows(e, &rows, &mut sess.cache)?;
3629        bump!(verify);
3630        plap!(first_verify_ms);
3631
3632        // ---- 3b. PENALTY ARM (lane/spec-exclusions-20260902, MEMRA_SPEC_PENALTY=1) ----
3633        // The target rows are penalized IN A COPY over the per-row EVOLVING window: row r
3634        // (the target at drafts[r]'s slot; row 0 = the anchor's) sees `pen_win ++
3635        // drafts[..r]` — the tokens committed before that position on every path where the
3636        // row is consulted, same-round accepts included — through `penalize_logits_rows_inc`
3637        // (the dspark accept walk's kernel, host-order arithmetic). Every p read below —
3638        // the greedy argmaxes, the sampled stats/gather, the bonus row, the reject-slot
3639        // residual — points at that copy, so the accepted stream is the plain penalized
3640        // sampler's target: greedy byte-identical (gate 15), sampled distribution-exact.
3641        // The anchor joins the window first (it was emitted last round, committed by this
3642        // one: the dspark `pen_hist.push(last)` convention); accepted drafts join after the
3643        // accept. `pen: None` = no copy, no launch — `plogits` IS `vlogits`, the pre-lane
3644        // program byte for byte.
3645        let pvl: Option<CudaSlice<f32>> = match sess.pen {
3646            Some(p) => {
3647                sess.pen_hist.push(sess.anchor);
3648                let w0 = sess.pen_hist.len().saturating_sub(p.win());
3649                let pen_win = &sess.pen_hist[w0..];
3650                let mut hist: Vec<u32> = Vec::with_capacity(pen_win.len() + drafts.len());
3651                hist.extend_from_slice(pen_win);
3652                hist.extend_from_slice(&drafts);
3653                let n_win = pen_win.len();
3654                let hd = eh.htod_u32_v(&hist)?;
3655                let mut buf = eh.clone_dtod(&vlogits)?;
3656                eh.penalize_logits_rows_inc(
3657                    &mut buf,
3658                    &hd,
3659                    n_win,
3660                    p.rep,
3661                    p.freq,
3662                    p.present,
3663                    n_vocab,
3664                    rows.len(),
3665                    p.win(),
3666                )?;
3667                Some(buf)
3668            }
3669            None => None,
3670        };
3671        let plogits: &CudaSlice<f32> = pvl.as_ref().unwrap_or(&vlogits);
3672        // The dspark accept walk penalizes internally off `sp.pen_on()`; under the arm it
3673        // receives the ALREADY penalized rows and a penalty-neutral config, so the round has
3674        // exactly one penalty pass. With the arm off this is `sp` itself, untouched.
3675        let sp_accept: Option<SpecSampling> = sp.map(|s| match sess.pen {
3676            Some(_) => SpecSampling {
3677                penalty_last_n: 0,
3678                ..*s
3679            },
3680            None => *s,
3681        });
3682
3683        // ---- 4. accept ----
3684        // ZERO-DRAFT SAMPLED ROUND (PMIN0): the verify batch is just the anchor row —
3685        // m=1 = a plain decode step, exactly the llama.cpp gating spec.rs vendored. The
3686        // bonus is the full-accept filtered-Gumbel draw from that one row through the
3687        // session's Philox stream (identical in distribution to the plain sampled step
3688        // this round degenerates to). Greedy zero-draft rounds ride the general arm
3689        // below (j=0, bonus = the row-0 device argmax).
3690        let (j, bonus) = if let (true, Some(sp)) = (drafts.is_empty(), sp) {
3691            (
3692                0,
3693                self.glm5_sampled_bonus(eh, sess, sp, plogits, 0, n_vocab)?,
3694            )
3695        } else {
3696            match (sp, &qside) {
3697                (None, _) => {
3698                    // Greedy longest matching prefix (the DFlash2 probe's rule); bonus = the
3699                    // target's own argmax at the first non-accepted slot. Byte-deterministic —
3700                    // the instrument the spec-vs-plain identity gates pin.
3701                    //
3702                    // DEVICE ACCEPT ARGMAXES (loop-port 1, the K=1 flip): per verify row, a
3703                    // device argmax into one [t] slot buffer, then ONE tiny u32 readback —
3704                    // replacing the (K+1) x n_vocab logits DtoH + (K+1) host argmax scans
3705                    // (~2.4 MB + a host walk over 600k floats at K=3 on the real head; the
3706                    // 3way arithmetic needs 0.67 ms off the fixed round cost to flip K=1).
3707                    // `argmax_token_device_col` carries the host argmax's tie-break contract
3708                    // bit for bit (lowest index wins, argmax_gate-validated), so the accept
3709                    // walk commits the SAME tokens in the SAME order — the byte-identity
3710                    // batteries below stay the proof.
3711                    let t = rows.len();
3712                    let mut vam_d = eh.alloc_u32_zeroed(t)?;
3713                    for r in 0..t {
3714                        eh.argmax_token_device_col(plogits, r, n_vocab, &mut vam_d, r)?;
3715                    }
3716                    let vam = eh.dtoh_u32(&vam_d)?;
3717                    let mut j = 0usize;
3718                    while j < drafts.len() && drafts[j] == vam[j] {
3719                        j += 1;
3720                    }
3721                    if knobs.accept_probe {
3722                        self.glm5_accept_probe(eh, sess.rounds, plogits, &drafts, &vam, j)?;
3723                    }
3724                    (j, vam[j])
3725                }
3726                (
3727                    Some(sp),
3728                    Glm5DraftQ::Mtp {
3729                        draft_idx,
3730                        draft_logits,
3731                        draft_stats,
3732                    },
3733                ) => self.glm5_sampled_accept(
3734                    eh,
3735                    sess,
3736                    sp,
3737                    plogits,
3738                    &drafts,
3739                    draft_idx,
3740                    draft_logits,
3741                    draft_stats,
3742                    d2t,
3743                    drafts.len(),
3744                )?,
3745                (Some(_), Glm5DraftQ::Selector { prop, dl }) => {
3746                    // The q38 serve route's rejection walk, VERBATIM (`dspark_accept_sampled`):
3747                    // `rows` = [anchor, drafts..] is its cand contract, verify row j arbitrates
3748                    // rows[j+1], the bonus draws from row k on full accept, and the reject-slot
3749                    // residual uses the selector's sparse candidate-set q. Philox counters are
3750                    // this session's — randomness never repeats across bursts.
3751                    let (m, next) = crate::dflash::dspark_accept_sampled(
3752                        eh,
3753                        plogits,
3754                        &rows,
3755                        rows.len(),
3756                        n_vocab,
3757                        dl,
3758                        prop,
3759                        sp_accept.as_ref().expect("sampled arm carries its config"),
3760                        &[],
3761                        &mut sess.sctr,
3762                        &mut sess.uctr,
3763                    )?;
3764                    let next = crate::spec::guard_vocab_token(
3765                        next,
3766                        n_vocab,
3767                        &format!(
3768                            "glm5 dflash2 sampled verify bonus at round {} j={m}",
3769                            sess.rounds
3770                        ),
3771                    )?;
3772                    (m, next)
3773                }
3774                (Some(_), Glm5DraftQ::None) => {
3775                    unreachable!("sampled round without a retained q side")
3776                }
3777            }
3778        };
3779        // Accepted drafts join the penalty window (the bonus joins as next round's anchor).
3780        if sess.pen.is_some() {
3781            sess.pen_hist.extend_from_slice(&drafts[..j]);
3782        }
3783        bump!(accept);
3784        plap!(first_accept_ms);
3785
3786        // ---- 5. commit j drafts + the bonus token ----
3787        let mut round_tokens: Vec<u32> = Vec::with_capacity(j + 1);
3788        round_tokens.extend_from_slice(&drafts[..j]);
3789        round_tokens.push(bonus);
3790
3791        // ---- 6. rollback the trunk to the accepted prefix ----
3792        let keep = j + 1;
3793        if knobs.disable_rollback {
3794            // RED-ARM INSTRUMENT: move pos, leave every state plane at post-row-K.
3795            sess.cache.pos = ckpt.pos + keep;
3796        } else {
3797            self.glm5_verify_rollback(e, &mut sess.cache, &ckpt, keep)?;
3798        }
3799        bump!(roll);
3800        plap!(first_roll_ms);
3801
3802        // ---- 7+8. draft-source state maintenance, SOURCE-KEYED ----
3803        match &mut sess.draft {
3804            Glm5DraftState::NativeMtp => {
3805                // MTP plane: len reset to the committed boundary (chain rows out), then
3806                // re-seed the pending pairs (token at pos0+i, collapsed row i-1).
3807                self.glm5_mtp_plane_reset(e, &mut sess.cache, mtp_committed_len)?;
3808                for i in 1..=keep {
3809                    let tok = round_tokens[i - 1];
3810                    let h = self.glm5_seed_row(eh, &collapsed, rows.len(), i - 1)?;
3811                    sess.pending.push((tok, h));
3812                }
3813            }
3814            Glm5DraftState::Dflash2 { pending, taps, .. } => {
3815                // The kept verify rows' tap features (rows 0..keep = [anchor, accepted
3816                // drafts]) become next round's drafter context — the probe's
3817                // `F_feat[new_lo:start]` advance. The drafter's own KV block rows were
3818                // transient (forward_round never moves kv.len), so no drafter rollback
3819                // exists to run. The trunk-side MTP plane was never touched.
3820                // Device-staged rows drain HERE — the round's one post-walk sync point
3821                // for tap features (loop-port 1).
3822                let mut sink = sess
3823                    .cache
3824                    .hc_taps
3825                    .take()
3826                    .ok_or("glm5 dflash verify tap sink vanished")?;
3827                self.glm5_tap_drain(e, &mut sink)?;
3828                let row_w = taps.len() * n_embd;
3829                pending.extend_from_slice(&sink.rows[..keep * row_w]);
3830            }
3831        }
3832        // Cache-row bookkeeping: the trunk committed rows [anchor, drafts[..j]] (keep =
3833        // j+1), so `committed` gains exactly those tokens; the BONUS is the new live
3834        // anchor — emitted this round, consumed by the trunk as the NEXT round's row 0
3835        // (the dspark `last` convention). Invariant at every round boundary:
3836        // `cache.pos == committed.len()`, token-for-token.
3837        sess.committed.push(sess.anchor);
3838        sess.committed.extend_from_slice(&drafts[..j]);
3839        sess.anchor = bonus;
3840        bump!(maint);
3841        plap!(first_maint_ms);
3842        if let Some(pf) = sess.prof.as_mut()
3843            && pclk.is_some()
3844        {
3845            pf.first_round_tokens = round_tokens.len();
3846        }
3847        if let Some(ph) = phase {
3848            ph.rounds += 1;
3849        }
3850        Ok((round_tokens, drafts.len()))
3851    }
3852
3853    /// THE ACCEPTANCE-RACE FIX (lane/glm5-accrace 2026-09-01): order the CALLER's stream
3854    /// behind EVERY stage stream — the exit mirror of
3855    /// [`crate::pp::PpNRt::fence_stages_behind`], built from
3856    /// [`crate::pp::PpNRt::publish_all_to`] (event waits, never a device sync, so the stage
3857    /// streams keep running).
3858    ///
3859    /// Call OUTSIDE any `rt.enter` scope: `e.stream()` must resolve to the caller's stream,
3860    /// not a stage's. Door shut or the same-stream seam (`MEMRA_PP_STREAMS=0`): a no-op by
3861    /// construction, so single-device and STREAMS=0 behaviour is untouched.
3862    fn glm5_publish_stages(&self, e: &Engine) -> Res<()> {
3863        if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
3864            let rt = crate::pp::PpNRt::get(e)?;
3865            let dst = e.stream();
3866            rt.publish_all_to(&dst)?;
3867        }
3868        Ok(())
3869    }
3870
3871    /// GATE INSTRUMENT (lane/glm5-accrace; contract in [`Glm5SpecKnobs::accept_probe`]):
3872    /// one stderr line per greedy round pairing the DEVICE accept row against a HOST
3873    /// argmax over the same buffer, plus a per-row (argmax, row hash) census so two runs
3874    /// of the same deterministic fixture can be diffed round-for-round.
3875    fn glm5_accept_probe(
3876        &self,
3877        eh: &Engine,
3878        round: usize,
3879        vlogits: &CudaSlice<f32>,
3880        drafts: &[u32],
3881        vam: &[u32],
3882        j: usize,
3883    ) -> Res<()> {
3884        let n_vocab = self.output.out_features();
3885        let t = vam.len();
3886        let host = eh.dtoh(vlogits)?;
3887        let mut hvam: Vec<u32> = Vec::with_capacity(t);
3888        let mut rows_census: Vec<String> = Vec::with_capacity(t);
3889        for r in 0..t {
3890            let row = &host[r * n_vocab..(r + 1) * n_vocab];
3891            let am = argmax(row) as u32;
3892            hvam.push(am);
3893            // FNV-1a over the row's f32 BITS: a bit-level fingerprint, so a run-to-run
3894            // diff is exact rather than eyeballed at some print precision.
3895            let mut h: u64 = 0xcbf2_9ce4_8422_2325;
3896            for v in row {
3897                for b in v.to_bits().to_le_bytes() {
3898                    h ^= u64::from(b);
3899                    h = h.wrapping_mul(0x100_0000_01b3);
3900                }
3901            }
3902            rows_census.push(format!("{r}:{am}:{h:016x}"));
3903        }
3904        eprintln!(
3905            "[glm5-accrace] round={round} t={t} j={j} keep={} drafts={drafts:?} \
3906             dev_vam={vam:?} host_vam={hvam:?} agree={} rows=[{}]",
3907            j + 1,
3908            hvam == vam,
3909            rows_census.join(" ")
3910        );
3911        Ok(())
3912    }
3913
3914    /// ONE round's drafts from the DFlash2 source (module doc, DRAFT SOURCE SEAM) — the
3915    /// shipped q38 selector round, re-aimed at glm5's hc-contract features:
3916    ///
3917    ///   1. ingest pending committed feature rows into the drafter's own ctx KV (chunked
3918    ///      at 256 rows — the qwen depth-OOM bound; round 1 carries the whole prompt);
3919    ///   2. block forward `[anchor, MASK x b-1]` at absolute positions over the cached ctx
3920    ///      (`forward_round` — block K/V transient, exactly the reference crop);
3921    ///   3. draft logits = trunk lm_head over rows 1..b (mask-fill harvest — the DFlash2
3922    ///      family census; FR-Spec trim consumed exactly as the dspark serve arm does);
3923    ///   4. selector walk: greedy chain, or the sampled candidate-set walk whose recorded
3924    ///      q (`DsparkDraftSample::Selector`) the shared rejection accept consumes.
3925    ///
3926    /// Drafts are truncated to `k` (the chain is sequential, so a prefix is well-formed),
3927    /// then to the CONFIDENCE prefix when `p_min` is armed (loop-port 2, the tau-slot
3928    /// form): the selector's recorded per-slot q — `q_chosen` on the sampled walk, its
3929    /// T=1 twin on the greedy walk — gates each slot through `glm5_conf_keep`, so the
3930    /// low-confidence tail never enters the verify batch (a truncated round rides down
3931    /// the `31.6 + 20.1*K` line; rejection sampling stays exact for any proposal prefix).
3932    /// `knobs.draft_override` applies after — the gate instrument, never serving.
3933    #[allow(clippy::too_many_arguments)]
3934    // allow: the parameter list mirrors the round contract plus the resolved gate pair
3935    fn glm5_dflash_round_drafts(
3936        &self,
3937        eh: &Engine,
3938        sess: &mut Glm5SpecSession,
3939        k: usize,
3940        sp: Option<&SpecSampling>,
3941        knobs: &mut Glm5SpecKnobs<'_>,
3942        p_min: f32,
3943        pmin0: bool,
3944        pclk: Option<&mut ProfClock>,
3945    ) -> Res<(Vec<u32>, Glm5DraftQ)> {
3946        let dr = self
3947            .glm5_dflash
3948            .as_ref()
3949            .ok_or("glm5 dflash draft state without a loaded drafter")?;
3950        let draft = &dr.draft;
3951        let c = &draft.cfg;
3952        let b = c.block_size;
3953        let n_embd = self.cfg.n_embd as usize;
3954        let n_vocab = self.output.out_features();
3955        let Glm5SpecSession {
3956            draft: state,
3957            cache,
3958            anchor,
3959            sctr: _,
3960            uctr,
3961            rounds,
3962            prof,
3963            ..
3964        } = sess;
3965        let Glm5DraftState::Dflash2 { kv, pending, taps } = state else {
3966            return Err("glm5_dflash_round_drafts on a native-mtp session".into());
3967        };
3968        let anchor = *anchor;
3969
3970        // ---- 1. ingest pending committed feature rows (positions kv.len..) ----
3971        // Under the default placement the PROMPT rows were ingested at session creation
3972        // and only the kept verify rows of the previous round arrive here; under
3973        // `MEMRA_GLM5_DRAFT_PRIME_LAZY=1` round 1 carries the whole prompt (the pre-lane
3974        // literal, the round-0 wall boot A measured at depth).
3975        let row_w = taps.len() * n_embd;
3976        let mut pclk = pclk;
3977        let n_new = pending.len() / row_w;
3978        let st = if n_new > 0 {
3979            let rows = std::mem::take(pending);
3980            self.glm5_dflash_ingest_rows(eh, draft, kv, &rows, row_w, pclk.as_deref_mut())?
3981        } else {
3982            DraftIngestStats::default()
3983        };
3984        // FIRST-TOKEN PROFILE: the lazy arm's round-1 prompt ingest lands in the
3985        // drafter-prime buckets here (a creation-time ingest already wrote them and
3986        // leaves nothing prompt-sized pending); the clock is re-based either way.
3987        if let (Some(ck), Some(pf)) = (pclk, prof.as_mut()) {
3988            let tail = ck.lap(eh, eh);
3989            if *rounds == 0 && n_new > 0 && pf.draft_prime_arm.is_empty() {
3990                st.write(pf, "eager-lazy");
3991                pf.draft_prime_ms += tail;
3992            }
3993        }
3994        let start = cache.pos;
3995        debug_assert_eq!(
3996            kv.len, start,
3997            "drafter ctx rows must equal committed trunk rows at a round boundary"
3998        );
3999
4000        // ---- 2. block forward over the cached ctx (decode-exact matmul scope: the m=8
4001        // drafter GEMMs otherwise fall into the prefill-GEMM class — the dspark round's
4002        // measured fix; RAII so a `?` exit never latches exact engine-wide) ----
4003        let exact_scope = eh.exact_scope(true);
4004        let mut block: Vec<u32> = vec![c.mask_token_id; b];
4005        block[0] = anchor;
4006        let noise = eh.htod(&self.embd.try_gather(n_embd, &block)?)?;
4007        let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
4008        let dh = draft.forward_round(eh, kv, &noise, &pos_block)?;
4009
4010        // ---- 3. draft logits over the mask-fill harvest rows 1..b ----
4011        let nd = b - 1;
4012        let mut rows_buf = eh.uninit(nd * n_embd)?;
4013        {
4014            let dv = eh.view(&dh, b * n_embd);
4015            let tail = dv.slice(n_embd..b * n_embd);
4016            eh.copy_view_into(&mut rows_buf, 0, &tail, nd * n_embd)?;
4017        }
4018        // TRIMMED DRAFT HEAD (`glm5_dflash_trim`): the dspark serve arm's resolution, the
4019        // FR-Spec self-trim on the MTP struct when it loaded (gathered rows of the target's
4020        // own head), else the `dflash_trim` slab the loader builds for a DFlash2 route that
4021        // boots WITHOUT the NextN block (lane/frspec-dflash2-20260902, the serving shape;
4022        // before it the `MEMRA_FRSPEC_TRIM` contract was a silent no-op here). Neither = the
4023        // full target head, stated in the boot receipt. The draft logits `dl` arrive
4024        // `[nd x n_ranks]`; the propose fns remap candidate RANK ids through d2t before the
4025        // selector walk; verify stays full-vocab, so the slab moves acceptance only.
4026        let trim = self.glm5_dflash_trim();
4027        let (dl_head, dl_vocab) = match trim {
4028            Some((head, d2t)) => (head, d2t.len()),
4029            None => (&self.output, n_vocab),
4030        };
4031        // skip_d2t_remap red arm (the q38 defect made loud): candidates stay RANK ids.
4032        let trim_d2t = trim.filter(|_| !knobs.skip_d2t_remap).map(|(_, d2t)| d2t);
4033        let dl = eh.matmul(dl_head, &rows_buf, nd)?;
4034
4035        // ---- 4. selector walk (greedy chain / sampled candidate-set walk) ----
4036        let (mut drafts, slot_q, qside) = match sp {
4037            None => {
4038                let (path, q) = draft
4039                    .dflash2_propose_greedy_q(eh, &dl, &rows_buf, nd, dl_vocab, anchor, trim_d2t)?;
4040                (path, q, Glm5DraftQ::None)
4041            }
4042            Some(sp) => {
4043                let (path, q_chosen, cand, q_rows) = draft.dflash2_propose_sampled(
4044                    eh, &dl, &rows_buf, nd, dl_vocab, anchor, sp.temp, sp.seed, uctr, trim_d2t,
4045                )?;
4046                let top_k = draft
4047                    .dflash2
4048                    .as_ref()
4049                    .ok_or("glm5 dflash drafter lost its DFlash2 head")?
4050                    .top_k;
4051                (
4052                    path,
4053                    q_chosen.clone(),
4054                    Glm5DraftQ::Selector {
4055                        prop: DsparkDraftSample::Selector {
4056                            cand,
4057                            q_rows,
4058                            q_chosen,
4059                            top_k,
4060                        },
4061                        dl,
4062                    },
4063                )
4064            }
4065        };
4066        drop(exact_scope);
4067        drafts.truncate(k);
4068        // TAU-SLOT CONFIDENCE TRUNCATION (loop-port 2): the low-confidence tail never
4069        // enters verify. Slot-indexed prefix reads keep the retained Selector q side
4070        // consistent (cand/q_rows/q_chosen are per-slot; the accept walk reads slots
4071        // 0..drafts.len()-1 only). p_min unset = today's rounds, untouched.
4072        if p_min > 0.0 {
4073            let kc = glm5_conf_keep(&slot_q[..drafts.len()], p_min, pmin0);
4074            drafts.truncate(kc);
4075        }
4076        if let Some(over) = knobs.draft_override.as_mut() {
4077            for (ki, d) in drafts.iter_mut().enumerate() {
4078                *d = over(*rounds, ki, *d);
4079            }
4080        }
4081        Ok((drafts, qside))
4082    }
4083
4084    /// SAMPLED ACCEPT (module doc): the rejection-sampling walk `u_j * q_j(x_j) < p_j(x_j)`
4085    /// over the verify logit rows — memra's existing sampled spec contract (the
4086    /// MEMRA_SPEC_TEMP route / dspark sampled-admission walk), plugged in at exactly the
4087    /// accept seam; walk and rollback unchanged. p and q take the SAME filter transforms
4088    /// (`filter_stats` + `softmax_gather_filtered`, distribution-exact for the filtered
4089    /// target); the accept-test uniforms come from `spec::host_u01` on the session's `uctr`
4090    /// (tag 0xFFFF_FFFE) and every device draw (draft chain, full-accept bonus, residual
4091    /// resample) advances the session's `sctr` — counters persist on the session so
4092    /// randomness never repeats across bursts. Returns `(j, bonus)`. `e` is the HEAD
4093    /// engine (the round resolves it): the verify rows and retained draft logits live on
4094    /// the last stage under a split.
4095    #[allow(clippy::too_many_arguments)]
4096    // allow: the parameter list mirrors the accept seam's inputs (verify rows + the draft
4097    // chain's retained q side); bundling into a struct would hide the p/q pairing
4098    fn glm5_sampled_accept(
4099        &self,
4100        e: &Engine,
4101        sess: &mut Glm5SpecSession,
4102        sp: &SpecSampling,
4103        vlogits: &CudaSlice<f32>,
4104        drafts: &[u32],
4105        draft_idx: &[u32],
4106        draft_logits: &[CudaSlice<f32>],
4107        draft_stats: &[(f32, f32, f32)],
4108        d2t: Option<&[u32]>,
4109        k: usize,
4110    ) -> Res<(usize, u32)> {
4111        let n_vocab = self.output.out_features();
4112        let d_vocab = d2t.map(|m| m.len()).unwrap_or(n_vocab);
4113        // FILTERED p_j: one batched stats pass over verify rows 0..k-1 (row j is the target
4114        // distribution at draft j's slot), then one batched gather of the drafted tokens.
4115        let rows_i: Vec<i32> = (0..k as i32).collect();
4116        let rows_d = e.htod_i32(&rows_i)?;
4117        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(k)?, e.zeros(k)?, e.zeros(k)?);
4118        e.filter_stats(
4119            vlogits, n_vocab, &rows_d, &mut th_d, &mut z_d, &mut mx_d, n_vocab, k, sp.temp,
4120            sp.top_k, sp.top_p, sp.min_p,
4121        )?;
4122        let ids_d = e.htod_u32_v(drafts)?;
4123        let mut pj_d = e.zeros(k)?;
4124        e.softmax_gather_filtered(
4125            vlogits, n_vocab, &ids_d, &rows_d, &th_d, &z_d, &mut pj_d, n_vocab, k, sp.temp,
4126        )?;
4127        let pj = e.dtoh(&pj_d)?;
4128        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
4129
4130        // The walk: FILTERED q_j from the retained draft logits (rank id for trimmed heads),
4131        // host Philox accept test per slot.
4132        let mut j = 0usize;
4133        while j < k {
4134            let (_qmx, qth, qz) = draft_stats[j];
4135            let idsd = e.htod_u32_v(&[draft_idx[j]])?;
4136            let rows0 = e.htod_i32(&[0])?;
4137            let thd = e.htod(&[qth])?;
4138            let zd = e.htod(&[qz])?;
4139            let mut outd = e.zeros(1)?;
4140            e.softmax_gather_filtered(
4141                &draft_logits[j],
4142                d_vocab,
4143                &idsd,
4144                &rows0,
4145                &thd,
4146                &zd,
4147                &mut outd,
4148                d_vocab,
4149                1,
4150                sp.temp,
4151            )?;
4152            let qj = e.dtoh(&outd)?[0];
4153            let u = crate::spec::host_u01(sp.seed, sess.uctr);
4154            sess.uctr = sess.uctr.wrapping_add(1);
4155            if (u as f64) * (qj as f64) < pj[j] as f64 {
4156                j += 1;
4157            } else {
4158                break;
4159            }
4160        }
4161
4162        // Bonus: full accept draws a filtered Gumbel sample from the LAST verify row
4163        // (`glm5_sampled_bonus` — shared with the PMIN0 zero-draft round); rejection at j
4164        // resamples the residual norm(max(0, fp_j - fq_j)) — with a trimmed draft head, q
4165        // scatters back to full vocab first (`scatter_trim_logits`).
4166        if j == k {
4167            return Ok((
4168                j,
4169                self.glm5_sampled_bonus(e, sess, sp, vlogits, k, n_vocab)?,
4170            ));
4171        }
4172        let mut col = e.zeros(n_vocab)?;
4173        let bonus = {
4174            let vv = e.view(vlogits, (k + 1) * n_vocab);
4175            let row = vv.slice(j * n_vocab..(j + 1) * n_vocab);
4176            e.copy_view_into(&mut col, 0, &row, n_vocab)?;
4177            let p_stats = (mxv[j], thv[j], zv[j]);
4178            let q_stats = draft_stats[j];
4179            let sc = sess.sctr;
4180            sess.sctr = sess.sctr.wrapping_add(1);
4181            let mut sample_tok = e.alloc_u32_zeroed(1)?;
4182            match d2t {
4183                Some(map) => {
4184                    let map_d = e.htod_u32_v(map)?;
4185                    let mut q_full = e.zeros(n_vocab)?;
4186                    e.scatter_trim_logits(&draft_logits[j], &map_d, &mut q_full, d_vocab, n_vocab)?;
4187                    e.residual_sample_filtered(
4188                        &col,
4189                        Some(&q_full),
4190                        n_vocab,
4191                        sp.temp,
4192                        sp.seed,
4193                        sc,
4194                        p_stats,
4195                        q_stats,
4196                        &mut sample_tok,
4197                    )?;
4198                }
4199                None => {
4200                    e.residual_sample_filtered(
4201                        &col,
4202                        Some(&draft_logits[j]),
4203                        n_vocab,
4204                        sp.temp,
4205                        sp.seed,
4206                        sc,
4207                        p_stats,
4208                        q_stats,
4209                        &mut sample_tok,
4210                    )?;
4211                }
4212            }
4213            e.dtoh_u32(&sample_tok)?[0]
4214        };
4215        let bonus = crate::spec::guard_vocab_token(
4216            bonus,
4217            n_vocab,
4218            &format!("glm5 sampled verify bonus at round {} j={j}", sess.rounds),
4219        )?;
4220        Ok((j, bonus))
4221    }
4222
4223    /// One filtered-Gumbel bonus draw from verify row `row` through the session's device
4224    /// Philox stream — the sampled FULL-ACCEPT bonus, and the entire accept of a PMIN0
4225    /// zero-draft round (whose verify batch is just the anchor row: m=1 = a plain sampled
4226    /// decode step). Advances `sctr` exactly once; byte-for-byte the pre-extraction
4227    /// full-accept arm of `glm5_sampled_accept`.
4228    fn glm5_sampled_bonus(
4229        &self,
4230        e: &Engine,
4231        sess: &mut Glm5SpecSession,
4232        sp: &SpecSampling,
4233        vlogits: &CudaSlice<f32>,
4234        row: usize,
4235        n_vocab: usize,
4236    ) -> Res<u32> {
4237        let mut col = e.zeros(n_vocab)?;
4238        let vv = e.view(vlogits, (row + 1) * n_vocab);
4239        let src = vv.slice(row * n_vocab..(row + 1) * n_vocab);
4240        e.copy_view_into(&mut col, 0, &src, n_vocab)?;
4241        let rows0 = e.htod_i32(&[0])?;
4242        let (mut bth, mut bz, mut bmx) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
4243        e.filter_stats(
4244            &col, n_vocab, &rows0, &mut bth, &mut bz, &mut bmx, n_vocab, 1, sp.temp, sp.top_k,
4245            sp.top_p, sp.min_p,
4246        )?;
4247        let (th, mx) = (e.dtoh(&bth)?[0], e.dtoh(&bmx)?[0]);
4248        let mut pb = e.zeros(n_vocab)?;
4249        e.gumbel_perturb_filtered(&col, &mut pb, n_vocab, sp.seed, sess.sctr, sp.temp, mx, th)?;
4250        sess.sctr = sess.sctr.wrapping_add(1);
4251        let td = e.argmax_token_device(&pb, n_vocab)?;
4252        crate::spec::guard_vocab_token(
4253            e.dtoh_u32_one(&td)?,
4254            n_vocab,
4255            &format!(
4256                "glm5 sampled verify bonus at round {} (row {row})",
4257                sess.rounds
4258            ),
4259        )
4260    }
4261}
4262
4263/// One filtered Gumbel draw from a draft-head logit row through the session's device Philox
4264/// stream — the sampled route's PROPOSAL. Returns the drawn RANK id and the row's filtered
4265/// stats `(row_max, threshold_e, renorm_mass)`, which the accept walk's q gather and the
4266/// rejection residual both reuse (the q side must be the distribution the draft was actually
4267/// drawn from, or rejection sampling is not exact for the filtered target).
4268fn glm5_sampled_draft(
4269    e: &Engine,
4270    dl: &CudaSlice<f32>,
4271    d_vocab: usize,
4272    sp: &SpecSampling,
4273    sctr: &mut u32,
4274) -> Res<(u32, (f32, f32, f32))> {
4275    let rows0 = e.htod_i32(&[0])?;
4276    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
4277    e.filter_stats(
4278        dl, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab, 1, sp.temp, sp.top_k,
4279        sp.top_p, sp.min_p,
4280    )?;
4281    let (th, z, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
4282    let mut pb = e.zeros(d_vocab)?;
4283    e.gumbel_perturb_filtered(dl, &mut pb, d_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
4284    *sctr = sctr.wrapping_add(1);
4285    let td = e.argmax_token_device(&pb, d_vocab)?;
4286    let idx =
4287        crate::spec::guard_vocab_token(e.dtoh_u32_one(&td)?, d_vocab, "glm5 sampled draft draw")?;
4288    Ok((idx, (mx, th, z)))
4289}
4290
4291/// glm5_next SERVED speculative session (lane/glm5-spec-routing, 2026-08-30): the state one
4292/// request's spec decoding carries across worker bursts — the dspark/gemma session twins'
4293/// shape. The session OWNS its trunk cache (the worker's `s.cache` stays `None`); at every
4294/// burst boundary the invariant is `cache.pos == committed.len()` with each committed row's
4295/// trunk state exactly what a plain prime of that sequence would hold (the accept walk's
4296/// basis, pinned by the tparallel gate), plus ONE emitted-but-uncommitted `anchor` token
4297/// (the next round's verify row 0 — the dspark `last` convention).
4298pub struct Glm5SpecSession {
4299    cache: Cache,
4300    /// Every token whose trunk state the cache holds, in order (prompt + committed
4301    /// generation). EXCLUDES the live `anchor`.
4302    pub committed: Vec<u32>,
4303    /// The last emitted token, not yet consumed by the trunk — round anchor / verify row 0.
4304    anchor: u32,
4305    /// The prime's boundary token is emitted exactly once, by the first burst.
4306    anchor_emitted: bool,
4307    /// Committed `(token, h_seed)` pairs not yet fed to the MTP draft plane; the last
4308    /// feed's logits double as the next round's first draft (the re-warm contract).
4309    /// NATIVE-MTP arm only; the DFlash2 source keeps its own pending rows in `draft`.
4310    pending: Vec<(u32, CudaSlice<f32>)>,
4311    /// The session's pinned draft source + its state (module doc, DRAFT SOURCE SEAM).
4312    draft: Glm5DraftState,
4313    /// `None` / `temp <= 0` = greedy byte-contract route. Fixed for the session — the
4314    /// worker's admission owns the sampler identity.
4315    sampling: Option<SpecSampling>,
4316    /// PENALTY ARM (lane/spec-exclusions-20260902, `MEMRA_SPEC_PENALTY=1`): the request's
4317    /// penalty config when it carries one — independent of `sampling` (a greedy request
4318    /// with penalties is `sampling: None, pen: Some`). `None` = no penalty pass anywhere on
4319    /// the round, the pre-lane program byte for byte.
4320    pen: Option<Glm5Penalty>,
4321    /// The session-spanning penalty history (`pen_window_seed` over the prompt, then the
4322    /// anchor at every round start and the accepted drafts after every accept — the dspark
4323    /// `pen_hist` convention); the verify rows penalize over its last `win()` entries.
4324    /// Empty with the arm off.
4325    pen_hist: Vec<u32>,
4326    /// Session-continuity Philox counters (never reset across bursts): `sctr` = device
4327    /// sampling events (boundary, draft chain, bonus, residual), `uctr` = host accept-test
4328    /// uniforms (`spec::host_u01`, tag 0xFFFF_FFFE).
4329    sctr: u32,
4330    uctr: u32,
4331    /// Verify rounds completed over the session lifetime (the worker's per-burst
4332    /// rounds-delta receipt, the dspark `rounds` convention).
4333    pub rounds: usize,
4334    /// Of `rounds`, those whose drafts came through a RANK-TRIMMED draft head (the DFlash2
4335    /// slab / MtpHead self-trim / native trimmed chain; lane/frspec-dflash2-20260902). Equals
4336    /// `rounds` when a trim is loaded, 0 otherwise, the counter the `[glm5-acc]` receipt
4337    /// and the rig gate read; a process-wide twin is `glm5_rank_trimmed_draft_rounds`.
4338    pub rank_trimmed_rounds: usize,
4339    done: bool,
4340    max_ctx: usize,
4341    /// MTP draft-plane layer index — `Some` on the native-MTP arm only.
4342    mtp_il: Option<usize>,
4343    /// Prompt-boundary capture for the DEFERRED prefix publication (lane/glm5-prefix-latent2,
4344    /// 2026-09-01; the dspark `prefix_capture` pattern): taken at session creation before any
4345    /// burst mutates the recurrent/tail state, drained by the worker's sweep. `None` when the
4346    /// worker did not request capture or when any boundary invariant refused.
4347    prefix_capture: Option<crate::spec::SpecBoundaryCapture>,
4348    /// FIRST-TOKEN PROFILE (lane/b200-spec-ttft-20260902): `Some` iff `MEMRA_SPEC_PROF=1`
4349    /// at creation; filled through session creation and round 1 of the first burst, then
4350    /// drained by the worker's one `[spec-prof]` line (`take_first_token_prof`).
4351    prof: Option<SpecFirstTokenProf>,
4352    /// DEPTH LOG (lane/spec-route-depth-20260902): the first `SPEC_PROF_ROUNDS` rounds'
4353    /// attribution rows, `Some` iff `MEMRA_SPEC_PROF=1`; the worker prints fresh rows after
4354    /// every burst and the summary once the log fills or the session ends.
4355    prof_rounds: Option<SpecRoundsLog>,
4356}
4357
4358impl Glm5SpecSession {
4359    /// The depth log, for the worker's `[spec-prof-rounds]` / `[spec-prof-summary]` lines.
4360    pub fn round_log_mut(&mut self) -> Option<&mut SpecRoundsLog> {
4361        self.prof_rounds.as_mut()
4362    }
4363    /// Rounds the depth log keeps (the print cadence's cap).
4364    pub fn round_log_cap() -> usize {
4365        SPEC_PROF_ROUNDS
4366    }
4367    /// Drafter ctx rows currently ingested (`None` on the native arm).
4368    pub fn draft_kv_len(&self) -> Option<usize> {
4369        match &self.draft {
4370            Glm5DraftState::Dflash2 { kv, .. } => Some(kv.len),
4371            _ => None,
4372        }
4373    }
4374    /// GATE ACCESSOR (rig gate 15): the first `rows` rows of every drafter K and V plane,
4375    /// host-side, `row_floats` = n_kv * head_dim. `None` on the native arm.
4376    #[allow(clippy::type_complexity)]
4377    // allow: (k planes, v planes) is the natural shape for a bit-identity diff
4378    pub fn draft_kv_rows_host(
4379        &self,
4380        e: &Engine,
4381        rows: usize,
4382        row_floats: usize,
4383    ) -> Option<(Vec<Vec<f32>>, Vec<Vec<f32>>)> {
4384        let Glm5DraftState::Dflash2 { kv, .. } = &self.draft else {
4385            return None;
4386        };
4387        let n = rows * row_floats;
4388        let take = |planes: &[CudaSlice<f32>]| -> Option<Vec<Vec<f32>>> {
4389            planes
4390                .iter()
4391                .map(|p| e.dtoh_view(&p.slice(0..n)).ok())
4392                .collect()
4393        };
4394        Some((take(&kv.k)?, take(&kv.v)?))
4395    }
4396    /// Context capacity of the session's cache (the server's ContextFull guard).
4397    pub fn cache_max_ctx(&self) -> usize {
4398        self.max_ctx
4399    }
4400    /// Drain the first-token profile (doc on the field). `None` with the profile off or
4401    /// once drained — the worker prints exactly one line per request.
4402    pub fn take_first_token_prof(&mut self) -> Option<SpecFirstTokenProf> {
4403        self.prof.take()
4404    }
4405    /// Drain the prompt-boundary capture (doc on the field; the dspark
4406    /// `take_prefix_capture` twin — the worker publishes it against `cache_ref`).
4407    pub fn take_prefix_capture(&mut self) -> Option<crate::spec::SpecBoundaryCapture> {
4408        self.prefix_capture.take()
4409    }
4410    /// True when the deferred prefix capture can publish NOW: a capture exists AND the
4411    /// DFlash2 drafter KV already covers the boundary. glm5 defers the prompt's feature
4412    /// ingest to round 1 (unlike dspark's at-creation ingest), so a drain that fired before
4413    /// the first burst would export an empty tail and waste the capture — the worker's
4414    /// sweep polls this instead.
4415    pub fn prefix_capture_ready(&self) -> bool {
4416        self.prefix_capture
4417            .as_ref()
4418            .is_some_and(|c| match &self.draft {
4419                Glm5DraftState::Dflash2 { kv, .. } => kv.len >= c.pos,
4420                _ => false,
4421            })
4422    }
4423    /// The drafter's readable KV tail at `upto` rows (the dspark `export_tail` seam) —
4424    /// `None` on the native arm or when the tail cannot cover the drafter window.
4425    pub fn export_draft_tail(
4426        &self,
4427        e: &Engine,
4428        upto: usize,
4429    ) -> Option<crate::dflash::DflashKvTail> {
4430        match &self.draft {
4431            Glm5DraftState::Dflash2 { kv, .. } => kv.export_tail(e, upto),
4432            _ => None,
4433        }
4434    }
4435    /// The session's trunk cache, for the worker's deferred prefix publication (the
4436    /// append-only-below-boundary slices) — never for mutation.
4437    pub fn cache_ref(&self) -> &Cache {
4438        &self.cache
4439    }
4440    /// Trunk rows currently committed (== `committed.len()` at burst boundaries).
4441    pub fn pos(&self) -> usize {
4442        self.cache.pos
4443    }
4444    /// EOS committed or the context guard tripped: the next burst would emit nothing.
4445    pub fn finished(&self) -> bool {
4446        self.done
4447    }
4448    /// True when the session is a legal demotion source (loop-port fold-in, map #8):
4449    /// GREEDY, UNPENALIZED only. A sampled session's committed stream depends on its
4450    /// session-owned Philox counters, and the plain batched sampler is a different random
4451    /// program mid-request (the exact exclusion the MTP and dspark sweeps carry). A
4452    /// penalized session (`pen: Some`, lane/spec-exclusions-20260902) is excluded the same
4453    /// way: the demotion flush (`glm5_spec_into_demoted`) argmaxes the boundary row RAW —
4454    /// no penalty pass — so demoting one would silently drop the request's penalties for
4455    /// that token, the exact failure class `glm5_penalty_admit`'s refusal exists to
4456    /// prevent, and dspark's structural reason for keeping penalized-greedy off spec
4457    /// entirely (revuto finding on the re-land, 2026-09-03). Penalized sessions therefore
4458    /// stay on spec until they end, like sampled ones.
4459    pub fn demote_eligible(&self) -> bool {
4460        self.sampling.is_none() && self.pen.is_none()
4461    }
4462}
4463
4464impl HybridModel {
4465    /// ONE-WAY DEMOTION HANDOFF for the glm5 session (loop-port fold-in — the map's #8,
4466    /// the `SpecSession::into_demoted` / `DsparkSpecSession::into_demoted` twin): consume
4467    /// the session and hand `(cache, next_pred)` to the plain batched-decode path, so a
4468    /// spec session admitted on a quiet box stops serializing the tick when load arrives
4469    /// (the spec-gate HIGH sweep's ship-safety lever; dspark receipt: "c=8 429.6 = parity
4470    /// (pre-lane -37%)").
4471    ///
4472    /// THE ANCHOR IS THE CARRIED-PENDING SHAPE: glm5 emits each round's bonus immediately
4473    /// (`round_tokens` include it) while the trunk consumes it only as the NEXT round's
4474    /// row 0 — so at every burst boundary the session holds ONE emitted-but-uncommitted
4475    /// token. Handing the cache over as-is would leave it one row short of the public
4476    /// stream, and `device_next` re-emitting the anchor would duplicate a served token.
4477    /// The flush below is `spec_flush_pending`'s exact analogue: ONE plain T=1 decode
4478    /// step commits the anchor (byte-identical to the never-drafted chain — the
4479    /// tparallel gate's accept-j-then-continue identity IS this claim), and its argmax
4480    /// becomes the handoff's `next_pred` — a token the batched path emits and feeds
4481    /// exactly as it would its own. One trunk pass, once per demotion, never per burst.
4482    ///
4483    /// ONE-WAY BY DESIGN: the draft state (MTP pending pairs / DFlash2 drafter KV and
4484    /// feature rows) and the Philox counters are DROPPED, freeing their VRAM; there is
4485    /// no cheap symmetric re-promotion (the spec.rs law, verbatim). Sampled sessions
4486    /// refuse loudly (`demote_eligible`; the worker's sweep excludes them first).
4487    pub fn glm5_spec_into_demoted(
4488        &self,
4489        e: &Engine,
4490        mut sess: Glm5SpecSession,
4491    ) -> Res<(Cache, u32)> {
4492        if !sess.demote_eligible() {
4493            let why = if sess.sampling.is_some() {
4494                "sampled sessions stay on spec until they end (session-owned Philox vs the \
4495                 worker sampler is an unmeasured distributional seam — the MTP sweep's \
4496                 exclusion, verbatim)"
4497            } else {
4498                "penalized sessions stay on spec until they end (this flush's plain argmax \
4499                 carries no penalty pass; demoting would silently drop the request's \
4500                 penalties for that token, lane/spec-exclusions-20260902)"
4501            };
4502            return Err(format!("glm5 demote: {why}").into());
4503        }
4504        if sess.cache.pos + 1 > sess.max_ctx {
4505            return Err(format!(
4506                "glm5 demote: no room to flush the live anchor ({} + 1 > ctx {})",
4507                sess.cache.pos, sess.max_ctx
4508            )
4509            .into());
4510        }
4511        let logits = self.decode_step(e, sess.anchor, &mut sess.cache)?;
4512        sess.committed.push(sess.anchor);
4513        let next = argmax(&logits) as u32;
4514        Ok((sess.cache, next))
4515    }
4516}
4517
4518/// Buckets of one eager drafter ingest (`glm5_dflash_ingest_rows`), written into the
4519/// first-token profile under the arm name that ran it.
4520#[derive(Default, Debug, Clone, Copy)]
4521struct DraftIngestStats {
4522    h2d_ms: f64,
4523    feat_ms: f64,
4524    kv_ms: f64,
4525    rows: usize,
4526    chunks: usize,
4527}
4528
4529impl DraftIngestStats {
4530    fn write(&self, pf: &mut SpecFirstTokenProf, arm: &'static str) {
4531        pf.draft_prime_ms = self.h2d_ms + self.feat_ms + self.kv_ms;
4532        pf.draft_prime_h2d_ms = self.h2d_ms;
4533        pf.draft_prime_feat_ms = self.feat_ms;
4534        pf.draft_prime_kv_ms = self.kv_ms;
4535        pf.draft_prime_rows = self.rows;
4536        pf.draft_prime_chunks = self.chunks;
4537        pf.draft_prime_arm = arm;
4538    }
4539}
4540
4541/// Gate instruments for `generate_spec_glm5_gated`. Documented as instruments: no serving
4542/// path constructs a non-default value.
4543#[derive(Default)]
4544pub struct Glm5SpecKnobs<'a> {
4545    /// `(round, draft_index, greedy_draft) -> draft` — deterministic forced-accept /
4546    /// forced-reject rounds for the end-to-end gate.
4547    pub draft_override: Option<&'a mut dyn FnMut(usize, usize, u32) -> u32>,
4548    /// RED ARM ONLY: skip the state rollback (pos still moves). A corrupted draft must then
4549    /// leave post-row-K KDA state and un-truncated latent rows behind — the end-to-end gate
4550    /// asserts the tape DIVERGES from plain decode (or the kpool residency tripwire fires).
4551    pub disable_rollback: bool,
4552    /// RED ARM ONLY: with an FR-Spec trim loaded, use the draft argmax RANK id as the vocab
4553    /// id (the q38 skipped-remap defect: 0/248 acceptance with every exactness gate green).
4554    /// The gate asserts the drafted sequence diverges from the untrimmed arm's while the
4555    /// output tape STAYS byte-identical to plain decode — the silent failure made loud.
4556    pub skip_d2t_remap: bool,
4557    /// GATE INSTRUMENT for the confidence gate (loop-port 2): `Some((p_min, pmin0))`
4558    /// overrides the `MEMRA_SPEC_PMIN`/`MEMRA_SPEC_PMIN0` env pair for this call — the
4559    /// env statics latch once per process, so the byte-identity gate drives its PMIN
4560    /// arms through here instead of the environment. `None` = the serving resolution.
4561    pub pmin_override: Option<(f32, bool)>,
4562    /// GATE INSTRUMENT (lane/glm5-accrace): trace every GREEDY round's accept decision to
4563    /// stderr as one `[glm5-accrace]` line — round, t, j, the drafts, the DEVICE argmax
4564    /// row (`argmax_token_device_col` + the one u32 readback the accept walk consumes), a
4565    /// HOST argmax over the same `vlogits` buffer, and a per-row (argmax, FNV-1a hash of
4566    /// the row's f32 bits) census.
4567    ///
4568    /// TWO THINGS IT SEPARATES, which is why it exists: (a) `dev != host` means the device
4569    /// accept path published a value the logits buffer does not justify (a readback/scratch
4570    /// race); (b) `dev == host` with a row hash that moves between two runs of the same
4571    /// deterministic fixture means the verify logits themselves were computed over
4572    /// corrupted state (an upstream walk/rollback race). The host read is issued AFTER the
4573    /// device path's own `dtoh_u32` has already synchronized the consuming stream, so the
4574    /// probe can only observe the race, never mask it.
4575    ///
4576    /// Never a serving surface: no serving path constructs a non-default value.
4577    pub accept_probe: bool,
4578}
4579
4580#[cfg(test)]
4581mod conf_keep_tests {
4582    use super::glm5_conf_keep;
4583
4584    /// The spec.rs chain-break semantics, pinned CPU-side (loop-port 2): break at the
4585    /// first sub-threshold slot; slot 0 survives a miss unless PMIN0.
4586    #[test]
4587    fn conf_keep_matches_the_spec_rs_break_semantics() {
4588        // Gate off: everything kept.
4589        assert_eq!(glm5_conf_keep(&[0.1, 0.1], 0.0, true), 2);
4590        // All confident: everything kept.
4591        assert_eq!(glm5_conf_keep(&[0.9, 0.8, 0.7], 0.5, false), 3);
4592        // Break mid-chain at the first miss; the confident tail after it never rides
4593        // (prefix truncation — the accept rule could never commit past the gap anyway).
4594        assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, false), 1);
4595        assert_eq!(glm5_conf_keep(&[0.9, 0.2, 0.9], 0.5, true), 1);
4596        // Slot-0 miss: survives without PMIN0 (the j > 0 arm of the break condition), and
4597        // does NOT latch — slot 1 is judged on its own confidence (the spec.rs chain
4598        // evaluates each slot's p independently)...
4599        assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, false), 2);
4600        // ...but a sub-threshold slot past 0 still breaks.
4601        assert_eq!(glm5_conf_keep(&[0.2, 0.2], 0.5, false), 1);
4602        // PMIN0 arms the zero-draft round.
4603        assert_eq!(glm5_conf_keep(&[0.2, 0.9], 0.5, true), 0);
4604        // Boundary: q == p_min is NOT below it (strict <, the spec.rs test).
4605        assert_eq!(glm5_conf_keep(&[0.5, 0.5], 0.5, true), 2);
4606        // Empty chain: nothing to keep.
4607        assert_eq!(glm5_conf_keep(&[], 0.5, true), 0);
4608    }
4609}