Skip to main content

memra_engine/
hybrid_forward.rs

1//! Hybrid forward pass (Stage-1, f32, prefill, single sequence). Per layer dispatches to a
2//! linear-attention (Gated DeltaNet) or full-attention mixer, then SwiGLU FFN. Matches
3//! llama.cpp src/models/qwen35.cpp node-for-node.
4
5// lane/clippy-zero-restore-20260901: index loops here mirror the llama.cpp reference
6// node-for-node (header above); iterator reshapes are not bit-neutral by inspection.
7#![allow(clippy::needless_range_loop)]
8
9use crate::Engine;
10use crate::cache::Cache;
11use cudarc::driver::CudaSlice;
12use memra_gguf::config::{ModelConfig, SwigluClamp};
13
14/// Resident trunk transients for the eager prime (piecewise-graph foundation; see
15/// HybridModel::prime_slabs). Every live buffer prefix is fully overwritten before use per prime;
16/// capacity beyond the current token count must never cross a shape-sensitive boundary.
17pub struct PrimeSlabs {
18    pub t_cap: usize,
19    pub h: CudaSlice<f32>,
20    pub x1: CudaSlice<f32>,
21    pub z: CudaSlice<f32>,
22    pub act: CudaSlice<f32>,
23    pub xa: CudaSlice<f32>,
24    pub xb: CudaSlice<f32>,
25    pub h16: CudaSlice<u8>,
26    pub z16: CudaSlice<u8>,
27    /// piecewise boundary slabs (increment 2): GEMM outputs land here so the
28    /// downstream captured segments see fixed addresses.
29    pub gate: CudaSlice<f32>, // t * n_ff_max
30    pub up: CudaSlice<f32>,      // t * n_ff_max
31    pub ffn_out: CudaSlice<f32>, // t * n_embd
32    /// piecewise increment 3: per-layer S-glue segment graphs (down-add + next
33    /// attn-norm, ALL-slab IO, zero in-graph allocations -> keeperless capture is
34    /// clean). Baked at this t_cap; replay only when t == t_cap. seg_glue[il] fires
35    /// between layer il and il+1 (ping-pong parity is deterministic per il).
36    pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
37    /// increment 5 (core-split edition): the mixer out-GEMM writes _into_ `mixed`
38    /// directly (no staging copy — the increment-4 copy route was refuted), making
39    /// S-mid [add + post-norm] all-slab and capturable.
40    pub mixed: CudaSlice<f32>,
41    pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
42    pub seg_t: usize,
43}
44
45// Split prime ranges cannot enter the full-range segment-graph arm, and every slab access
46// is serialized by its device mutex after binding that device's CUDA context on the thread.
47unsafe impl Send for PrimeSlabs {}
48
49/// Shared-expert gate+up at t==1: NVFP4 fused2 (the ornith15/qwen35moe NVFP4 mints keep
50/// gate/up_shexp uniformly NVFP4, so the Q8-only fused2 never fired there and the pair fell
51/// to two mr2 singles + two re-quantizes of the same z — 2 of the 8 unfused launches/layer
52/// the orndecode B=1 census ranked at 17.1%), else the Q8_0 fused2 (the Q8 35B mint), else
53/// two singles. ONE helper for all three shexp dispatch sites — the MEMRA_GDN_MMA
54/// three-read-sites defect is the precedent for not inlining this thrice. Fusion law
55/// everywhere: per (tensor,row) the fused seg body is verbatim, so fused == singles
56/// bit-identically, and the shared (zq, zd) is the same quantize each single recomputes.
57fn shexp_gate_up_t1(
58    e: &Engine,
59    gate_shexp: &crate::model::GpuTensor,
60    up_shexp: &crate::model::GpuTensor,
61    z: &CudaSlice<f32>,
62    zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
63) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
64    let is_nvfp4 = |w: &crate::model::GpuTensor| matches!(w, crate::model::GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_NVFP4);
65    if is_nvfp4(gate_shexp) && is_nvfp4(up_shexp) {
66        // Reuse the caller's t==1 z-quantize when one exists (the zq8 seam the dev arm
67        // already consumes) — the helper's own quantize is the identical kernel on the
68        // identical input, so this drops one launch per MoE layer without moving a byte.
69        let pair = match zq8 {
70            Some((zq, zd)) => e.matmul_nvfp4_fused2(gate_shexp, up_shexp, zq, zd, 1)?,
71            None => {
72                let (zq, zd) = e.quantize_q8_1(z, 1, gate_shexp.in_features())?;
73                e.matmul_nvfp4_fused2(gate_shexp, up_shexp, &zq, &zd, 1)?
74            }
75        };
76        if let Some(pair) = pair {
77            return Ok(pair);
78        }
79    }
80    match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
81        Some(pair) => Ok(pair),
82        None => Ok((e.matmul(gate_shexp, z, 1)?, e.matmul(up_shexp, z, 1)?)),
83    }
84}
85
86fn active_matrix_values(
87    available: usize,
88    rows: usize,
89    columns: usize,
90    label: &str,
91) -> Result<usize, String> {
92    let required = rows
93        .checked_mul(columns)
94        .ok_or_else(|| format!("{label} shape overflows: {rows}x{columns}"))?;
95    if available < required {
96        return Err(format!(
97            "{label} has {available} values, fewer than the active {rows}x{columns} ({required})"
98        ));
99    }
100    Ok(required)
101}
102
103fn step_grouped_decode_shape(prefill: bool, tokens: usize) -> bool {
104    !prefill && tokens == 1
105}
106
107fn parse_step_ep_grouped_prefill(value: Option<&str>) -> Result<bool, String> {
108    match value {
109        None | Some("") | Some("0") => Ok(false),
110        Some("1") => Ok(true),
111        Some(value) => Err(format!(
112            "MEMRA_STEP_EP_GROUPED_PREFILL={value:?} is invalid; expected 0 or 1"
113        )),
114    }
115}
116
117fn step_ep_grouped_prefill_enabled() -> Result<bool, String> {
118    parse_step_ep_grouped_prefill(
119        std::env::var("MEMRA_STEP_EP_GROUPED_PREFILL")
120            .ok()
121            .as_deref(),
122    )
123}
124
125fn step_grouped_prefill_shape(enabled: bool, prefill: bool, tokens: usize) -> bool {
126    enabled && prefill && (PRIME_MIN_T..=crate::cache::PRIME_CHUNK_MAX_TOKENS).contains(&tokens)
127}
128
129fn parse_step_tp_prefill(value: Option<&str>) -> Result<bool, String> {
130    match value {
131        None | Some("") | Some("0") => Ok(false),
132        Some("1") => Ok(true),
133        Some(value) => Err(format!(
134            "MEMRA_STEP_TP_PREFILL={value:?} is invalid; expected 0 or 1"
135        )),
136    }
137}
138
139fn step_tp_prefill_enabled() -> Result<bool, String> {
140    parse_step_tp_prefill(std::env::var("MEMRA_STEP_TP_PREFILL").ok().as_deref())
141}
142
143fn validate_step_prime_batch_modes(tp_prefill: bool, grouped_prefill: bool) -> Result<(), String> {
144    if grouped_prefill && !tp_prefill {
145        return Err("MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into());
146    }
147    if tp_prefill {
148        return Err(
149            "Step TP4 cross-request prime batching did not clear the live-server performance \
150             gate; use per-session grouped prefill"
151                .into(),
152        );
153    }
154    Ok(())
155}
156
157fn step_tp_prefill_shape(
158    enabled: bool,
159    tokens: usize,
160    ranks: usize,
161    native_p2p: bool,
162    has_rank_local_attention: bool,
163    fp8_kv: bool,
164) -> bool {
165    // TP2 admitted 2026-08-25 behind the same off-by-default door. The prefill body is
166    // rank-count-generic (every geometry check divides by `ranks`); only this shape ever
167    // named 4. TP4's high-context NO-GO was a TRANSPORT verdict — token-row column
168    // gathers plus remote O blocks issue ~61,440 peer copies per 4K attention layer on
169    // that placement — which is not evidence about a 2-card native-P2P placement that
170    // reduces O rank-locally. TP2 is UNQUALIFIED until its own prefill argmax + TTFT
171    // receipts land; the door stays off by default.
172    enabled
173        && tokens >= PRIME_MIN_T
174        && matches!(ranks, 2 | 4)
175        && native_p2p
176        && has_rank_local_attention
177        && !fp8_kv
178}
179
180fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
181    std::iter::repeat_with(|| None).take(n).collect()
182}
183
184fn prime_cache_stage_for_layer(fence: &[usize], layer: usize) -> usize {
185    debug_assert!(fence.len() >= 3);
186    match fence[1..fence.len() - 1].binary_search(&layer) {
187        Ok(index) => index + 1,
188        Err(index) => index,
189    }
190}
191
192fn move_prime_cache_layers<T>(
193    parent: &mut [Option<T>],
194    stages: &mut [Vec<Option<T>>],
195    fence: &[usize],
196) {
197    assert_eq!(stages.len() + 1, fence.len());
198    assert!(stages.iter().all(|stage| stage.len() == parent.len()));
199    for (layer, value) in parent.iter_mut().enumerate() {
200        let stage = prime_cache_stage_for_layer(fence, layer);
201        debug_assert!(stages[stage][layer].is_none());
202        stages[stage][layer] = value.take();
203    }
204}
205
206#[cfg(test)]
207fn restore_prime_cache_layers<T>(
208    parent: &mut [Option<T>],
209    stages: &mut [Vec<Option<T>>],
210    fence: &[usize],
211) {
212    assert_eq!(stages.len() + 1, fence.len());
213    assert!(stages.iter().all(|stage| stage.len() == parent.len()));
214    for (layer, value) in parent.iter_mut().enumerate() {
215        let stage = prime_cache_stage_for_layer(fence, layer);
216        debug_assert!(value.is_none());
217        *value = stages[stage][layer].take();
218    }
219}
220
221/// Temporarily move a PP cache's layer state into independently-owned stage shells. The stage
222/// walkers then receive disjoint `&mut Cache` values and can run on separate host threads without
223/// aliasing. GPU buffers are moved, not copied; Drop restores every layer and publishes the last
224/// position completed by every stage.
225struct PrimeCacheStages<'a> {
226    parent: &'a mut Cache,
227    fence: Vec<usize>,
228    stages: Vec<std::sync::Mutex<Cache>>,
229    committed: bool,
230}
231
232impl<'a> PrimeCacheStages<'a> {
233    fn new(parent: &'a mut Cache, fence: &[usize]) -> Self {
234        let n = parent.kv.len();
235        assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
236        assert_eq!(parent.tp_kv.len(), n, "cache layer vectors disagree");
237        assert_eq!(parent.latent.len(), n, "cache layer vectors disagree");
238        let n_stages = fence.len().checked_sub(1).expect("PP cache fence is empty");
239        assert!((2..=4).contains(&n_stages), "PP cache needs 2..=4 stages");
240        assert_eq!(fence[0], 0, "PP cache fence must start at layer zero");
241        assert!(
242            fence.windows(2).all(|pair| pair[0] < pair[1]),
243            "PP cache fence must be strictly increasing"
244        );
245        assert!(fence[n_stages] <= n, "PP cache fence exceeds {n} layers");
246
247        let mut latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
248        let mut g5_recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
249        let mut g5_latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
250        let mut kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
251        let mut tp_kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
252        let mut recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
253        move_prime_cache_layers(&mut parent.kv, &mut kv, fence);
254        move_prime_cache_layers(&mut parent.tp_kv, &mut tp_kv, fence);
255        move_prime_cache_layers(&mut parent.recur, &mut recur, fence);
256
257        move_prime_cache_layers(&mut parent.latent, &mut latent, fence);
258        move_prime_cache_layers(&mut parent.glm5_tp_recur, &mut g5_recur, fence);
259        move_prime_cache_layers(&mut parent.glm5_tp_latent_peer, &mut g5_latent, fence);
260        let pos = parent.pos;
261        let max_ctx = parent.max_ctx;
262        // Indexed rather than zipped: six per-stage layer vectors (main's kv/tp_kv/recur plus
263        // this lane's latent/glm5_tp_recur/glm5_tp_latent_peer) do not read as a zip chain, and
264        // a nested-tuple pattern is exactly where a field silently lands on the wrong stage.
265        let stages = (0..n_stages)
266            .map(|stage| {
267                std::sync::Mutex::new(Cache {
268                    kv: std::mem::take(&mut kv[stage]),
269                    tp_kv: std::mem::take(&mut tp_kv[stage]),
270                    recur: std::mem::take(&mut recur[stage]),
271                    latent: std::mem::take(&mut latent[stage]),
272                    glm5_tp_recur: std::mem::take(&mut g5_recur[stage]),
273                    glm5_tp_latent_peer: std::mem::take(&mut g5_latent[stage]),
274                    pos,
275                    max_ctx,
276                    tainted: false,
277                    last_logits_dev: None,
278                    dflash_taps: None,
279                    hc_taps: None,
280                    // Per-stage split caches start with no captured graphs: a run graph bakes the
281                    // state pointers of the cache it was captured against, and this one is new.
282                    glm5_decode_graph: None,
283                })
284            })
285            .collect();
286        Self {
287            parent,
288            fence: fence.to_vec(),
289            stages,
290            committed: false,
291        }
292    }
293
294    fn pp2_parts(&mut self) -> (&mut Cache, &mut Cache) {
295        assert_eq!(self.stages.len(), 2);
296        let (stage0, stage1) = self.stages.split_at_mut(1);
297        (
298            stage0[0]
299                .get_mut()
300                .unwrap_or_else(|poisoned| poisoned.into_inner()),
301            stage1[0]
302                .get_mut()
303                .unwrap_or_else(|poisoned| poisoned.into_inner()),
304        )
305    }
306
307    fn stages(&self) -> &[std::sync::Mutex<Cache>] {
308        &self.stages
309    }
310
311    fn commit(&mut self) {
312        self.committed = true;
313    }
314}
315
316impl Drop for PrimeCacheStages<'_> {
317    fn drop(&mut self) {
318        let n = self.parent.kv.len();
319        for i in 0..n {
320            let stage = prime_cache_stage_for_layer(&self.fence, i);
321            let source = self.stages[stage]
322                .get_mut()
323                .unwrap_or_else(|poisoned| poisoned.into_inner());
324            debug_assert!(self.parent.kv[i].is_none());
325            debug_assert!(self.parent.tp_kv[i].is_none());
326            debug_assert!(self.parent.recur[i].is_none());
327            debug_assert!(self.parent.latent[i].is_none());
328            self.parent.kv[i] = source.kv[i].take();
329            self.parent.tp_kv[i] = source.tp_kv[i].take();
330            self.parent.recur[i] = source.recur[i].take();
331            self.parent.latent[i] = source.latent[i].take();
332            self.parent.glm5_tp_recur[i] = source.glm5_tp_recur[i].take();
333            self.parent.glm5_tp_latent_peer[i] = source.glm5_tp_latent_peer[i].take();
334        }
335        self.parent.pos = self
336            .stages
337            .iter_mut()
338            .map(|stage| {
339                stage
340                    .get_mut()
341                    .unwrap_or_else(|poisoned| poisoned.into_inner())
342                    .pos
343            })
344            .min()
345            .unwrap_or(self.parent.pos);
346        if !self.committed {
347            self.parent.mark_tainted();
348        }
349    }
350}
351
352/// Fail-stop transaction marker for concat-prime paths. These paths mutate several independent
353/// caches before their final epilogue can fail; an error must make every member permanently
354/// ineligible for retry/reuse rather than replaying a queue over partially advanced state.
355struct CacheTaintGuard {
356    caches: Vec<*mut Cache>,
357    committed: bool,
358}
359
360impl CacheTaintGuard {
361    fn arm(caches: &mut [&mut Cache]) -> Self {
362        Self {
363            caches: caches
364                .iter_mut()
365                .map(|cache| *cache as *mut Cache)
366                .collect(),
367            committed: false,
368        }
369    }
370
371    fn commit(&mut self) {
372        self.committed = true;
373    }
374}
375
376impl Drop for CacheTaintGuard {
377    fn drop(&mut self) {
378        if self.committed {
379            return;
380        }
381        for cache in &self.caches {
382            // SAFETY: `arm` receives the function's unique cache references. The guard never
383            // escapes that call or dereferences them until unwind/return after active borrows end.
384            unsafe { (&mut **cache).mark_tainted() };
385        }
386    }
387}
388
389#[derive(Debug, Clone, Copy, PartialEq, Eq)]
390struct PrimePpWaveSlot {
391    wave: usize,
392    slot: usize,
393}
394
395#[derive(Debug)]
396enum PrimePpSignal {
397    Slot(PrimePpWaveSlot),
398    Error(String),
399}
400
401#[derive(Default)]
402struct PrimePpWaveCredits {
403    next_wave: usize,
404    pending: std::collections::VecDeque<PrimePpWaveSlot>,
405}
406
407impl PrimePpWaveCredits {
408    fn release_required(&self) -> Option<PrimePpWaveSlot> {
409        (self.pending.len() == 2).then(|| self.pending[0])
410    }
411
412    fn record_release(&mut self, released: PrimePpWaveSlot) -> Result<(), String> {
413        let expected =
414            self.pending.front().copied().ok_or_else(|| {
415                "prime PP received a slot release with no pending wave".to_string()
416            })?;
417        if released != expected {
418            return Err(format!(
419                "prime PP slot release {:?} does not match oldest pending {:?}",
420                released, expected
421            ));
422        }
423        self.pending.pop_front();
424        Ok(())
425    }
426
427    fn record_send(&mut self, sent: PrimePpWaveSlot) -> Result<(), String> {
428        if sent.wave != self.next_wave {
429            return Err(format!(
430                "prime PP sent wave {} while wave {} was next",
431                sent.wave, self.next_wave
432            ));
433        }
434        if sent.slot >= 2 {
435            return Err(format!(
436                "prime PP boundary returned invalid slot {}",
437                sent.slot
438            ));
439        }
440        if self.pending.iter().any(|pending| pending.slot == sent.slot) {
441            return Err(format!(
442                "prime PP reused slot {} before its exact-wave release",
443                sent.slot
444            ));
445        }
446        self.pending.push_back(sent);
447        self.next_wave += 1;
448        Ok(())
449    }
450}
451
452fn recv_prime_pp_signal(
453    receiver: &std::sync::mpsc::Receiver<PrimePpSignal>,
454    expected: PrimePpWaveSlot,
455    exact_slot: bool,
456    label: &str,
457) -> Result<PrimePpWaveSlot, String> {
458    match receiver.recv() {
459        Ok(PrimePpSignal::Error(error)) => Err(error),
460        Ok(PrimePpSignal::Slot(received))
461            if received.wave == expected.wave
462                && (!exact_slot || received.slot == expected.slot) =>
463        {
464            if received.slot >= 2 {
465                Err(format!(
466                    "{label}: wave {} carried invalid slot {}",
467                    received.wave, received.slot
468                ))
469            } else {
470                Ok(received)
471            }
472        }
473        Ok(PrimePpSignal::Slot(received)) => Err(format!(
474            "{label}: expected wave/slot {:?}, received {:?}",
475            expected, received
476        )),
477        Err(_) => Err(format!(
478            "{label}: channel closed while waiting for wave {}",
479            expected.wave
480        )),
481    }
482}
483
484fn send_prime_pp_signal(
485    sender: &std::sync::mpsc::Sender<PrimePpSignal>,
486    signal: PrimePpSignal,
487    label: &str,
488) -> Result<(), String> {
489    sender
490        .send(signal)
491        .map_err(|_| format!("{label}: channel closed"))
492}
493
494struct PrimePpWave<'a> {
495    start: usize,
496    end: usize,
497    tokens: &'a [u32],
498}
499
500struct PrimePpStageChannels {
501    incoming: Option<std::sync::mpsc::Receiver<PrimePpSignal>>,
502    release_upstream: Option<std::sync::mpsc::Sender<PrimePpSignal>>,
503    outgoing: std::sync::mpsc::Sender<PrimePpSignal>,
504    released_downstream: std::sync::mpsc::Receiver<PrimePpSignal>,
505}
506
507impl PrimePpStageChannels {
508    fn notify_failure(&self, error: &str) {
509        if let Some(upstream) = &self.release_upstream {
510            let _ = upstream.send(PrimePpSignal::Error(error.to_string()));
511        }
512        let _ = self.outgoing.send(PrimePpSignal::Error(error.to_string()));
513    }
514}
515
516/// The DSA k-pool indexer's resident state, borrowed for one `mla_attn_core` call.
517///
518/// TWO PLANES, DIFFERENT LIFETIMES. `state` is the packed `[k_norm | gate]` row per cached token
519/// (`LatentKvLayer::index_rows`); it is append-only and grows with the cache. `pool_keys` is the
520/// collapsed key per COMPLETE pool of `pool` such rows, and it is the residency win: a pool's key
521/// is final the moment its last row lands, so pools `[0, *ready)` are never recomputed and each
522/// call builds only the pools its own tokens completed. `ready` is written back through the
523/// borrow, so the caller must persist it alongside the buffers.
524///
525/// `pool_keys` is `Option` because its size needs the indexer's `pool`, which the state plan does
526/// not carry — `mla_kpool_indices` allocates it on first use and leaves it resident thereafter.
527/// A caller that hands over a fresh `None` every call (the stateless arm) gets the old
528/// rebuild-everything behaviour, which is exactly right when the state itself is per-call.
529/// The MLA core's PRE/POST handoff buffers, one set per session, reused by every MLA layer
530/// (lane/glm5-mla-capture-20260904, door `MEMRA_MLA_SEG_WS`). WHY: a captured PRE graph and a
531/// captured POST graph must hand each other buffers at STABLE addresses; allocations inside ONE
532/// capture are fine (they become graph alloc nodes, which is how the KDA runs already capture),
533/// but two graphs cannot hand each other a fresh allocation. Every MLA layer of a glm5 model
534/// shares one `MlaGeom`, so one set serves all eleven, sequentially, exactly as the eager walk's
535/// fresh buffers do. BYTE-IDENTICAL by construction: the same kernels write the same values in
536/// the same order to a different address, and every consumer reads the value, never the address
537/// (door W's contract, `MEMRA_VERIFY_WS`).
538/// Engagement counter for `MEMRA_MLA_SEG_WS`; gates take a delta.
539pub static MLA_SEG_WS_DISPATCHES: std::sync::atomic::AtomicU64 =
540    std::sync::atomic::AtomicU64::new(0);
541
542/// Snapshot of [`MLA_SEG_WS_DISPATCHES`].
543pub fn mla_seg_ws_dispatches() -> u64 {
544    MLA_SEG_WS_DISPATCHES.load(std::sync::atomic::Ordering::Relaxed)
545}
546
547pub(crate) struct MlaSegWs {
548    pub q_nope: CudaSlice<f32>,
549    pub q_pe: CudaSlice<f32>,
550    pub q_an: CudaSlice<f32>,
551    pub c_kv_n: CudaSlice<f32>,
552    pub k_pe: CudaSlice<f32>,
553    /// `(nh, dn, dr, r, q_lora)` this set was sized for; another geometry refuses by name.
554    pub sig: (usize, usize, usize, usize, usize),
555}
556
557impl MlaSegWs {
558    /// Size one set for `t = 1` decode at this layer's geometry.
559    pub(crate) fn new(
560        e: &Engine,
561        nh: usize,
562        dn: usize,
563        dr: usize,
564        r: usize,
565        q_lora: usize,
566    ) -> Result<Self, Box<dyn std::error::Error>> {
567        Ok(Self {
568            q_nope: e.uninit(nh * dn)?,
569            q_pe: e.uninit((nh * dr).max(1))?,
570            q_an: e.uninit(q_lora)?,
571            c_kv_n: e.uninit(r)?,
572            k_pe: e.uninit(dr.max(1))?,
573            sig: (nh, dn, dr, r, q_lora),
574        })
575    }
576}
577
578/// The three planes segment MID reads, from either the owned `MlaPreOut` or the pooled
579/// [`MlaSegWs`].
580pub(crate) struct MlaMidIn<'a> {
581    pub q_an: &'a CudaSlice<f32>,
582    pub c_kv_n: &'a CudaSlice<f32>,
583    pub k_pe: &'a CudaSlice<f32>,
584}
585
586/// The MLA core's PRE segment outputs (lane/glm5-mla-segments-20260904): everything the
587/// append and the attention read from the projections.
588pub(crate) struct MlaPreOut {
589    pub q_nope: CudaSlice<f32>,
590    pub q_pe: CudaSlice<f32>,
591    pub q_an: CudaSlice<f32>,
592    pub c_kv_n: CudaSlice<f32>,
593    pub k_pe: CudaSlice<f32>,
594}
595
596pub struct IndexerPlanes<'a> {
597    pub state: &'a mut CudaSlice<f32>,
598    pub pool_keys: &'a mut Option<CudaSlice<f32>>,
599    pub ready: &'a mut usize,
600    /// PHYSICAL rows of `state` when it is a TAIL RING; 0 when the plane is flat (one row per
601    /// cached token, absolute addressing). `mla_kpool_indices` rounds this DOWN to a multiple of
602    /// the indexer's `pool` — the state plan does not carry `pool`, so the allocator cannot — and
603    /// proves the liveness bound against the rounded value before it appends.
604    pub state_ring_rows: usize,
605    /// Token capacity of the session, which sizes `pool_keys`. It is NOT derivable from
606    /// `state.len()` once `state` is a ring: the ring holds one call's tail, the pool-key plane
607    /// holds the whole context collapsed `pool`-to-one.
608    pub capacity_tokens: usize,
609}
610
611/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
612pub(crate) struct AttnPre {
613    pub q: cudarc::driver::CudaSlice<f32>,
614    pub k: cudarc::driver::CudaSlice<f32>,
615    pub v: cudarc::driver::CudaSlice<f32>,
616    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
617}
618
619/// task #18: one sequence's GDN prep outputs (the scan inputs).
620pub(crate) struct GdnPrep {
621    pub hk: usize,
622    pub q_l2: cudarc::driver::CudaSlice<f32>,
623    pub k_l2: cudarc::driver::CudaSlice<f32>,
624    pub v_g: cudarc::driver::CudaSlice<f32>,
625    pub beta: cudarc::driver::CudaSlice<f32>,
626    pub g_log: cudarc::driver::CudaSlice<f32>,
627    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
628    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
629}
630
631/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
632pub(crate) struct VerifyStreamScratch {
633    pub pos_d: CudaSlice<i32>,
634    pub row_ctrs: Vec<CudaSlice<i32>>,
635}
636use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
637
638struct MoeInputTraceWriter {
639    dir: std::path::PathBuf,
640    index: std::fs::File,
641    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
642}
643
644static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
645    std::sync::OnceLock::new();
646
647/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
648/// per-expert launch chain). See `moe_gdec_token`.
649fn gdec_enabled() -> bool {
650    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
651    *E.get_or_init(|| {
652        std::env::var("MEMRA_MOE_GDEC")
653            .map(|v| v != "0")
654            .unwrap_or(true)
655    })
656}
657
658/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
659/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
660/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
661/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
662/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
663/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
664/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
665/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
666/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
667/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
668fn moe_slab_enabled() -> bool {
669    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
670}
671
672/// `MEMRA_MOE_FUSED_EPI` — the glm5_next fused MoE epilogue (sigmoid-routed, PRE-clamped SwiGLU,
673/// per-expert macro fold) collapsed into one launch pair per token-layer.
674///
675/// DEFAULT OFF, deliberately (docs/FLAGS.md carries the row and the reasons). The arm is proven
676/// EXACT against `memra_reference` by `tests/glm5_moe_epilogue_gpu.rs`, but it has no throughput
677/// receipt: the rig is correctness-only by law and the 190.7 GB artifact has never been on it, so
678/// the launch-count claim is arithmetic from source and nothing has been measured on serving
679/// hardware. Unmeasured behavior does not default ON.
680///
681/// Read PER CALL, not latched in a `OnceLock`: the acceptance gate flips both arms inside one
682/// test process (the interleave unit is a model load, not a boot), and a latched flag would make
683/// the second arm silently a copy of the first.
684fn moe_fused_epi_enabled() -> bool {
685    std::env::var("MEMRA_MOE_FUSED_EPI")
686        .map(|v| v != "0")
687        .unwrap_or(false)
688}
689
690/// `MEMRA_HC_DECODE_WS` — the persistent hc-glue decode workspace (lane/glm5-decode-diet
691/// lever 2): the T=1 hc decode walk lands its glue transients (mixes, gates, comb, collapse
692/// y, both norm scratches, the per-site post output) in one per-engine `HyperDecodeWs`
693/// instead of ~12 fresh `cuMemAllocAsync`+free pairs per layer per token (the launch-diet
694/// census's 2,358-calls/token class). Same kernels, same call order, same operand bytes —
695/// byte identity ON/OFF gated by `tests/hc_decode_ws_gpu.rs`.
696///
697/// DEFAULT OFF, deliberately (docs/FLAGS.md row): the alloc-call reduction is proven on the
698/// rig by counter receipt, but the ms/token value is arithmetic against the box's measured
699/// launch/alloc constants — nothing has been measured on serving hardware yet. Unmeasured
700/// behavior does not default ON.
701///
702/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent).
703/// memra#131 cell 6: the graph self-check runs a second, workspace-free eager reference next to
704/// the workspace one; this flag forces the allocating walk for that reference only.
705pub(crate) static HC_WS_FORCE_PLAIN: std::sync::atomic::AtomicBool =
706    std::sync::atomic::AtomicBool::new(false);
707
708fn hyper_decode_ws_on() -> bool {
709    !HC_WS_FORCE_PLAIN.load(std::sync::atomic::Ordering::Relaxed)
710        && std::env::var("MEMRA_HC_DECODE_WS").as_deref() == Ok("1")
711}
712
713/// Engagement counter for the workspace walk — the receipt the gate and any box A/B arm
714/// must show (engagement lines are receipts, never inferred).
715pub static HC_DECODE_WS_DISPATCHES: std::sync::atomic::AtomicU64 =
716    std::sync::atomic::AtomicU64::new(0);
717
718/// Engagement counter for door `MEMRA_GLM5_Q8_FUSE` (lane/b200-q8-fuse-20260902) — every
719/// call into `rms_norm_zq8_f32` from the mHC T=1 decode walk increments this, and the FIRST
720/// increment prints `[glm5-q8-fuse] engaged` once per boot. Needed so a box A/B or nsys
721/// census can prove the fused kernel actually ran rather than inferring it from a green
722/// diff (wiring-assertions law).
723pub static GLM5_Q8_FUSE_DISPATCHES: std::sync::atomic::AtomicU64 =
724    std::sync::atomic::AtomicU64::new(0);
725
726/// `MEMRA_MLA_TC_PREFILL` — the glm5_next tensor-core MLA prefill chain
727/// (lane/glm5-mla-tc-prefill, 2026-08-30): at prefill widths the three per-position f32
728/// kernels the launch-diet census named (`memra_mla_attn_gathered_kernel` 139 ms +
729/// `memra_mla_absorb_q_kernel` 44.5 ms + `memra_mla_decompress_v_kernel` 43.6 ms per
730/// layer-chunk, 75.8% of a 98%-GPU-busy cold prime) are replaced by two strided-batched
731/// bf16 tensor-core GEMMs (absorb / decompress, one launch each) and one gathered
732/// flash-attention MMA kernel (`fa_mla_gathered_bf16`). Selection, the latent cache, the
733/// q/kv projections, and decode are UNTOUCHED.
734///
735/// DEFAULT ON (owner acceptance 2026-08-30, "why not? i dont see why not", on the two-box
736/// A/B receipts): interleaved x5 fresh boots per arm on BOTH the Server-Edition and
737/// Workstation-Edition 4-card boxes, zero violations, zero argmax flips across 20 boots,
738/// TTFD -62%..-69% (7.45->2.83 s @4.6k / 6.58->2.51 s), prefill 619-724 -> 1629-2255 tok/s,
739/// vendor-default sampled twin -66/-67%, decode untouched, engagement receipted in every ON
740/// boot with no cuBLASLt declines (docs/FLAGS.md row carries the pointers). The numeric
741/// config remains band-gated (bf16 operands, f32 accumulate — the fa_prefill/MEMRA_PP_BF16
742/// class, `tests/mla_tc_prefill_gpu.rs`, never bit). `MEMRA_MLA_TC_PREFILL=0` is the
743/// rollback seam.
744///
745/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent): the gate
746/// flips both arms inside one test process, and a latched flag would make the second arm
747/// silently a copy of the first.
748fn mla_tc_prefill_enabled() -> bool {
749    std::env::var("MEMRA_MLA_TC_PREFILL")
750        .map(|v| v != "0")
751        .unwrap_or(true)
752}
753
754/// `MEMRA_B200_PRIME_V2` — the B200 (sm_100a) mHC PRIME schedule door
755/// (lane/b200-prefill-roofline-20260902, **default OFF**). It arms two changes to the
756/// HyperConnections prime and nothing else; decode, spec, batch and every non-hyper trunk are
757/// untouched by construction (the only two call sites are `hyper_prime_ranges` and
758/// `prime_cache_hyper`'s pp branch).
759///
760/// WHY, measured. `prime_chunk_tokens` applies the `PRIME_PIPE_MICROBATCHES = 8` geometry
761/// whenever a PP-2 stage fence exists, so a 4096-token glm5 prime splits into EIGHT calls. That
762/// geometry belongs to [`HybridModel::prime_cache_pp2_pipelined`], the SERIAL trunk's
763/// microbatched PP-2 prime — the hyper walk has never called it, so glm5 pays the split and
764/// collects none of the overlap. The split is not free: every one of the 288 experts is hit at
765/// any width past a few hundred tokens, so a MoE layer reads its WHOLE 4.08 GB slab once per
766/// CALL. Eight calls read 1,369.8 GB where one reads 171.2 GB; the HBM floor rises 21.4 ->
767/// 171.2 ms, the 4k prime crosses from compute-bound to memory-bound (854 -> 107 FLOP/B against
768/// a B200 ridge of 275), and pairs-per-expert falls 113.8 -> 14.2, under the grouped GEMM's
769/// `MEMRA_F16G_SK_CROSS` 64-row tile crossover. The 41.9k prime, whose chunk clamps at
770/// `PRIME_CHUNK_MAX_TOKENS` instead, runs 3.3x more efficient per FLOP on the same binary
771/// (87.0 vs 26.6 TFLOP/s) — the same fact from the other side. Arithmetic:
772/// `research/b200-prefill-roofline-20260902/roofline.py`.
773///
774/// ARM 1 (schedule): `hyper_prime_ranges` takes the natural chunk
775/// (`PRIME_CHUNK_MAX_TOKENS`, the same clamp the 41.9k prime already gets) instead of the
776/// microbatch geometry. An explicit `MEMRA_PRIME_CHUNK` stays authoritative and turns this arm
777/// off for that call, so the operator override is never overridden.
778/// ARM 2 (pipeline): with two or more chunks the ppN prime runs
779/// [`HybridModel::prime_cache_hyper_pp2_pipelined`] — stage 0 of chunk k+1 on device 0 while
780/// stage 1 of chunk k runs on device 1, on two scoped host threads over disjoint per-stage
781/// caches. Two threads are REQUIRED, not a style choice: the grouped MoE prefill's sigmoid
782/// host-oracle drains the stage stream once per layer, so two CUDA streams on one host thread
783/// serialize however the calls are ordered (the same finding `prime_cache_pp2_pipelined`
784/// records for the serial trunk).
785///
786/// NUMERIC CLASS. Arm 1 is NOT bit-identical and no chunk-size change on this trunk ever was:
787/// `Engine::linear` — the cuBLASLt f32 mixes GEMM in `hyper::pre` — is not m-invariant, the
788/// arms diverge at ROW 0 where no cross-token state can reach, and `hyper_prime_ranges`' own
789/// header documents the near-tie. The bar is the calibrated band
790/// `tests/glm5_chunked_prime_gpu.rs` already holds the chunked prime to (relative maxdiff
791/// <= 2e-5, five orders below the 1.813e0 signature of a real chunk-invariance defect) plus
792/// argmax equality. Arm 2 IS bit-identical by construction: same ranges, same per-chunk
793/// program, same operand bytes — only which host thread issues a stage and when changes.
794///
795/// Read PER CALL (the `MEMRA_MLA_TC_PREFILL` rollback-seam precedent): `glm5-prime-v2-gate`
796/// flips both arms inside one process, and a latched flag would make the second arm a copy of
797/// the first.
798/// Where `hyper_range_prime` finds the DFlash2 hc tap sink for the walk it is running.
799///
800/// The serial walks read it from the cache, exactly as before. The PIPELINED prime cannot:
801/// [`PrimeCacheStages`] gives each stage a cache shell with `hc_taps: None`, and the two stage
802/// threads are on DIFFERENT CHUNKS at the same moment, so the sink's single `base` field cannot
803/// describe both. So the pipelined walk shares ONE sink and passes each stage its own base.
804///
805/// SHARING RATHER THAN SPLITTING, because splitting is unaffordable: a prime sink spans the
806/// whole prompt (`HcTapSink::new(taps, n_embd, plen)`), which at depth is gigabytes of host
807/// rows, and a per-stage copy would double it.
808///
809/// SHARING IS SAFE BY CONSTRUCTION, and this is the argument the bit-identity claim rests on.
810/// A tap write lands at `(base - origin + r) * n_taps * hidden + slot * hidden` for r in 0..t.
811/// Two things make the two stages' writes disjoint, and either alone would be enough:
812///
813///   * DIFFERENT SLOT COLUMNS. A tapped layer belongs to exactly one stage — the same fact
814///     `glm5_tap_drain` already relies on to pick a slot's owning engine — so stage 0 writes
815///     only the columns of layers below the fence and stage 1 only those above it.
816///   * DIFFERENT ROWS. The stages run different chunks, so their `base` windows do not overlap.
817///
818/// The mutex is therefore not protecting against a data race on the values; it is what makes
819/// the shared `&mut` legal, and it is held only for the row memcpy — the contraction and the
820/// device-to-host readback happen outside it.
821pub(crate) enum HcTapArm<'a, 's: 'a> {
822    /// Read the sink from the cache. Every serial walk; byte-identical to the pre-lane path.
823    FromCache,
824    /// The pipelined mHC prime: one shared sink, this stage's chunk base.
825    Shared(&'a std::sync::Mutex<&'s mut crate::cache::HcTapSink>, usize),
826}
827
828/// One named line per process when the pipelined mHC prime arm declines. A door that is armed
829/// and inert must say so: the grouped MoE prefill's own `decline_once` exists because this
830/// exact shape (flag announced ON, arm returning early on every layer) cost a 24k prime its
831/// speed while the boot log said the door was on.
832fn hyper_pipe_decline_once(reason: &str) {
833    static DECLINED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
834    if !DECLINED.swap(true, std::sync::atomic::Ordering::Relaxed) {
835        eprintln!(
836            "[hyper-prime-pipe] DECLINED: {reason} -> the serial per-chunk stage walk serves \
837             this prime (MEMRA_B200_PRIME_V2 arm 1 is unaffected; logged once per process)"
838        );
839    }
840}
841
842pub fn b200_prime_v2_on() -> bool {
843    std::env::var("MEMRA_B200_PRIME_V2").as_deref() == Ok("1")
844}
845
846/// Chunks primed through the pipelined mHC PP-2 walk. The engagement receipt arm 2's
847/// non-vacuity assertion reads: bit identity between two walks that both ran SERIALLY is
848/// vacuous, so the gate requires this to advance.
849/// Prime calls whose schedule was decided by `MEMRA_B200_PRIME_V2` arm 1 (the natural chunk).
850/// The engagement counter the gate reads: an armed door that changed no schedule is the exact
851/// failure this lane was written to catch, so "the flag was set" is never the receipt.
852pub static HYPER_PRIME_NATURAL_SCHEDULES: std::sync::atomic::AtomicU64 =
853    std::sync::atomic::AtomicU64::new(0);
854
855pub static HYPER_PRIME_PIPELINED_CHUNKS: std::sync::atomic::AtomicU64 =
856    std::sync::atomic::AtomicU64::new(0);
857
858/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
859/// default flip. `=0` selects the established path, while any other explicit value enables the
860/// grouped research arm for the current call.
861fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
862    std::env::var("MEMRA_MOE_GROUPED")
863        .map(|value| value != "0")
864        .unwrap_or(false)
865}
866
867/// `MEMRA_MOE_GROUPED_PREFILL`: the glm5_next expert-grouped MoE PREFILL arm, token-sort by
868/// expert (host CSR, the `moe_align_block_size` shape), then ONE grouped tensor-core GEMM per
869/// projection per layer-chunk over the resident NVFP4 bank, with the sigmoid `noaux_tc` routing,
870/// the PRE-clamped SwiGLU epilogue and the per-expert `weight_scale_2` macro fold the fused
871/// epilogue lane qualified for this family.
872///
873/// DEFAULT ON since 2026-08-29 (owner acceptance; `=0` is the rollback seam). The flip carries
874/// its receipts, per the flag-default law: reference-band + routing-exactness gate green
875/// (`tests/glm5_moe_grouped_prefill_gpu.rs`; grouped GEMM is measured non-bit-stable, so byte
876/// identity is not the honest bar; routing sel/w stay bit-identical by construction, the same
877/// `moe_router_logits` + `moe_route_sigmoid_cfg` invocation as the sequential arm), plus the
878/// interleaved x5 box A/B on the serving card class: TTFD 54.2 -> 7.5 s / 65.5 -> 8.9 /
879/// 75.9 -> 10.3 at 4.6/5.5/6.5k-token real prompts (85 -> 616-639 tok/s prefill, decode
880/// unchanged, sampled vendor-default twin green, engagement 42/42). The one greedy first-token
881/// flip (B5550) sits at a position the 8-draw vendor-default census measured as SOFT in both
882/// arms (the OFF arm itself draws the ON arm's token there) and was accepted by the OWNER on
883/// 2026-08-29, the MEMRA_BF16_MMV acceptance class. Receipts:
884/// `research/glm53-flash-bringup-20260827/moe-grouped-prefill-receipts/` (`box-ab-20260829/`).
885///
886/// Read PER CALL, not latched: the acceptance gate flips both arms inside one test process.
887fn moe_grouped_prefill_enabled() -> bool {
888    std::env::var("MEMRA_MOE_GROUPED_PREFILL")
889        .map(|v| v != "0")
890        .unwrap_or(true)
891}
892
893/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
894/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
895fn moe_prefetch_enabled() -> bool {
896    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
897    *E.get_or_init(|| {
898        std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
899            || crate::spill_pread::worker_enabled()
900    })
901}
902
903/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
904/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
905/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
906/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
907fn moe_page_prefetch_window() -> usize {
908    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
909    *W.get_or_init(|| {
910        page_prefetch_window_from_values(
911            std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
912            std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
913                .ok()
914                .as_deref(),
915        )
916    })
917}
918
919fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
920    if !enabled {
921        return 0;
922    }
923    raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
924}
925
926/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
927/// full window; each later position adds one expert at the far edge. Thus widening the window does
928/// not repeatedly issue `MADV_WILLNEED` for the same range.
929fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
930    if window == 0 || position >= len {
931        return len..len;
932    }
933    let (start, count) = if position == 0 {
934        (1, window)
935    } else {
936        (position.saturating_add(window), 1)
937    };
938    let start = start.min(len);
939    start..start.saturating_add(count).min(len)
940}
941
942/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
943/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
944fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
945    let position = current.map_or(0, |position| position.saturating_add(1));
946    (position < order_len).then_some(position)
947}
948
949/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
950/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
951/// window. Position zero primes the current expert too: its three independent reads can run in
952/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
953fn worker_prefetch_window() -> usize {
954    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
955    *WINDOW.get_or_init(|| {
956        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
957        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
958            .ok()
959            .and_then(|value| value.parse::<usize>().ok())
960            .unwrap_or(automatic.max(1))
961    })
962}
963
964/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
965/// this includes the current expert when the window is seeded so all three current projections
966/// enter the CPU pool together.
967fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
968    if window == 0 || position >= len {
969        return len..len;
970    }
971    let (start, count) = if position == 0 {
972        (0, window)
973    } else {
974        (position.saturating_add(window).saturating_sub(1), 1)
975    };
976    let start = start.min(len);
977    start..start.saturating_add(count).min(len)
978}
979
980/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
981/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
982/// expert weight pointers come from the per-layer device table. Requires the fused router (the
983/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
984fn moe_dev_enabled() -> bool {
985    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
986    *E.get_or_init(|| {
987        std::env::var("MEMRA_MOE_DEV")
988            .map(|v| v != "0")
989            .unwrap_or(true)
990            && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
991    })
992}
993
994/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
995/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
996/// Where the verify-rows MoE pair's routed selection lives for one layer-call
997/// (lane/glm5-moe-loc door D). `Host` is the shipped arm: the router's pinned readback gave the
998/// host `sel`/`w`, and the host builds the pointer/scale tables. `Dev` is door D's arm: the
999/// router's own device `sel_idx`/`sel_w` are still live, so the tables are built where they are
1000/// and the readback (2 DtoH + a full `cuStreamSynchronize` per MoE layer-call) never happens.
1001/// ONE launch path consumes both — only the table build differs, term-for-term identically.
1002enum VrowsSel<'a> {
1003    Host(&'a [u32], &'a [f32]),
1004    Dev(&'a CudaSlice<i32>, &'a CudaSlice<f32>),
1005}
1006
1007pub(crate) fn sigmoid_router_enabled() -> bool {
1008    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1009    *E.get_or_init(|| {
1010        std::env::var("MEMRA_SIG_ROUTER")
1011            .map(|v| v != "0")
1012            .unwrap_or(true)
1013    })
1014}
1015
1016/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
1017/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
1018/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
1019/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
1020fn moe_q8_enabled() -> bool {
1021    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1022    *E.get_or_init(|| {
1023        std::env::var("MEMRA_MOE_Q8")
1024            .map(|v| v != "0")
1025            .unwrap_or(true)
1026    })
1027}
1028
1029/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
1030/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
1031fn expert_dp4a_supported(qt: i32) -> bool {
1032    qt == crate::QT_Q4_0
1033        || qt == crate::QT_IQ3_S
1034        || qt == crate::QT_IQ4_XS
1035        || qt == crate::QT_Q3_K
1036        || qt == crate::QT_Q4_K
1037        || qt == crate::QT_Q6_K
1038}
1039
1040fn q8_expert_supported(qt: i32) -> bool {
1041    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
1042    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
1043    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
1044    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
1045    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
1046    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1047    let kq = *KQ.get_or_init(|| {
1048        std::env::var("MEMRA_MOE_Q8_KQ")
1049            .map(|v| v != "0")
1050            .unwrap_or(true)
1051    });
1052    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
1053    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
1054    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
1055    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
1056    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
1057    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
1058    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
1059        .map(|v| v != "0")
1060        .unwrap_or(true);
1061    qt == crate::QT_IQ3_S
1062        || qt == crate::QT_IQ4_XS
1063        || (nvfp4_q8 && qt == crate::QT_NVFP4)
1064        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
1065}
1066
1067/// ModelOpt `W4A16_NVFP4` is weight-only: feeding its expert weights a q8_1 activation changes
1068/// the declared numeric program to W4A8. Keep the established q8 path for every other artifact,
1069/// but force Hy3 W4A16 experts through the BF16-activation `qmatvec_view` oracle.
1070fn q8_expert_supported_for_model(cfg: &ModelConfig, qt: i32) -> bool {
1071    let weight_only_nvfp4 = cfg.hy3.as_ref().is_some_and(|hy3| hy3.weight_only_nvfp4);
1072    q8_expert_supported(qt) && !(weight_only_nvfp4 && qt == crate::QT_NVFP4)
1073}
1074
1075fn moe_q8_enabled_for_model(cfg: &ModelConfig, m: &MoeWeights) -> bool {
1076    m.has_uniform_expert_layout()
1077        && moe_q8_enabled()
1078        && q8_expert_supported_for_model(cfg, m.gate_exps.qtype)
1079        && q8_expert_supported_for_model(cfg, m.up_exps.qtype)
1080        && q8_expert_supported_for_model(cfg, m.down_exps.qtype)
1081}
1082
1083#[cfg(test)]
1084mod w4a16_dispatch_tests {
1085    use super::q8_expert_supported_for_model;
1086    use memra_gguf::config::{HfConfig, ModelConfig};
1087
1088    #[test]
1089    fn hy3_w4a16_never_admits_q8_activations() {
1090        let hf = HfConfig::parse(
1091            r#"{"model_type":"hy_v3","num_hidden_layers":2,"hidden_size":8,
1092            "num_attention_heads":2,"intermediate_size":16,"vocab_size":32,
1093            "max_position_embeddings":32,
1094            "quantization_config":{"quant_method":"modelopt","quant_algo":"W4A16_NVFP4"}}"#,
1095        );
1096        let cfg = ModelConfig::from_hf(&hf);
1097        assert!(!q8_expert_supported_for_model(&cfg, crate::QT_NVFP4));
1098        assert!(q8_expert_supported_for_model(&cfg, crate::QT_IQ4_XS));
1099    }
1100}
1101
1102/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
1103/// k-quant tensors must fall to the _em dot path instead.
1104fn q8_expert_dec_supported(qt: i32) -> bool {
1105    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
1106}
1107
1108/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
1109/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
1110/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
1111/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
1112/// q35 layers, which is why that cell measured FLAT.
1113fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
1114    match qt {
1115        crate::QT_Q4_0 => in_f.is_multiple_of(32),
1116        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
1117            in_f.is_multiple_of(256)
1118        }
1119        // NVFP4 (block 64) added lane/moebatch-q35moe 2026-08-21: the ornith15 expert bank is
1120        // uniform NVFP4, which passed the pairs q8 gate but missed BOTH batched doors
1121        // (use_mma's dec set and this table), so 14.7k-token prefill rode the per-pair _em
1122        // fallback — 88.6% of the prime wall (prime-anatomy receipt).
1123        crate::QT_NVFP4 | crate::QT_NVFP4_V2 => in_f.is_multiple_of(64),
1124        _ => false,
1125    }
1126}
1127
1128/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
1129/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
1130fn moe_prewarm_enabled() -> bool {
1131    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1132    *E.get_or_init(|| {
1133        std::env::var("MEMRA_MOE_PREWARM")
1134            .map(|v| v != "0")
1135            .unwrap_or(true)
1136    })
1137}
1138
1139/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
1140/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
1141/// can then vote for and exercise those experts on GPU before the cache is frozen.
1142fn cpu_expert_profile_admit_enabled() -> bool {
1143    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1144    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
1145}
1146
1147/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
1148/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
1149/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
1150pub const PRIME_MIN_T: usize = 16;
1151
1152/// CUDA grid dimensions y and z cap at 65,535 (an architecture constant on every compute
1153/// capability we target; only grid.x is 2^31-1). Several prime-path kernels launch with the
1154/// call's token count in grid.y — the fused GDN conv (`ssm_conv1d_gdn_state_f32`,
1155/// lib.rs `ssm_conv1d_gdn_state_pad`), the dp4a matvec family taken at out_f < 128
1156/// (`qmatvec*_fast`/`qmatvec_dp4a_named`, which is where the GDN ssm_beta/ssm_alpha
1157/// projections land), and the MoE `router_gemv` — so ONE prime call must never carry more
1158/// tokens than this. A monolithic prime above it is a guaranteed
1159/// `DriverError(CUDA_ERROR_INVALID_VALUE)` at launch: measured on ornith-1.5 serving with
1160/// MEMRA_PRIME_CHUNK=0 (cold 64,984 PASS / 65,643 FAIL, darklanes
1161/// research/ornith-move-20260829 F2; re-hit on prod by the 2026-09-01 stress campaign at
1162/// 66,045/79,717/82,440), and the same wall was hit and chunk-walked away by the glm5 1M
1163/// lane's ppN prime.
1164pub const CUDA_GRID_YZ_MAX: usize = 65_535;
1165
1166/// Widest prime range that stays launch-legal through the ring-off tail fold of
1167/// `fixed_prime_chunk_ranges_for_ring`: a trailing remainder shorter than `PRIME_MIN_T`
1168/// folds INTO the previous range, widening it by up to `PRIME_MIN_T - 1` tokens, so the
1169/// cap keeps `chunk + PRIME_MIN_T - 1 <= CUDA_GRID_YZ_MAX`. With this value every
1170/// t <= 65,535 still schedules as the identical single monolithic range (t <= chunk, or
1171/// the fold collapses the split), so behavior below the CUDA wall is byte-for-byte
1172/// unchanged — only prompts that today CANNOT launch get chunked.
1173pub const PRIME_CHUNK_LAUNCH_CAP: usize = CUDA_GRID_YZ_MAX - (PRIME_MIN_T - 1);
1174
1175/// The explicit-`MEMRA_PRIME_CHUNK` chunk width, ring-aware. Extracted pure for tests.
1176/// Ring ON keeps the historical clamp to `PRIME_CHUNK_MAX_TOKENS`. Ring OFF preserves the
1177/// operator's value except at the CUDA launch wall: `0` ("monolithic") now means
1178/// "monolithic up to `PRIME_CHUNK_LAUNCH_CAP`", and any larger explicit value is capped
1179/// there too — an uncapped value above the wall never produced output, only
1180/// CUDA_ERROR_INVALID_VALUE (see `CUDA_GRID_YZ_MAX`).
1181fn explicit_prime_chunk(parsed: usize, ring_on: bool) -> usize {
1182    if ring_on {
1183        if parsed == 0 {
1184            crate::cache::PRIME_CHUNK_MAX_TOKENS
1185        } else {
1186            parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
1187        }
1188    } else if parsed == 0 {
1189        PRIME_CHUNK_LAUNCH_CAP
1190    } else {
1191        parsed.min(PRIME_CHUNK_LAUNCH_CAP)
1192    }
1193}
1194
1195/// MEMRA_STEP_GEMM_PRIME_SUFFIX: does a CONTINUATION prime (`cache.pos > 0` — a rewound
1196/// session's suffix, or a prompt remainder split across scheduler ticks) ride the batched
1197/// GEMM prime, like a fresh prompt does?
1198///
1199/// DEFAULT ON since 2026-08-29, by decision, under the flip bar the OFF-era FLAGS row
1200/// wrote down (never byte identity — a prime-decomposition m-dependence that EVERY
1201/// measured prime path shares, walk included, bars that gate for all of them):
1202///  1. vendor-default sampled rows: the blind, rubric-pre-registered 8-turn quality A/B
1203///     (research/step37-sampled-quality-20260828, 72/72 valid rows, engagement receipts
1204///     per row) — WARM-GEMM sits inside COLD's own self-spread at t4 and t8 (t8 carried
1205///     at n=16; the round-1 walk-over-gemm signal collapsed at p~0.91).
1206///  2. the 8-turn cache-on twin: warm TTFT 0.58 s (door) vs 7.15 s (walk) on the real
1207///     warm serving shape, zero faults.
1208///  3. the batched prime's own standard: acceptance 0.80-0.86 across all arms with the
1209///     door arm highest at t8; interleaved arms; zero ILLEGAL/#87/panics in 19 boots.
1210///
1211/// Precondition shipped first: the SWA-ring checkpoint restore fix (c9a617ca99) — real
1212/// session reuse crosses the grow path before any door question matters.
1213/// Why it is worth it, measured: the walk continuation costs 5.5978 ms/suffix-token
1214/// against this path's 0.99 ms/token (five-point sweep, R^2 0.9976), a 7.97x suffix
1215/// slope collapse.
1216///
1217/// The `seq_end` fix beneath is NOT gated on this door — it is unconditional, because
1218/// the chunk-local `seq_end` it replaced is wrong for a fresh prompt of 4096+k tokens
1219/// (k in [PRIME_MIN_T, 512)) with no continuation anywhere in sight.
1220///
1221/// `=0` is the kill switch (continuations back on the walk, fresh primes keep the fast
1222/// path); `=1` forces; `MEMRA_STEP_GEMM_PRIME=0` remains the whole-path seam. Read per
1223/// call, not cached — probes flip it in process.
1224fn step_gemm_prime_suffix_on() -> bool {
1225    std::env::var("MEMRA_STEP_GEMM_PRIME_SUFFIX").as_deref() != Ok("0")
1226}
1227
1228/// Widest tick the MoE DEV per-token program serves (lane/orndecode-20260822). PRIME_MIN_T
1229/// doubled as the dev-arm's upper bound on the assumption that t==16 only ever meant real
1230/// prefill; the exact-16 decode tier broke that assumption — at B=16 the MoE stage crossed
1231/// onto the t>=MMA_T grouped/kq GEMM program (m_e ~1.6 rows/expert: 52.6% of the tick at
1232/// ~104 us/launch) or the `_em` per-pair fallback (67.7 us), both catastrophically slower
1233/// than the dev q8 kernels that serve B<=8 (8.8 us gate_up covering a token's whole expert
1234/// set). Decode widths 2..=16 now ride dev; the grouped/pairs prefill programs start at 17.
1235/// gate2/gate3 byte batteries at B=12/16 are the qualification (bit-checked vs isolated).
1236const MOE_DEV_MAX_T: usize = 16;
1237const PRIME_PIPE_MICROBATCHES: usize = 8;
1238const PRIME_PIPE_MIN_CHUNK: usize = 128;
1239const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
1240const PRIME_PIPE_LINEAR_WORK: usize = 8;
1241
1242fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
1243    crate::pp::prime_pp_on()
1244        && !crate::pp::pp2_streams_off()
1245        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
1246}
1247
1248fn prime_ppn_wave_auto_geometry(n_layers: usize) -> bool {
1249    crate::pp::prime_pp_on()
1250        && !crate::pp::pp2_streams_off()
1251        && crate::pp::pp_wave_on() == Ok(true)
1252        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| matches!(cuts.len(), 4 | 5))
1253}
1254
1255fn prime_pipeline_auto_geometry(n_layers: usize) -> bool {
1256    prime_pp2_auto_geometry(n_layers) || prime_ppn_wave_auto_geometry(n_layers)
1257}
1258
1259/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative up to the
1260/// CUDA launch wall (`PRIME_CHUNK_LAUNCH_CAP`; 0 = monolithic up to that wall — see
1261/// `explicit_prime_chunk`).
1262/// Pipelined PP primes use the measured PP-2 geometry: up to eight microchunks, never below
1263/// 128 tokens, while the legacy 4096-token cap remains the long-context bound. PP-3/4 inherit
1264/// only the geometry when their separate MEMRA_PP_WAVE door is explicitly open.
1265pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
1266    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
1267        let parsed = value
1268            .parse::<usize>()
1269            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
1270        return explicit_prime_chunk(parsed, crate::cache::swa_ring_on());
1271    }
1272    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
1273    if prime_pipeline_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
1274        chunk.min(
1275            t.div_ceil(PRIME_PIPE_MICROBATCHES)
1276                .max(PRIME_PIPE_MIN_CHUNK),
1277        )
1278    } else {
1279        chunk
1280    }
1281}
1282
1283fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
1284    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
1285}
1286
1287fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
1288    if chunk == 0 || t <= chunk {
1289        return vec![(0, t)];
1290    }
1291    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
1292    let mut start = 0usize;
1293    while start < t {
1294        let mut end = (start + chunk).min(t);
1295        if t - end > 0 && t - end < PRIME_MIN_T {
1296            if ring_on {
1297                let shifted = t - PRIME_MIN_T;
1298                end = if shifted > start { shifted } else { t };
1299            } else {
1300                end = t;
1301            }
1302        }
1303        ranges.push((start, end));
1304        start = end;
1305    }
1306    ranges
1307}
1308
1309fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
1310    let prefix = prefix as u128;
1311    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
1312}
1313
1314fn dynamic_prime_chunk_ranges(
1315    t: usize,
1316    fixed_chunk: usize,
1317    fixed: &[(usize, usize)],
1318) -> Vec<(usize, usize)> {
1319    let n = fixed.len();
1320    if n < 3 {
1321        return fixed.to_vec();
1322    }
1323
1324    let max_first = t - (n - 1) * PRIME_MIN_T;
1325    let first = fixed_chunk
1326        .div_ceil(2)
1327        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
1328        .min(max_first);
1329    let mut ranges = Vec::with_capacity(n);
1330    ranges.push((0, first));
1331
1332    let first_work = prime_chunk_work(first, t);
1333    let work_span = prime_chunk_work(t, t) - first_work;
1334    let denominator = (n - 1) as u128;
1335    let mut previous = first;
1336    for boundary in 1..n - 1 {
1337        let target = first_work * denominator + work_span * (boundary as u128);
1338        let remaining = n - 1 - boundary;
1339        let mut low = previous + PRIME_MIN_T;
1340        let mut high = t - remaining * PRIME_MIN_T;
1341        while low < high {
1342            let mid = low + (high - low) / 2;
1343            if prime_chunk_work(mid, t) * denominator >= target {
1344                high = mid;
1345            } else {
1346                low = mid + 1;
1347            }
1348        }
1349        ranges.push((previous, low));
1350        previous = low;
1351    }
1352    ranges.push((previous, t));
1353    ranges
1354}
1355
1356/// Internal prime ranges. A pipelined PP prime defaults to a short-fill, equal-modeled-time
1357/// schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
1358/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
1359///
1360/// `gdn_grid`: the model runs the chunked GDN WY scan (`HybridModel::gdn_prime_grid_on`) —
1361/// AUTO-scheduled internal boundaries are then snapped down to the WY-chunk grid
1362/// (`align_prime_ranges_to_gdn`; the spec-longctx grid law, extended from serve splits to
1363/// the PP prime microchunks). Explicit MEMRA_PRIME_CHUNK keeps its operator-authoritative
1364/// (fixed, unaligned) semantics — the FLAGS caveat documents that identity contract.
1365pub fn prime_chunk_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1366    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
1367    let chunk = prime_chunk_tokens(t, n_layers);
1368    let fixed = fixed_prime_chunk_ranges(t, chunk);
1369    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
1370        Ok(value) => value == "dynamic",
1371        Err(_) => true,
1372    };
1373    if explicit_chunk {
1374        return fixed;
1375    }
1376    let ranges = if !dynamic || !prime_pipeline_auto_geometry(n_layers) {
1377        fixed
1378    } else {
1379        dynamic_prime_chunk_ranges(t, chunk, &fixed)
1380    };
1381    // MEMRA_PRIME_GRID_ALIGN=0 is the shared rollback seam of the grid law (same env the
1382    // worker's serve-boundary alignment honors, read per call so gates can flip it
1383    // in-process): the legacy off-grid auto schedule — the toothed cell's broken arm.
1384    if gdn_grid && std::env::var("MEMRA_PRIME_GRID_ALIGN").as_deref() != Ok("0") {
1385        align_prime_ranges_to_gdn(&ranges, t, Engine::gdn_chunk_size())
1386    } else {
1387        ranges
1388    }
1389}
1390
1391/// The mHC prime schedule: how `prime_cache_hyper` SPLITS one prompt into calls.
1392///
1393/// THIS IS THE RULE THE CAPACITY GATE ASSERTS ON, and it is separate from
1394/// [`prime_chunk_ranges`] so the hyper walk's split can be read, gated and changed without
1395/// touching the serial trunk's. It DELEGATES to the serial schedule rather than deriving a
1396/// second one: the transient pressure it answers is the same pressure, and two schedules would
1397/// be two things to keep aligned with `MEMRA_PRIME_CHUNK`.
1398///
1399/// WHY THE SPLIT IS SEMANTICALLY INERT, term by term. Read this precisely: it says the split
1400/// computes the SAME PROGRAM, not that it computes the same BITS. The bit claim is FALSE on this
1401/// trunk and measured so — see the near-tie note at the end.
1402///
1403///   * **The mHC residual is strictly PER TOKEN.** `crate::hyper`'s contract is `mixes[t,:]`, an
1404///     RMS rescale over that token's own `streams*hidden` slab, a Sinkhorn per token per site,
1405///     a per-token collapse and a per-token post. The stream state is expanded at the start of a
1406///     call and collapsed at its end; it carries NOTHING between tokens. Splitting the token
1407///     axis cannot move a value.
1408///   * **KDA prefill is a SEQUENTIAL scan** (`kda.rs`: `memra_kda_scan_s128` runs prefill and
1409///     decode alike, the chunked UT transform is not the shipped path). A sequential recurrence
1410///     has no fold grid, so the GDN WY grid law has NO KDA analogue to violate and
1411///     `align_prime_ranges_to_gdn` has nothing to align. The conv ring carries across calls
1412///     already — it is the seam every decode step uses. **DEBT, named:** if the chunked KDA twin
1413///     ever becomes the prefill path, it acquires a fold grid and this schedule's internal
1414///     boundaries must be snapped to it exactly as the GDN ones are, or chunked prime stops
1415///     being bit-identical. The `gdn_grid` argument is the seam that change lands on.
1416///   * **The latent KV plane is f32** (`LatentKvLayer::rows`), so a later call reads earlier
1417///     calls' rows in the SAME numeric class it would have computed them in. There is no
1418///     analogue of the serial trunk's f32-vs-quantized-KV class edge — the thing that made
1419///     `MEMRA_PRIME_CHUNK` steer arithmetic until the 2026-08-05 grain-free fix.
1420///   * **The DSA pool keys are already incremental** (`index_pools_ready`): a pool key is a pure
1421///     function of its own `pool` state rows and the constant `kpool_ape`, final the instant the
1422///     pool's last row lands, so no boundary can move one. Selection is per query over resident
1423///     keys, with visibility keyed on the query's ABSOLUTE cache row.
1424///
1425/// NOT BIT-STABLE ACROSS CHUNK SIZES, and the cause is NOT the split. Measured on the rig
1426/// (`research/glm53-flash-bringup-20260827/1m-context-20260828/02`): the arms diverge at ROW 0,
1427/// which no cross-token state can reach, and `Engine::linear` — the cuBLASLt f32 `mixes` GEMM in
1428/// `hyper::pre` — is itself not m-invariant (m=32 vs m=200 moves 9601/12288 output bits at worst
1429/// 3.815e-6, the same worst the chunked prime reports; m=128 and m=199 vs m=200 are identical).
1430/// cuBLASLt reselects its algorithm by shape and the reduction order goes with it, which
1431/// `hyper.rs`'s header already concedes ("a serving trunk, not a byte-parity oracle"). So this
1432/// is a documented near-tie class the split EXPOSES, not one it creates. It is written into the
1433/// `MEMRA_PRIME_CHUNK` FLAGS row, and `glm5_chunked_prime_gpu` holds the split to a calibrated
1434/// band anchored on `memra_reference::execute` rather than on the monolithic sibling.
1435///
1436/// One thing the split provably cannot break here: POSITIONS. glm5_next is NoPE end to end
1437/// (`qk_rope_head_dim = 0`, `mla_use_nope`), and KDA is positionless, so `pos_d` reaches no
1438/// kernel on this path — a mutation that made it call-local instead of session-absolute moved
1439/// nothing at all.
1440///
1441/// A prompt at or under one chunk takes the monolithic body unchanged, and `MEMRA_PRIME_CHUNK=0`
1442/// restores the monolithic walk up to `PRIME_CHUNK_LAUNCH_CAP` (65,520 tokens; the CUDA
1443/// grid.y wall is 65,535 and the old walk always died at launch above it, 08-29 ppN receipts)
1444/// — the rollback seam, and the oracle arm the correctness gate compares against, both now
1445/// bounded by that cap; longer prompts split at the cap even under `0`.
1446pub fn hyper_prime_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1447    // DOOR `MEMRA_B200_PRIME_V2` arm 1 (default OFF; see the door's header for the arithmetic
1448    // and the numeric class). The delegation above inherits `PRIME_PIPE_MICROBATCHES` from a
1449    // pipeline this walk does not run; under the door the mHC prime takes the NATURAL chunk
1450    // instead. An explicit `MEMRA_PRIME_CHUNK` is still authoritative -- the operator override
1451    // must not be overridden by a door, and `prime_chunk_tokens` already honors it, so the
1452    // door simply stands aside rather than re-deciding.
1453    if b200_prime_v2_on() && std::env::var_os("MEMRA_PRIME_CHUNK").is_none() {
1454        return hyper_prime_ranges_natural(t, gdn_grid);
1455    }
1456    prime_chunk_ranges(t, n_layers, gdn_grid)
1457}
1458
1459/// The mHC prime schedule with the microbatch geometry removed: fixed ranges at
1460/// `PRIME_CHUNK_MAX_TOKENS`, the SAME clamp `prime_chunk_tokens` already hands a prompt too
1461/// long for the microbatch split (which is why the 41.9k prime is 3.3x more efficient per FLOP
1462/// than the 4k one on the same binary). The GDN grid law still applies: `align_prime_ranges_to_gdn`
1463/// is the same call the auto schedule makes, kept here so a chunked-KDA prefill twin cannot
1464/// acquire a fold grid on one path and not the other.
1465fn hyper_prime_ranges_natural(t: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1466    let ranges = fixed_prime_chunk_ranges(t, crate::cache::PRIME_CHUNK_MAX_TOKENS);
1467    let ranges = if gdn_grid && std::env::var("MEMRA_PRIME_GRID_ALIGN").as_deref() != Ok("0") {
1468        align_prime_ranges_to_gdn(&ranges, t, Engine::gdn_chunk_size())
1469    } else {
1470        ranges
1471    };
1472    HYPER_PRIME_NATURAL_SCHEDULES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1473    // ENGAGEMENT RECEIPT, once per process. The first box run of this door produced a 22-45%
1474    // TTFT win with NO line in the log to grep for it: "the flag was set" is not a receipt, and
1475    // a door that silently did nothing would read the same way. Print what the schedule ACTUALLY
1476    // became, not what was asked for.
1477    static SAID: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
1478    if !SAID.swap(true, std::sync::atomic::Ordering::Relaxed) {
1479        eprintln!(
1480            "[prime-v2] arm1 engaged t={t} chunk={} chunks={} widths={:?} (natural chunk \
1481             replaces the PRIME_PIPE_MICROBATCHES geometry; MEMRA_B200_PRIME_V2, logged once \
1482             per process)",
1483            crate::cache::PRIME_CHUNK_MAX_TOKENS,
1484            ranges.len(),
1485            ranges.iter().map(|&(a, b)| b - a).collect::<Vec<_>>(),
1486        );
1487    }
1488    ranges
1489}
1490
1491/// The largest number of token rows any ONE mHC prime call carries under
1492/// [`hyper_prime_ranges`]. Every per-call transient in the walk — the `t*streams*hidden` stream
1493/// state, the MLA query planes, and the DSA indexer's `t * n_pools` score plane — is
1494/// proportional to this, so it is the single number a capacity assertion needs.
1495/// Admission workspace shape of the NON-hyper prime (memra#144): the GDN / dense-MoE trunk
1496/// (ornith, qwen-hybrid) and the step-TP trunk (step-3.7). `prime_layers` keeps its trunk
1497/// transients in the retained prime slab pool sized by the CALL's row count
1498/// (`prime_slabs_get`): seven f32 `[t, n_embd]` planes (h, x1, z, xa, xb, ffn_out, mixed),
1499/// two f16 mirrors (h16, z16), three f32 `[t, n_ff_max]` planes (act, gate, up), plus the
1500/// per-call f16 activation (`a16`), the f16 pre-norm operand, and the `x`/`hn` copies the
1501/// tail takes. The whole-prompt `hiddens` stack (`[t, n_embd]` f32) lives for the prime.
1502/// Before this shape existed the admission line charged `0MB prefill-workspace` on these
1503/// paths, which is how a 164k ornith prompt was admitted onto a card that could not hold
1504/// its monolithic prime (darklanes research/memra146-ornith-repro-20260904: 22 GB more
1505/// resident after a 257k sweep with `MEMRA_PRIME_CHUNK=0` than with 4096) and how a 91k
1506/// step37 prompt stepped into a CUDA OOM after admission priced it at 3.1 GB.
1507#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1508pub struct PrimeWorkspaceShape {
1509    /// Bytes per row of ONE prime call (the slab planes above).
1510    pub call_row_bytes: usize,
1511    /// Bytes per PROMPT row that live for the whole prime (the returned hiddens stack).
1512    pub prompt_row_bytes: usize,
1513    /// Trunk depth, for the auto chunk geometry (`prime_chunk_tokens`).
1514    pub n_layers: usize,
1515}
1516
1517impl PrimeWorkspaceShape {
1518    /// The charge for `prompt_rows` with an explicit per-call row count (pure arithmetic,
1519    /// what the tests pin).
1520    pub fn admission_bytes_with_call_rows(&self, prompt_rows: usize, call_rows: usize) -> usize {
1521        let rows = prompt_rows.min(call_rows.max(1));
1522        self.call_row_bytes
1523            .saturating_mul(rows)
1524            .saturating_add(self.prompt_row_bytes.saturating_mul(prompt_rows))
1525    }
1526
1527    /// The charge under the deployment's prime chunking: `MEMRA_PRIME_CHUNK` (0 = monolithic
1528    /// up to the launch cap, so the call rows are the whole prompt) or the auto geometry.
1529    pub fn admission_bytes(&self, prompt_rows: usize) -> usize {
1530        self.admission_bytes_with_call_rows(
1531            prompt_rows,
1532            prime_chunk_tokens(prompt_rows, self.n_layers),
1533        )
1534    }
1535}
1536
1537pub fn hyper_prime_call_rows(t: usize, n_layers: usize, gdn_grid: bool) -> usize {
1538    hyper_prime_ranges(t, n_layers, gdn_grid)
1539        .iter()
1540        .map(|&(start, end)| end - start)
1541        .max()
1542        .unwrap_or(0)
1543}
1544
1545/// Per-request prefill WORKSPACE coefficients for a HyperConnections trunk, published to
1546/// admission (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper model: their
1547/// admission arithmetic is byte-identical to the pre-lane behavior.
1548///
1549/// These are the FORMULA behind the 262k 2-card cell's measured ~0.8 MiB/token/card prefill
1550/// wall (`research/glm53-flash-bringup-20260827/262k-2card-20260830/LANE.md`), not the slope
1551/// itself: each term is the size of a named allocation in the walk, summed per token of ONE
1552/// prime call. On GLM-5.3-Flash geometry (H=4096, S=4, F=2048, U=8, heads=64, qk=256, v=256,
1553/// topk=2048, P=4) `chunk_token_bytes` evaluates to ~0.86 MiB — the receipt's slope with the
1554/// conservative side up. The attribution table naming every term lives in
1555/// `research/glm53-flash-bringup-20260827/gpf-workspace-20260830/LANE.md` §1.
1556#[derive(Debug, Clone, Copy)]
1557pub struct HyperPrimeWorkspaceShape {
1558    /// Bytes of per-call prefill transients PER TOKEN OF ONE PRIME CALL: the double-buffered
1559    /// `[t, streams, hidden]` stream state + ppN boundary slots, the pre/norm transients, the
1560    /// grouped-MoE staging (CSR activations + three f32 partial planes + f16 mirrors + scatter
1561    /// planes), the MLA query/attention planes, the k-pool idx plane, and the prime-tail
1562    /// hidden/norm pair. Multiplied by [`hyper_prime_call_rows`] this bounds the workspace of
1563    /// the CHUNKED prime; on the monolithic rollback (`MEMRA_PRIME_CHUNK=0`) the call rows are
1564    /// the whole prompt up to `PRIME_CHUNK_LAUNCH_CAP` (65,535 launch-legal max; admission
1565    /// re-derives `hyper_prime_call_rows` so the arithmetic stays consistent either way) and
1566    /// the same product stays honest.
1567    pub chunk_token_bytes: usize,
1568    /// Bytes per PROMPT token that live for the WHOLE prime on the last stage: the returned
1569    /// pre-output_norm `hiddens` stack (`n_embd` f32), consumed by the MTP-spec `prompt_h`
1570    /// and the embed capture.
1571    pub prompt_bytes_per_token: usize,
1572    /// DSA k-pool group size `P`, or 0 when the model runs no k-pool indexer. The selection
1573    /// score plane of ONE call is `call_rows * (ctx / P)` f32 — the one prefill transient that
1574    /// stays COUPLED TO CONTEXT DEPTH after chunking (it is the allocation the 3-card 1M prime
1575    /// died on at 97.2 GiB).
1576    pub kpool_score_pool: usize,
1577    /// Trunk layer count, for re-deriving [`hyper_prime_call_rows`] at admission time with the
1578    /// same env-sensitive schedule the prime itself will walk.
1579    pub n_layers: usize,
1580    /// The model's own GDN grid-alignment input to the schedule.
1581    pub gdn_grid: bool,
1582}
1583
1584impl HyperPrimeWorkspaceShape {
1585    /// The admission charge for one request: workspace of the LARGEST prime call this request
1586    /// can produce, plus the ctx-coupled score plane at that call width, plus the prompt-long
1587    /// hiddens stack.
1588    ///
1589    /// Keyed on PROMPT rows, deliberately not on `ctx_cap`: every term here is a function of
1590    /// what the PRIME walks, and a `max_tokens`-omitted request carries a `ctx_cap` of the
1591    /// whole server window — charging the window would refuse every vendor-default short
1592    /// prompt on a deep-window box for workspace it never allocates. A continuation request's
1593    /// `prompt` is the full rendered conversation (the suffix optimization is internal
1594    /// reuse), so the score plane's `t_kv` is covered too.
1595    pub fn admission_bytes(&self, prompt_rows: usize) -> usize {
1596        let rows = hyper_prime_call_rows(prompt_rows, self.n_layers, self.gdn_grid);
1597        let chunk = self.chunk_token_bytes.saturating_mul(rows);
1598        let score = prompt_rows
1599            .checked_div(self.kpool_score_pool)
1600            .map(|pools| rows.saturating_mul(pools).saturating_mul(size_of::<f32>()))
1601            .unwrap_or(0);
1602        chunk
1603            .saturating_add(score)
1604            .saturating_add(self.prompt_bytes_per_token.saturating_mul(prompt_rows))
1605    }
1606}
1607
1608/// Snap AUTO prime-range internal boundaries DOWN to the GDN WY-chunk grid (lane/
1609/// hermes-perf-fixes, 2026-08-23 — the missing helper the PP-auto-ranges finding names).
1610///
1611/// THE LAW THIS EXTENDS (measured, research/multiturn-cache-20260821/
1612/// LONGCTX-EXACTNESS-20260821.md; the serve-split half already ships as the worker's
1613/// `grid_align_boundary`): under the chunked WY scan a prompt primed as two calls split at
1614/// L is bit-identical to the monolithic prime iff `L % gdn_chunk_size() == 0` — an off-grid
1615/// call start shifts the fold grid and materializes recurrent state at a point the
1616/// monolithic program never computes. The prime loop walks these ranges as separate
1617/// `prime_layers` calls, so INTERNAL microchunk boundaries are the same seam: the PP-2
1618/// auto geometry (`t.div_ceil(8).max(128)` fills, and every dynamic short-fill boundary)
1619/// lands off the 32-token grid for most prompt lengths, which is exactly the
1620/// chunk-value bit-identity the GDN lane falsified (FLAGS PRIME_CHUNK/SCHED caveat).
1621///
1622/// Boundaries only move DOWN (earlier is always semantically safe — same argument as the
1623/// worker's alignment); a boundary that collapses onto its predecessor is dropped (ranges
1624/// merge). The final range always ends at `t`. Aligning down only GROWS the tail
1625/// remainder, so the fixed-schedule tail-merge rule is never re-violated. Cost bound: at
1626/// most `c-1` tokens shift per boundary.
1627pub fn align_prime_ranges_to_gdn(
1628    ranges: &[(usize, usize)],
1629    t: usize,
1630    c: usize,
1631) -> Vec<(usize, usize)> {
1632    if c == 0 || ranges.len() < 2 {
1633        return ranges.to_vec();
1634    }
1635    let mut out: Vec<(usize, usize)> = Vec::with_capacity(ranges.len());
1636    let mut start = 0usize;
1637    for (i, &(_, end)) in ranges.iter().enumerate() {
1638        let e = if i + 1 == ranges.len() {
1639            t
1640        } else {
1641            end / c * c
1642        };
1643        if e > start {
1644            out.push((start, e));
1645            start = e;
1646        } // else: boundary collapsed onto its predecessor — merge into the next range
1647    }
1648    debug_assert_eq!(out.last().map(|&(_, e)| e), Some(t));
1649    out
1650}
1651
1652struct HeadSplit {
1653    pin: u64,
1654    w1: CudaSlice<u8>,
1655    hn1: CudaSlice<f32>,
1656    y1: CudaSlice<f32>,
1657    logits_e: CudaSlice<f32>,
1658    ev_hn: cudarc::driver::CudaEvent,
1659    ev_done: cudarc::driver::CudaEvent,
1660    raw_hn1: u64,
1661    raw_y1: u64,
1662    raw_logits_hi: u64,
1663    /// SAMPLED-TAIL scratch (perturbed row + the filter's threshold/z/max slots + the row
1664    /// index). Allocating these per token cost more than the split head saved: the first
1665    /// sampled-split measurement came in at 78.25 tok/s against 78.96 for the unsplit head,
1666    /// which is five allocations per token, not arithmetic.
1667    samp: Option<SampScratch>,
1668}
1669
1670struct SampScratch {
1671    pb: CudaSlice<f32>,
1672    th: CudaSlice<f32>,
1673    z: CudaSlice<f32>,
1674    mx: CudaSlice<f32>,
1675    rows: CudaSlice<i32>,
1676}
1677/// HEAD-SPLIT workspace (host + device twins share it).
1678static HEAD_SPLIT_WS: std::sync::Mutex<Option<HeadSplit>> = std::sync::Mutex::new(None);
1679
1680/// DEV1-LOCAL ROUTER replicas (MEMRA_DEV1_ROUTER): per-layer (gate_inp_f32, exp_probs_b,
1681/// active_experts) on rank1 + a shared logits scratch. Deterministic kernels on identical
1682/// input bits — rank1's local selection is bit-equal to the root's.
1683#[allow(clippy::type_complexity)]
1684static DEV1_ROUTER_REPS: std::sync::Mutex<
1685    Option<(
1686        std::collections::HashMap<u16, (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<u8>)>,
1687        Option<CudaSlice<f32>>,
1688    )>,
1689> = std::sync::Mutex::new(None);
1690
1691/// SHEXP-ON-DEV1 workspace (MEMRA_SHEXP_DEV1): replica weights + scratch on rank1, the
1692/// down row lands ROOT-resident over P2P (single store pass), and apply adds it on e
1693/// behind ev_done. (pins, wg1, wu1, wd1, act1, sh_root, ev_z, ev_done).
1694#[allow(clippy::type_complexity)]
1695static SHEXP_D1_REPS: std::sync::Mutex<
1696    Option<std::collections::HashMap<u16, (CudaSlice<u8>, CudaSlice<u8>, CudaSlice<u8>)>>,
1697> = std::sync::Mutex::new(None);
1698#[allow(clippy::type_complexity)]
1699static SHEXP_D1_WS: std::sync::Mutex<
1700    Option<(
1701        (usize, usize),
1702        CudaSlice<f32>,
1703        CudaSlice<f32>,
1704        CudaSlice<f32>,
1705        cudarc::driver::CudaEvent,
1706        cudarc::driver::CudaEvent,
1707    )>,
1708> = std::sync::Mutex::new(None);
1709
1710/// SHEXP OVERLAP workspace (issue writes, apply reads): (device, n_embd, n_ff_sh, act, sh).
1711#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1712static SHEXP_OV_WS: std::sync::Mutex<
1713    Option<(usize, usize, usize, CudaSlice<f32>, CudaSlice<f32>)>,
1714> = std::sync::Mutex::new(None);
1715
1716impl HybridModel {
1717    /// Does this model's prime schedule live under the GDN WY-chunk grid law? True when the
1718    /// trunk has GDN (linear-attention) layers AND the chunked scan is on — the regime where
1719    /// an off-grid prime-call boundary shifts the WY fold grid (see
1720    /// `align_prime_ranges_to_gdn`). Attention-only models and the sequential scan
1721    /// (`MEMRA_GDN_CHUNKED=0`) are split-invariant, so the grid is a no-op contract there.
1722    pub fn gdn_prime_grid_on(&self) -> bool {
1723        Engine::gdn_chunked_enabled()
1724            && self
1725                .layers
1726                .iter()
1727                .any(|l| matches!(l.mixer, crate::hybrid::Mixer::Linear(_)))
1728    }
1729
1730    /// Can the step TP runtime run the DEVICE-RESIDENT activation path from this serving
1731    /// engine? Native P2P (peer copies replace the host staging) AND a shared root context
1732    /// (the device buffers must be addressable on both sides — the TP registry builds its
1733    /// own Engine per rank, so this is a real seam, not a formality).
1734    fn full_attn_tp_device_resident(e: &Engine, tp: &crate::hybrid::StepTpQkv) -> bool {
1735        tp.runtime.native_p2p() && tp.runtime.root_shares_ctx(e)
1736    }
1737
1738    pub(crate) fn full_attn_tp_qkv(
1739        &self,
1740        e: &Engine,
1741        fa: &FullAttnLayer,
1742        h: &CudaSlice<f32>,
1743        t: usize,
1744    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
1745        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1746            return Ok(None);
1747        };
1748        let values = active_matrix_values(
1749            h.len(),
1750            t,
1751            self.cfg.n_embd as usize,
1752            "Step TP QKV activation",
1753        )?;
1754        // DEVICE-RESIDENT NATIVE PATH (lane/hermes-perf-fixes, 2026-08-23 — the host-bounce
1755        // finding): the native-P2P arm used to dtoh the FULL hidden state per layer, run
1756        // from a host copy, gather q/k/v to host vectors, and htod all three back — a host
1757        // round-trip on every execute that the peer transport exists to remove. The
1758        // device twins are byte-identical by construction (the same bytes travel dtod
1759        // instead of dtoh+htod; kernels, peer copies, and gather order are shared code).
1760        // The host arm below remains the transport for !native_p2p (host staging IS that
1761        // transport) and for a root context this engine cannot address.
1762        if Self::full_attn_tp_device_resident(e, tp) {
1763            // Producer fence: h was written on THIS engine's stream; the TP ranks read it
1764            // on theirs (same context, different streams).
1765            e.stream().synchronize()?;
1766            let q = tp
1767                .runtime
1768                .bf16_column_parallel_resident_native_device(&tp.q, h, t)?;
1769            let k = tp
1770                .runtime
1771                .bf16_column_parallel_resident_native_device(&tp.k, h, t)?;
1772            let v = tp
1773                .runtime
1774                .bf16_column_parallel_resident_native_device(&tp.v, h, t)?;
1775            Self::full_attn_tp_log_once(tp, "qkv", "device-resident");
1776            return Ok(Some(vec![q, k, v]));
1777        }
1778        let host = e.dtoh_view(&h.slice(0..values))?;
1779        let q = if tp.runtime.native_p2p() {
1780            tp.runtime
1781                .bf16_column_parallel_resident_native(&tp.q, &host, t)?
1782        } else {
1783            tp.runtime
1784                .bf16_column_parallel_resident(&tp.q, &host, t)?
1785                .gathered
1786        };
1787        let k = if tp.runtime.native_p2p() {
1788            tp.runtime
1789                .bf16_column_parallel_resident_native(&tp.k, &host, t)?
1790        } else {
1791            tp.runtime
1792                .bf16_column_parallel_resident(&tp.k, &host, t)?
1793                .gathered
1794        };
1795        let v = if tp.runtime.native_p2p() {
1796            tp.runtime
1797                .bf16_column_parallel_resident_native(&tp.v, &host, t)?
1798        } else {
1799            tp.runtime
1800                .bf16_column_parallel_resident(&tp.v, &host, t)?
1801                .gathered
1802        };
1803        Self::full_attn_tp_log_once(tp, "qkv", "host-canonical");
1804        Ok(Some(vec![e.htod(&q)?, e.htod(&k)?, e.htod(&v)?]))
1805    }
1806
1807    /// One transport banner per (projection, transport) — the old per-call eprintln fired
1808    /// on EVERY layer of EVERY step, itself a decode-rate cost on the path this lane is
1809    /// unbouncing (the sibling grouped-EP path already learned this).
1810    fn full_attn_tp_log_once(tp: &crate::hybrid::StepTpQkv, proj: &str, activation: &'static str) {
1811        use std::sync::atomic::{AtomicBool, Ordering};
1812        static LOGGED: [AtomicBool; 4] = [
1813            AtomicBool::new(false),
1814            AtomicBool::new(false),
1815            AtomicBool::new(false),
1816            AtomicBool::new(false),
1817        ];
1818        let idx = 2 * usize::from(proj == "o") + usize::from(activation == "device-resident");
1819        if LOGGED[idx].swap(true, Ordering::Relaxed) {
1820            return;
1821        }
1822        eprintln!(
1823            "[step-tp-{proj}] execute layer={} devices={:?} projections={proj} \
1824             tensor_parallel=true attention_local=true kv_local=true transport={} \
1825             native_p2p={} bulk_p2p={} activation={activation} \
1826             output={} performance_claim=false (logged once per transport)",
1827            tp.layer,
1828            tp.devices,
1829            tp.runtime.transport_label(),
1830            tp.runtime.native_p2p(),
1831            tp.runtime.bulk_p2p(),
1832            if activation == "device-resident" {
1833                "root-resident"
1834            } else {
1835                "root-readback"
1836            },
1837        );
1838    }
1839
1840    pub(crate) fn full_attn_tp_o(
1841        &self,
1842        e: &Engine,
1843        fa: &FullAttnLayer,
1844        activation: &CudaSlice<f32>,
1845        tokens: usize,
1846    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1847        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1848            return Ok(None);
1849        };
1850        // DEVICE-RESIDENT NATIVE PATH — the O-projection half of the same finding: no DtoH
1851        // of the attention output, no host O staging, root-resident reduction consumed in
1852        // place (byte-identical shared core: `step_bf16_row_native_reduce_from_root`).
1853        if Self::full_attn_tp_device_resident(e, tp) {
1854            e.stream().synchronize()?; // producer fence, as the QKV half
1855            let output = tp
1856                .runtime
1857                .step_bf16_row_parallel_resident_native_device(&tp.o, activation, tokens)?;
1858            Self::full_attn_tp_log_once(tp, "o", "device-resident");
1859            return Ok(Some(output));
1860        }
1861        let host = e.dtoh(activation)?;
1862        let output = if tp.runtime.native_p2p() {
1863            tp.runtime
1864                .step_bf16_row_parallel_resident_native(&tp.o, &host, tokens)?
1865        } else {
1866            tp.runtime
1867                .step_bf16_row_parallel_resident(&tp.o, &host, tokens)?
1868        };
1869        Self::full_attn_tp_log_once(tp, "o", "host-canonical");
1870        Ok(Some(e.htod(&output)?))
1871    }
1872
1873    fn full_attn_o(
1874        &self,
1875        e: &Engine,
1876        fa: &FullAttnLayer,
1877        activation: &CudaSlice<f32>,
1878        tokens: usize,
1879    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1880        match self.full_attn_tp_o(e, fa, activation, tokens)? {
1881            Some(output) => Ok(output),
1882            None => e.matmul(&fa.wo, activation, tokens),
1883        }
1884    }
1885
1886    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
1887    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
1888    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
1889    /// (it forces a dtoh + host hash per layer).
1890    fn prime_trace_path() -> Option<&'static str> {
1891        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
1892        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
1893            .as_deref()
1894    }
1895
1896    /// PRIME ANATOMY (diagnostic): `MEMRA_PRIME_ANATOMY=1` synchronizes the stream around
1897    /// each prime_layers stage and accumulates wall time per stage class, printed after
1898    /// every prime_layers call (cumulative across chunks/reps). The per-stage syncs
1899    /// serialize launch/execute overlap, so the summed total exceeds the naked prime wall —
1900    /// attribution ratios only, never a measured default run. Non-seg serial arm only.
1901    fn prime_anatomy_on() -> bool {
1902        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1903        *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
1904    }
1905
1906    fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
1907        static S: [std::sync::atomic::AtomicU64; 5] = [
1908            std::sync::atomic::AtomicU64::new(0), // 0 mixer full-attn
1909            std::sync::atomic::AtomicU64::new(0), // 1 mixer linear-attn (GDN)
1910            std::sync::atomic::AtomicU64::new(0), // 2 ffn MoE (router + experts + shexp)
1911            std::sync::atomic::AtomicU64::new(0), // 3 ffn dense
1912            std::sync::atomic::AtomicU64::new(0), // 4 norms/adds/glue
1913        ];
1914        &S
1915    }
1916
1917    /// Fail closed on a path that has not been taught the mHC residual program.
1918    ///
1919    /// A serial residual on an hc model is not a degraded answer, it is a DIFFERENT function
1920    /// computed at full speed and full confidence — the exact failure `crate::hyper` exists to
1921    /// prevent. Every trunk entry point that has not been converted calls this first, so the
1922    /// unconverted set is a list of named refusals rather than a list of silent wrong answers.
1923    pub(crate) fn refuse_hyper(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
1924        if let Some(topology) = self.hyper.as_ref() {
1925            return Err(format!(
1926                "{path} runs a serial residual, but this model's ModelPlan declares \
1927                 ResidualTopology::HyperConnections{{ streams: {}, collapse: {:?} }}. Refusing: \
1928                 that path would compute a different model. Converted paths: forward, \
1929                 forward_last, prime_cache, decode_step, and the batched serving chain \
1930                 decode_step_batch / _sampled / _lean / _masked.",
1931                topology.streams, topology.collapse
1932            )
1933            .into());
1934        }
1935        Ok(())
1936    }
1937
1938    /// The FFN branch of one hc site, from an already-normed `[t, hidden]` input.
1939    ///
1940    /// Split out because under hyper-connections the FFN's input is `rms_norm(hc_pre(x))`, not
1941    /// `rms_norm(x + attn)` — the fused add+norm+quantize forms the serial paths use have no
1942    /// residual to fold, so this is the unfused dispatch by construction.
1943    /// `zq8`: an optional pre-quantized q8_1 pair for `z` (door `MEMRA_GLM5_Q8_FUSE`,
1944    /// lane/b200-q8-fuse-20260902 — the caller's norm producer folded the standalone
1945    /// `quantize_q8_1(z, ...)` launch into itself). `None` preserves the unfused chain:
1946    /// `moe_ffn_il_zq8` re-quantizes `z` itself, byte-identical either way.
1947    #[allow(clippy::too_many_arguments)] // allow: mirrors the mixer/FFN dispatch contract every sibling walk in this file shares; bundling into a struct is a refactor, not a lint fix, and would touch every call site for no behavior change.
1948    pub(crate) fn hyper_ffn_branch(
1949        &self,
1950        e: &Engine,
1951        layer: &crate::hybrid::HybridLayer,
1952        z: &CudaSlice<f32>,
1953        t: usize,
1954        il: usize,
1955        prefill: bool,
1956        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
1957    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1958        match &layer.ffn {
1959            crate::hybrid::Ffn::Dense {
1960                ffn_gate,
1961                ffn_up,
1962                ffn_down,
1963            } => {
1964                let n_ff = ffn_gate.out_features();
1965                let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], z, t)?;
1966                let up = g2.pop().unwrap();
1967                let gate = g2.pop().unwrap();
1968                let mut act = e.uninit(t * n_ff)?;
1969                // A dense FFN reads the SHEXP clamp array — see forward()'s note.
1970                Self::ffn_act_lim(
1971                    e,
1972                    &self.cfg,
1973                    &gate,
1974                    &up,
1975                    1.0,
1976                    1.0,
1977                    self.cfg.clamp_shexp_at(il as u32),
1978                    &mut act,
1979                    t * n_ff,
1980                )?;
1981                e.matmul(ffn_down, &act, t)
1982            }
1983            crate::hybrid::Ffn::Moe(m) => {
1984                if prefill {
1985                    self.moe_ffn_il_prefill(e, m, z, t, il as u16)
1986                } else {
1987                    self.moe_ffn_il_zq8(e, m, z, zq8, t, il as u16)
1988                }
1989            }
1990        }
1991    }
1992
1993    /// Stateless prefill under the mHC residual (`crate::hyper`), the hc twin of `forward` /
1994    /// `forward_last`. The mixers, the FFNs and the norms are the SAME calls the serial paths
1995    /// make; only the residual program around them changes, which is the whole point — a mixer
1996    /// never sees the stream dimension.
1997    fn forward_hyper(
1998        &self,
1999        e: &Engine,
2000        tokens: &[u32],
2001        last_only: bool,
2002    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2003        let topology = *self
2004            .hyper
2005            .as_ref()
2006            .ok_or("forward_hyper on a model with no HyperConnections topology")?;
2007        // M2 ppN door for the mHC trunk. The generic arm's door lives in `decode_step_h`;
2008        // this walk is reached BEFORE it (decode.rs routes `hyper.is_some()` first), so the
2009        // hc walks own their own door. Loud refusal, never silent fallback: an unqualified
2010        // pipeline rewrite errors here rather than running a single-engine walk over weights
2011        // the loader has already sharded across devices.
2012        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2013            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
2014                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2015            }
2016            return self.forward_hyper_ppn(e, tokens, last_only, &topology, &fence);
2017        }
2018        let n_embd = self.cfg.n_embd as usize;
2019        let t = tokens.len();
2020        let eps = self.cfg.rms_eps;
2021        let pos: Vec<i32> = (0..t as i32).collect();
2022        let pos_d = e.htod_i32(&pos)?;
2023
2024        let embedded = self.embed(e, tokens)?;
2025        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
2026        let trace = memra_reference::hidden_trace::enabled();
2027        if trace {
2028            memra_reference::hidden_trace::emit_tokens(tokens);
2029            let streams = x.len() / (t * n_embd);
2030            memra_reference::hidden_trace::emit_last_row(
2031                "expand",
2032                -1,
2033                t,
2034                streams * n_embd,
2035                &e.dtoh(&x)?,
2036            );
2037        }
2038
2039        x = self.hyper_range_forward(e, &topology, x, 0, self.layers.len(), &pos_d, t, trace)?;
2040
2041        // SHARED EXIT with the ppN twin: one trunk exit, so the split and unsplit arms cannot
2042        // drift apart in the head. That is what makes `glm5-hyper-ppn-gate`'s bit-identity bar
2043        // a structural property rather than a coincidence of two maintained copies.
2044        self.hyper_head_logits(e, &topology, &x, t, n_embd, eps, last_only)
2045    }
2046
2047    /// Stateful prefill under the mHC residual: `prime_cache_overlaid`'s contract (leave a
2048    /// decode-ready cache behind, return last-row logits + the pre-output_norm hidden seed and
2049    /// stack) over the hc layer program.
2050    ///
2051    /// Deliberately UNCHUNKED and UNCAPTURED. The serial prime's chunking, prime slabs, S-mid
2052    /// graph capture and core-split arms are all keyed to the serial residual's transient set;
2053    /// re-deriving them for a stream state is a tuning lane, not a correctness one, and this
2054    /// path is the one the reference gate pins. Long prompts therefore hold `T*streams*hidden`
2055    /// f32 of stream state — 4x the serial trunk's — and that ceiling is the named cost of the
2056    /// simple form.
2057    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2058    fn prime_cache_hyper(
2059        &self,
2060        e: &Engine,
2061        tokens: &[u32],
2062        cache: &mut Cache,
2063        queued_after: usize,
2064        overlay: Option<&crate::vision::EmbedOverlay>,
2065    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2066        let topology = *self
2067            .hyper
2068            .as_ref()
2069            .ok_or("prime_cache_hyper on a model with no HyperConnections topology")?;
2070        // M2 ppN door for the mHC prime (see `forward_hyper`'s note). glm5_next has no
2071        // batched prime arm: the split twin is a straight layer-range split with the
2072        // [t, streams, hidden] state on the wire.
2073        //
2074        // CHUNKED, like the single-engine walk below (lane/glm53-1m-demo, 2026-08-29 — the
2075        // follow-up the previous note named). The monolithic ppN prime carried the WHOLE
2076        // prompt as one call, which capped it three independent ways on the 4x96 GB box:
2077        // per-call transients proportional to t OOM'd from ~32k tokens, and every launch
2078        // that places t in grid.y (kda_conv_silu, kda_gate, rms_norm over rows, the router)
2079        // hits the CUDA 65,535 grid.y ceiling from t=65,536 (measured: instant
2080        // CUDA_ERROR_INVALID_VALUE at a 128,566-token prime, receipts in
2081        // research/glm53-flash-bringup-20260827/1m-demo-20260829/). The chunk loop reuses
2082        // the SAME schedule as the single-engine walk (`hyper_prime_ranges`), so per-chunk
2083        // t is bounded and — because the per-chunk staged walk is bit-identical to the
2084        // per-chunk unsplit walk (glm5_hyper_ppn_gate arm 2) — the chunked ppN prime
2085        // composes to bit-identity with the chunked single-engine prime over the same
2086        // schedule. `queued_after + (t - end)` keeps the REQUEST-level `seq_end` invariant
2087        // across chunks (each call recomputes pos0+start + (end-start) + rest = pos0 + t +
2088        // queued_after). A prompt at or under one chunk takes the monolithic ppN body
2089        // unchanged, and `MEMRA_PRIME_CHUNK=0` restores the monolithic walk up to
2090        // PRIME_CHUNK_LAUNCH_CAP (grid.y-legal ceiling; above it the old walk always died
2091        // at launch) — the same rollback seam the single-engine chunk walk documents.
2092        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2093            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
2094                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2095            }
2096            // The mixed-embedding overlay rides the ppN twin too (lane/glm5-vision-default-on,
2097            // 2026-08-30): the splice is an EMBEDDING-INTAKE transform, and the ppN walk embeds
2098            // on stage 0 only — every later stage receives the already-expanded stream state.
2099            // Under the chunked ppN prime each chunk takes the overlay WINDOWED to its own
2100            // call-relative range (`EmbedOverlay::window`, the same rebase seam the serve
2101            // prefill tick uses), so splice placement is chunk-schedule-invariant. Gated by
2102            // glm5-hyper-ppn-gate's overlay arm (bit-identity vs the substituted-token truth,
2103            // red arm = shifted spans).
2104            let n_embd = self.cfg.n_embd as usize;
2105            let t = tokens.len();
2106            if cache.pos + t > cache.max_ctx {
2107                return Err("prime_cache: prompt exceeds cache max_ctx".into());
2108            }
2109            let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
2110            // Device-resident tap ingest hooks (lane/spec-route-depth-20260902): a no-op
2111            // unless the cache's tap sink carries an ingest consumer; then each completed
2112            // range's tap rows go straight into the drafter KV at the range boundary, the
2113            // one point where every stage's writes for the range have retired.
2114            if ranges.len() == 1 {
2115                self.glm5_taps_range_begin(cache, 0);
2116                let out = self.prime_cache_hyper_ppn(
2117                    e,
2118                    tokens,
2119                    cache,
2120                    queued_after,
2121                    &topology,
2122                    &fence,
2123                    overlay,
2124                )?;
2125                self.glm5_taps_range_done(e, cache, 0, t)?;
2126                return Ok(out);
2127            }
2128            // DOOR `MEMRA_B200_PRIME_V2` arm 2 (default OFF): overlap stage 0 of chunk k+1 with
2129            // stage 1 of chunk k. Every conjunct is load-bearing and every decline is named --
2130            // an arm that quietly did not run is the failure mode this lane's own roofline was
2131            // written to catch.
2132            if b200_prime_v2_on() && !crate::pp::pp2_streams_off() && fence.len() == 3 {
2133                if !crate::pp::prime_pipe_on() {
2134                    hyper_pipe_decline_once("MEMRA_PRIME_PIPE=0 (the operator rollback seam)");
2135                } else if overlay.is_some() {
2136                    hyper_pipe_decline_once(
2137                        "a vision embedding overlay is present: the splice is a stage-0 \
2138                         embedding-intake transform whose gate ran on the serial ppN body",
2139                    );
2140                } else {
2141                    // The DFlash2 hc tap sink used to refuse here, and that refusal was the
2142                    // whole product-route cost: on the pair at 256,756 tokens the SPEC route
2143                    // measured 152.95 s with the door on against 153.12 s off (arm 2 never
2144                    // engaged), while the PLAIN route got 69.08 s against 128.87 s. The sink is
2145                    // now SHARED across the two stage threads instead of being lost to
2146                    // `PrimeCacheStages`' `hc_taps: None` shells — see `HcTapArm` for why that
2147                    // is sound and stays byte-identical to the serial walk's sink.
2148                    let seq_end = cache.pos + t + queued_after;
2149                    let mut sink = cache.hc_taps.take();
2150                    let out = {
2151                        let lock = sink.as_mut().map(std::sync::Mutex::new);
2152                        self.prime_cache_hyper_pp2_pipelined(
2153                            e,
2154                            tokens,
2155                            cache,
2156                            seq_end,
2157                            &topology,
2158                            &ranges,
2159                            &fence,
2160                            lock.as_ref(),
2161                        )
2162                    };
2163                    // Restore the sink whatever happened: a failed prime still owns its
2164                    // draft-source rows, and losing them silently is the shape this change
2165                    // exists to remove.
2166                    if let Some(mut s) = sink {
2167                        // Parity with the serial walk, which leaves `base` at the last chunk's
2168                        // absolute start after the loop.
2169                        if let Some(&(last, _)) = ranges.last() {
2170                            s.base = cache.pos.saturating_sub(t).saturating_add(last);
2171                        }
2172                        cache.hc_taps = Some(s);
2173                    }
2174                    return out;
2175                }
2176            }
2177            let mut hiddens = e.uninit(t * n_embd)?;
2178            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
2179            for &(start, end) in &ranges {
2180                let ov = overlay.and_then(|o| o.window(start, end - start));
2181                self.glm5_taps_range_begin(cache, start);
2182                let (l, hs, x) = self.prime_cache_hyper_ppn(
2183                    e,
2184                    &tokens[start..end],
2185                    cache,
2186                    queued_after + (t - end),
2187                    &topology,
2188                    &fence,
2189                    ov.as_ref(),
2190                )?;
2191                self.glm5_taps_range_done(e, cache, start, end)?;
2192                e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
2193                last = Some((l, hs));
2194                // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
2195                // so the device finished it. Stamp the odometer /health reads, so a BUSY
2196                // worker mid-long-prefill is never mistaken for a wedged one.
2197                crate::progress::note_prime_rows(end - start);
2198            }
2199            let (logits, h_seed) =
2200                last.expect("hyper_prime_ranges never returns an empty schedule");
2201            return Ok((logits, h_seed, hiddens));
2202        }
2203        let n_embd = self.cfg.n_embd as usize;
2204        let t = tokens.len();
2205        if cache.pos + t > cache.max_ctx {
2206            return Err("prime_cache: prompt exceeds cache max_ctx".into());
2207        }
2208        // The REQUEST's absolute end position, computed ONCE before the walk: every chunk sees
2209        // the same value whatever the chunk size, which is the tick-seg law the serial loop
2210        // above carries verbatim (`+ queued_after` closes the serve-split axis).
2211        let seq_end = cache.pos + t + queued_after;
2212        let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
2213        // Device-resident tap ingest hooks (doc at the ppN twin above).
2214        if ranges.len() == 1 {
2215            self.glm5_taps_range_begin(cache, 0);
2216            let out = self.prime_chunk_hyper(e, tokens, cache, seq_end, 0, overlay)?;
2217            self.glm5_taps_range_done(e, cache, 0, t)?;
2218            return Ok(out);
2219        }
2220        let mut hiddens = e.uninit(t * n_embd)?;
2221        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
2222        for &(start, end) in &ranges {
2223            self.glm5_taps_range_begin(cache, start);
2224            let (l, hs, x) =
2225                self.prime_chunk_hyper(e, &tokens[start..end], cache, seq_end, start, overlay)?;
2226            self.glm5_taps_range_done(e, cache, start, end)?;
2227            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
2228            last = Some((l, hs));
2229            // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
2230            // so the device finished it. Stamp the odometer /health reads, so a BUSY
2231            // worker mid-long-prefill is never mistaken for a wedged one.
2232            crate::progress::note_prime_rows(end - start);
2233        }
2234        let (logits, h_seed) = last.expect("hyper_prime_ranges never returns an empty schedule");
2235        Ok((logits, h_seed, hiddens))
2236    }
2237
2238    /// Publish this model's prefill-workspace coefficients to admission
2239    /// (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper trunk — the server's
2240    /// admission arithmetic is then byte-identical to the pre-lane behavior for that family.
2241    ///
2242    /// Term-by-term, each anchored on a named allocation of the hyper prime walk (glm5 numbers
2243    /// in parentheses; the full attribution table is the lane doc's §1):
2244    ///   * stream state, double-buffered at `hyper::post`, PLUS the ppN boundary tx/rx pair of
2245    ///     the same `[t, streams, hidden]` payload: `4 * S * H * 4` (256 KiB/token);
2246    ///   * pre/norm transients (`hyper::pre` y + `rms_norm` h/z + ffn_out): `4 * H * 4`
2247    ///     (64 KiB/token);
2248    ///   * prime-tail hidden/norm pair (`collapse` + output-norm stack): `2 * H * 4`
2249    ///     (32 KiB/token);
2250    ///   * grouped-MoE prefill staging (`moe_ffn_grouped_prefill_sigmoid`): the f16 CSR
2251    ///     activations `U*2H`, three f32 partial planes `3*U*4F` (gate/up/act), the f16 down
2252    ///     mirror `U*2F`, the CSR-order down output + pair-order permute `2*U*4H`, and the
2253    ///     scatter target `4H` — `U*(10H + 14F) + 4H` (560 KiB/token);
2254    ///   * MLA query/attention planes: `heads * (qk_head_dim + v_head_dim) * 4`
2255    ///     (128 KiB/token) and the k-pool idx plane `(topk/P + 1) * 4` (~2 KiB/token).
2256    ///
2257    /// Validation against truth: at GLM-5.3-Flash geometry the sum is ~0.92 MiB per call
2258    /// token, against the 262k cell's MEASURED retained slope of ~0.8 MiB/token/card
2259    /// (vramwatch.csv: +6.3 GiB across the 8,072-token prime) — the formula sits above the
2260    /// measurement, never below it. The ctx-coupled score plane and the prompt-long hiddens
2261    /// stack are separate coefficients on the shape; see [`HyperPrimeWorkspaceShape`].
2262    /// memra#144: the workspace shape of the non-hyper prime; `None` on the hyper trunk,
2263    /// which publishes `hyper_prime_workspace_shape` instead. Mirrors `prime_layers`'
2264    /// `n_ff_max` rule (the widest dense FFN, never below `n_embd`).
2265    pub fn prime_workspace_shape(&self) -> Option<PrimeWorkspaceShape> {
2266        if self.hyper.is_some() {
2267            return None;
2268        }
2269        let h = self.cfg.n_embd as usize;
2270        let n_ff_max = self
2271            .layers
2272            .iter()
2273            .map(|l| match &l.ffn {
2274                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
2275                _ => h,
2276            })
2277            .max()
2278            .unwrap_or(h)
2279            .max(h);
2280        let f32b = std::mem::size_of::<f32>();
2281        // slab planes: 7 f32 + 2 f16 on n_embd, 3 f32 on n_ff; per-call extras: a16 (f16 ff),
2282        // the f16 pre-norm operand (n_embd), and the x / hn copies (2 f32 n_embd).
2283        let mut call_row_bytes =
2284            (7 + 2) * h * f32b + (2 + 1) * h * 2 + 3 * n_ff_max * f32b + n_ff_max * 2;
2285        // Routed-expert staging per row (the grouped f16 GEMM path in `moe_ffn_inner`): the
2286        // gathered inputs, three projection planes with f16 mirrors, the down output and
2287        // the scatter/partial planes, all sized by the USED expert count and the expert FFN
2288        // width. Same term the hyper shape charges. Calibrated 2026-09-04 on ornith (h 2048,
2289        // 8 of 256 experts, expert ff 512): the dense planes alone priced 120 KB/row against
2290        // a measured 345 KB/row monolithic peak; this term is the missing 221 KB/row.
2291        if let Some(moe) = self.cfg.moe.as_ref() {
2292            let u = moe.expert_used_count as usize;
2293            let f = moe.expert_ff_length as usize;
2294            call_row_bytes += u * (10 * h + 14 * f);
2295        }
2296        Some(PrimeWorkspaceShape {
2297            call_row_bytes,
2298            prompt_row_bytes: h * f32b,
2299            n_layers: self.layers.len(),
2300        })
2301    }
2302
2303    pub fn hyper_prime_workspace_shape(&self) -> Option<HyperPrimeWorkspaceShape> {
2304        let topology = self.hyper.as_ref()?;
2305        let h = self.cfg.n_embd as usize;
2306        let s = topology.streams;
2307        let f32b = std::mem::size_of::<f32>();
2308        // Stream state (x2) + ppN boundary slots (x2), pre/norm transients, prime tail.
2309        let mut chunk_token_bytes = 4 * s * h * f32b + 4 * h * f32b + 2 * h * f32b;
2310        if let Some(moe) = self.cfg.moe.as_ref() {
2311            let u = moe.expert_used_count as usize;
2312            let f = moe.expert_ff_length as usize;
2313            chunk_token_bytes += u * (10 * h + 14 * f) + 4 * h;
2314        }
2315        let mut kpool_score_pool = 0;
2316        if let Some(glm5) = self.cfg.glm5.as_ref() {
2317            let heads = self.cfg.n_head as usize;
2318            chunk_token_bytes +=
2319                heads * (glm5.qk_head_dim as usize + glm5.v_head_dim as usize) * f32b;
2320            if glm5.index_kpool > 0 {
2321                chunk_token_bytes += (glm5.index_topk as usize / glm5.index_kpool as usize + 1)
2322                    * std::mem::size_of::<i32>();
2323                kpool_score_pool = glm5.index_kpool as usize;
2324            }
2325        }
2326        Some(HyperPrimeWorkspaceShape {
2327            chunk_token_bytes,
2328            prompt_bytes_per_token: h * f32b,
2329            kpool_score_pool,
2330            n_layers: self.layers.len(),
2331            gdn_grid: self.gdn_prime_grid_on(),
2332        })
2333    }
2334
2335    /// One T=1 decode step under the mHC residual: `decode_step_h`'s contract over the hc layer
2336    /// program. The stream state is INTRA-STEP — expanded from the embedded row, collapsed for
2337    /// the logits — so no cache format changes and the mixers keep their own state exactly as
2338    /// they do on the serial path.
2339    pub(crate) fn decode_step_hyper(
2340        &self,
2341        e: &Engine,
2342        token: u32,
2343        cache: &mut Cache,
2344    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2345        let topology = *self
2346            .hyper
2347            .as_ref()
2348            .ok_or("decode_step_hyper on a model with no HyperConnections topology")?;
2349        // M2 ppN door for the mHC decode step (see `forward_hyper`'s note). This is the door
2350        // the GLM-5.3-Flash residency arc turns on: with it shut, every routed expert has to
2351        // fit beside card 0's trunk, which 171.2 GB of experts cannot do.
2352        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2353            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
2354                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2355            }
2356            return self.decode_step_hyper_ppn(e, token, cache, &topology, &fence);
2357        }
2358        let n_embd = self.cfg.n_embd as usize;
2359        let eps = self.cfg.rms_eps;
2360        let pos = cache.pos;
2361        let pos_d = e.htod_i32(&[pos as i32])?;
2362
2363        let embedded = e.htod(&self.embd.try_gather(n_embd, &[token])?)?;
2364        let mut x = crate::hyper::expand(e, &topology, &embedded, 1, n_embd)?;
2365
2366        x = self.hyper_range_decode(e, &topology, x, 0, self.layers.len(), &pos_d, pos, cache)?;
2367
2368        // SHARED EXIT with the ppN twin (see `forward_hyper`'s note).
2369        self.hyper_decode_tail(e, &topology, &x, n_embd, eps, cache)
2370    }
2371
2372    /// One hc layer RANGE `[lo, hi)` of the STATELESS prefill walk, driven by engine `e`.
2373    ///
2374    /// Extracted so the unsplit walk and every pipeline stage run the SAME code over their own
2375    /// range: the ppN arm's bit-identity claim is then structural, not a coincidence of two
2376    /// hand-kept-in-sync copies. `x` enters and leaves as the `[t, streams, hidden]` stream
2377    /// state, which is exactly the payload a stage boundary transports.
2378    #[allow(clippy::too_many_arguments)]
2379    fn hyper_range_forward(
2380        &self,
2381        e: &Engine,
2382        topology: &crate::hyper::HyperTopology,
2383        mut x: CudaSlice<f32>,
2384        lo: usize,
2385        hi: usize,
2386        pos_d: &CudaSlice<i32>,
2387        t: usize,
2388        trace: bool,
2389    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2390        let n_embd = self.cfg.n_embd as usize;
2391        let eps = self.cfg.rms_eps;
2392        for il in lo..hi {
2393            let layer = &self.layers[il];
2394            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2395                format!("layer {il} carries no hyper-connection weights under an hc plan")
2396            })?;
2397
2398            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2399            let mut h = e.uninit(t * n_embd)?;
2400            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2401            let mixed = match &layer.mixer {
2402                Mixer::Full(fa) => self.full_attn(e, fa, &h, pos_d, t, il)?,
2403                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
2404                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, pos_d, t, il)?,
2405                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
2406            };
2407            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2408            if trace {
2409                let index = il as i64;
2410                memra_reference::hidden_trace::emit_last_row(
2411                    "mixer",
2412                    index,
2413                    t,
2414                    n_embd,
2415                    &e.dtoh(&mixed)?,
2416                );
2417                let streams = x.len() / (t * n_embd);
2418                memra_reference::hidden_trace::emit_last_row(
2419                    "attn",
2420                    index,
2421                    t,
2422                    streams * n_embd,
2423                    &e.dtoh(&x)?,
2424                );
2425            }
2426
2427            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2428            let mut z = e.uninit(t * n_embd)?;
2429            e.rms_norm(
2430                &y,
2431                layer.post_attn_norm.float_data(),
2432                &mut z,
2433                n_embd,
2434                t,
2435                eps,
2436            )?;
2437            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true, None)?;
2438            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2439            if trace {
2440                let index = il as i64;
2441                memra_reference::hidden_trace::emit_last_row(
2442                    "ffn",
2443                    index,
2444                    t,
2445                    n_embd,
2446                    &e.dtoh(&ffn_out)?,
2447                );
2448                let streams = x.len() / (t * n_embd);
2449                memra_reference::hidden_trace::emit_last_row(
2450                    "layer",
2451                    index,
2452                    t,
2453                    streams * n_embd,
2454                    &e.dtoh(&x)?,
2455                );
2456            }
2457        }
2458        Ok(x)
2459    }
2460
2461    /// One hc layer RANGE `[lo, hi)` of the STATEFUL prime walk (see `hyper_range_forward`).
2462    /// Every mixer writes its own layer's cache state through `e`, so under the ppN door a
2463    /// stage's KDA conv ring / delta-rule state and its MLA latent rows + kpool indexer plane
2464    /// are written by the SAME engine `pp::new_cache` allocated them on.
2465    #[allow(clippy::too_many_arguments)]
2466    #[allow(clippy::too_many_arguments)] // allow: the list is the walk contract plus the tap arm
2467    fn hyper_range_prime(
2468        &self,
2469        e: &Engine,
2470        topology: &crate::hyper::HyperTopology,
2471        mut x: CudaSlice<f32>,
2472        lo: usize,
2473        hi: usize,
2474        pos_d: &CudaSlice<i32>,
2475        t: usize,
2476        cache: &mut Cache,
2477        seq_end: usize,
2478        taps: HcTapArm<'_, '_>,
2479    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2480        let n_embd = self.cfg.n_embd as usize;
2481        let eps = self.cfg.rms_eps;
2482        // MEMRA_PRIME_PROF=1: per-phase wall inside the mHC prime walk, sync-bounded (each mark
2483        // drains the stage stream, so absolute time inflates -- the SPLIT is the signal).
2484        //
2485        // WHY IT LIVES HERE TOO. The pre-existing `[prime-prof]` marks are in
2486        // `step35_prime_batch_layers`, the step37 BATCHED prime, which refuses on the first
2487        // non-`Mixer::Full` layer -- glm5_next never enters it, so `MEMRA_PRIME_PROF=1` on a
2488        // glm5 box printed nothing at all for the walk that actually ran. The four slots are
2489        // this walk's four phase classes, in issue order: the attention-site hc glue and norm,
2490        // the mixer (KDA scan / MLA+DSA), the FFN-site hc glue and norm, and the FFN itself
2491        // (grouped MoE, or the dense MLP on layers 0..first_k_dense). The per-layer MoE
2492        // breakdown is the `[moe-grouped-prefill-prof]` line the same flag already prints.
2493        let prof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
2494        let mut ph = [0f64; 4]; // 0 attn-site glue+norm, 1 mixer, 2 ffn-site glue+norm, 3 ffn
2495        let mut pt = std::time::Instant::now();
2496        let mark = |e: &Engine, acc: usize, pt: &mut std::time::Instant, ph: &mut [f64; 4]| {
2497            if prof {
2498                let _ = e.stream().synchronize();
2499                ph[acc] += pt.elapsed().as_secs_f64() * 1e3;
2500                *pt = std::time::Instant::now();
2501            }
2502        };
2503        for il in lo..hi {
2504            let layer = &self.layers[il];
2505            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2506                format!("layer {il} carries no hyper-connection weights under an hc plan")
2507            })?;
2508
2509            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2510            let mut h = e.uninit(t * n_embd)?;
2511            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2512            mark(e, 0, &mut pt, &mut ph);
2513            let mixed = match &layer.mixer {
2514                Mixer::Full(fa) => {
2515                    self.full_attn_prime(e, fa, &h, None, pos_d, t, cache, il, seq_end)?
2516                }
2517                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2518                Mixer::Mla(mla) if mla.tp.is_some() => {
2519                    self.mla_tp_attn_cached(e, mla, &h, pos_d, t, il, cache, false)?
2520                }
2521                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, t, il, cache)?,
2522                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2523                    e,
2524                    la,
2525                    &h,
2526                    t,
2527                    eps,
2528                    cache,
2529                    il,
2530                    crate::kda::ConvArm::Prefill,
2531                )?,
2532                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2533            };
2534            mark(e, 1, &mut pt, &mut ph);
2535            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2536
2537            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2538            let mut z = e.uninit(t * n_embd)?;
2539            e.rms_norm(
2540                &y,
2541                layer.post_attn_norm.float_data(),
2542                &mut z,
2543                n_embd,
2544                t,
2545                eps,
2546            )?;
2547            mark(e, 2, &mut pt, &mut ph);
2548            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true, None)?;
2549            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2550            // glm5 DFlash2 feature tap (lane/glm5-dflash-draft-src): the CONTRACTED
2551            // completed layer output, host-staged — one Option check when unarmed.
2552            match &taps {
2553                HcTapArm::FromCache => self.glm5_hc_tap(e, cache, topology, il, &x, t)?,
2554                HcTapArm::Shared(sink, base) => {
2555                    // Lock held for the row copy only (see `HcTapArm`'s header for why sharing
2556                    // is sound); the contraction and the readback inside are this stage's own.
2557                    let mut guard = sink
2558                        .lock()
2559                        .map_err(|_| "hc tap sink lock poisoned by a failed stage walk")?;
2560                    self.glm5_hc_tap_into(e, &mut guard, *base, topology, il, &x, t)?;
2561                }
2562            }
2563            mark(e, 3, &mut pt, &mut ph);
2564        }
2565        if prof {
2566            // Same shape as the step35 line (walk, token count, layer count, four ms fields);
2567            // the KEYS name this walk's phases rather than step37's, because a line labelled
2568            // `norm+qkv` for an mHC site would be a lie that reads like a measurement.
2569            eprintln!(
2570                "[prime-prof] walk=hyper t={t} layers={} attn_glue={:.0}ms mixer={:.0}ms \
2571                 ffn_glue={:.0}ms ffn={:.0}ms",
2572                hi - lo,
2573                ph[0],
2574                ph[1],
2575                ph[2],
2576                ph[3]
2577            );
2578        }
2579        Ok(x)
2580    }
2581
2582    /// One hc layer RANGE `[lo, hi)` of the T=1 decode step (see `hyper_range_forward`).
2583    #[allow(clippy::too_many_arguments)]
2584    fn hyper_range_decode(
2585        &self,
2586        e: &Engine,
2587        topology: &crate::hyper::HyperTopology,
2588        x: CudaSlice<f32>,
2589        lo: usize,
2590        hi: usize,
2591        pos_d: &CudaSlice<i32>,
2592        pos: usize,
2593        cache: &mut Cache,
2594    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2595        // DOOR `MEMRA_GLM5_DECODE_GRAPH` (lane/b200-glm5-graph-20260902, default ON since
2596        // 2026-09-04, `=0` disarms, read PER CALL): replay this stage's captured KDA-layer runs
2597        // instead of issuing them. Every
2598        // refusal shape returns None and falls through to the eager walk below, byte-identically
2599        // — see crates/memra-engine/src/glm5_decode_graph.rs for what is captured and why the
2600        // MLA/DSA layers are not.
2601        if crate::glm5_decode_graph_on() && self.glm5_decode_graph_ready(e, cache, lo, hi) {
2602            return self.hyper_range_decode_graphed(e, topology, x, lo, hi, pos_d, pos, cache);
2603        }
2604        // MEMRA_GLM5_GRAPH_TRACE (gate harness): split the eager walk at the SAME run boundaries
2605        // the graph arm replays at, so the two arms' checksums line up segment for segment. The
2606        // split is a loop split, not a program change — the same kernels in the same order.
2607        if crate::glm5_graph_trace_on() {
2608            return self.hyper_range_decode_eager_traced(e, topology, x, lo, hi, pos_d, pos, cache);
2609        }
2610        self.hyper_range_decode_eager(e, topology, x, lo, hi, pos_d, pos, cache)
2611    }
2612
2613    /// The eager T=1 hc decode walk over `[lo, hi)` — the program every other arm is measured
2614    /// against, and the fall-through of the decode-graph door above.
2615    #[allow(clippy::too_many_arguments)]
2616    pub(crate) fn hyper_range_decode_eager(
2617        &self,
2618        e: &Engine,
2619        topology: &crate::hyper::HyperTopology,
2620        mut x: CudaSlice<f32>,
2621        lo: usize,
2622        hi: usize,
2623        pos_d: &CudaSlice<i32>,
2624        pos: usize,
2625        cache: &mut Cache,
2626    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2627        // MEMRA_HC_DECODE_WS=1 (default OFF, read per call — rollback seam): the persistent-
2628        // workspace twin of this walk. Same kernels, same call order, same operand bytes;
2629        // only the hc-glue allocations (mixes/gates/comb/y/h/z/post-out, ~12 alloc+free pairs
2630        // per layer per token of the census's 2,358) disappear. Byte identity ON/OFF is gated
2631        // by hc_decode_ws_gpu.rs; refusal shapes fall through to the allocating walk below.
2632        if hyper_decode_ws_on() {
2633            return self.hyper_range_decode_ws(e, topology, x, lo, hi, pos_d, pos, cache);
2634        }
2635        let n_embd = self.cfg.n_embd as usize;
2636        let eps = self.cfg.rms_eps;
2637        for il in lo..hi {
2638            let layer = &self.layers[il];
2639            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2640                format!("layer {il} carries no hyper-connection weights under an hc plan")
2641            })?;
2642
2643            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, 1, n_embd)?;
2644            let mut h = e.uninit(n_embd)?;
2645            // MEMRA_GLM5_Q8_FUSE_ATTN (lane/glm5-attn-norm-zq8-20260904, default OFF): on a plain
2646            // KDA layer the attention-input norm emits its q8_1 view too, and the fused
2647            // six-projection launcher reads it instead of quantizing `h` again. Byte-identical
2648            // (rms_norm_zq8_f32 = rms_norm then quantize_q8_1; the launcher's quantize is the
2649            // same kernel). Every other mixer keeps the plain norm.
2650            let attn_q8 = if crate::glm5_q8_fuse_attn_on()
2651                && matches!(&layer.mixer, Mixer::Kda(la) if la.tp.is_none())
2652            {
2653                let pair =
2654                    e.rms_norm_zq8_f32(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
2655                if crate::GLM5_Q8_FUSE_ATTN_DISPATCHES
2656                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
2657                    == 0
2658                {
2659                    eprintln!(
2660                        "[glm5-q8-fuse-attn] engaged (rms_norm_zq8_f32 -> kda6, hyper_range_decode)"
2661                    );
2662                }
2663                Some(pair)
2664            } else {
2665                e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
2666                None
2667            };
2668            let mixed = match &layer.mixer {
2669                Mixer::Full(fa) => self.full_attn_decode(e, fa, &h, pos_d, pos, cache, il)?,
2670                Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il)?,
2671                Mixer::Mla(mla) if mla.tp.is_some() => {
2672                    self.mla_tp_attn_cached(e, mla, &h, pos_d, 1, il, cache, false)?
2673                }
2674                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache)?,
2675                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2676                    e,
2677                    la,
2678                    &h,
2679                    1,
2680                    eps,
2681                    cache,
2682                    il,
2683                    crate::kda::ConvArm::Decode,
2684                )?,
2685                Mixer::Kda(la) => crate::kda::kda_decode_cached_q8(
2686                    e,
2687                    la,
2688                    &h,
2689                    attn_q8.as_ref().map(|(q, d)| (q, d)),
2690                    eps,
2691                    cache,
2692                    il,
2693                )?,
2694            };
2695            x = crate::hyper::post(e, topology, &mixed, &x, &mix, 1, n_embd)?;
2696
2697            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, 1, n_embd)?;
2698            let mut z = e.uninit(n_embd)?;
2699            // Door MEMRA_GLM5_Q8_FUSE (lane/b200-q8-fuse-20260902, default OFF): fold the
2700            // FFN-input rms_norm and its consumer's standalone quantize_q8_1 launch into
2701            // ONE kernel. z (f32) is still needed by the router/shexp branches inside
2702            // hyper_ffn_branch, so the fused kernel emits both views; moe_ffn_il_zq8 takes
2703            // the pre-quantized pair when present and re-quantizes z itself when not —
2704            // byte-identical either way (rms_norm_zq8_f32's header carries the identity
2705            // argument against rms_norm then quantize_q8_1).
2706            let zq8 = if crate::glm5_q8_fuse_on() {
2707                let pair = e.rms_norm_zq8_f32(
2708                    &y,
2709                    layer.post_attn_norm.float_data(),
2710                    &mut z,
2711                    n_embd,
2712                    1,
2713                    eps,
2714                )?;
2715                if GLM5_Q8_FUSE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2716                    eprintln!("[glm5-q8-fuse] engaged (rms_norm_zq8_f32, hyper_range_decode)");
2717                }
2718                Some(pair)
2719            } else {
2720                e.rms_norm(
2721                    &y,
2722                    layer.post_attn_norm.float_data(),
2723                    &mut z,
2724                    n_embd,
2725                    1,
2726                    eps,
2727                )?;
2728                None
2729            };
2730            let ffn_out = self.hyper_ffn_branch(e, layer, &z, 1, il, false, zq8.as_ref())?;
2731            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, 1, n_embd)?;
2732        }
2733        Ok(x)
2734    }
2735
2736    /// The persistent-workspace twin of `hyper_range_decode` (MEMRA_HC_DECODE_WS, lever 2 of
2737    /// the decode diet). One `HyperDecodeWs` per engine (so each ppN stage owns its own,
2738    /// allocated on its own device); the walk TAKES it from the engine pool, rotates the
2739    /// stream state against `ws.xb` (an ownership swap, not a copy), and puts it back. The
2740    /// mixers and the FFN/MoE branches are the SAME calls with the SAME inputs — their
2741    /// internal allocations are untouched by this lever.
2742    #[allow(clippy::too_many_arguments)]
2743    fn hyper_range_decode_ws(
2744        &self,
2745        e: &Engine,
2746        topology: &crate::hyper::HyperTopology,
2747        x: CudaSlice<f32>,
2748        lo: usize,
2749        hi: usize,
2750        pos_d: &CudaSlice<i32>,
2751        pos: usize,
2752        cache: &mut Cache,
2753    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2754        let n_embd = self.cfg.n_embd as usize;
2755        let mut ws = match e.hyper_ws_take() {
2756            Some(ws) if ws.matches(topology, n_embd) => ws,
2757            _ => crate::hyper::HyperDecodeWs::new(e, topology, n_embd)?,
2758        };
2759        if HC_DECODE_WS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2760            eprintln!(
2761                "[hc-decode-ws] engaged streams={} hidden={n_embd} (persistent hc-glue \
2762                 workspace, per-engine pool; MEMRA_HC_DECODE_WS=1)",
2763                topology.streams
2764            );
2765        }
2766        let mut x = x;
2767        let out = self
2768            .hyper_range_decode_ws_body(e, topology, &mut x, lo, hi, pos_d, pos, cache, &mut ws)
2769            .map(|()| x);
2770        e.hyper_ws_put(ws);
2771        out
2772    }
2773
2774    /// The walk itself — `hyper_range_decode`'s loop with the hc glue landing in `ws`.
2775    /// KEPT CALL-FOR-CALL IN STEP with the allocating walk above: same kernels, same order
2776    /// (pre -> rms_norm -> mixer -> post -> pre -> rms_norm -> ffn -> post), so the
2777    /// byte-identity gate is a structural claim, not a coincidence.
2778    #[allow(clippy::too_many_arguments)]
2779    pub(crate) fn hyper_range_decode_ws_body(
2780        &self,
2781        e: &Engine,
2782        topology: &crate::hyper::HyperTopology,
2783        x: &mut CudaSlice<f32>,
2784        lo: usize,
2785        hi: usize,
2786        pos_d: &CudaSlice<i32>,
2787        pos: usize,
2788        cache: &mut Cache,
2789        ws: &mut crate::hyper::HyperDecodeWs,
2790    ) -> Result<(), Box<dyn std::error::Error>> {
2791        let n_embd = self.cfg.n_embd as usize;
2792        let eps = self.cfg.rms_eps;
2793        for il in lo..hi {
2794            let layer = &self.layers[il];
2795            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2796                format!("layer {il} carries no hyper-connection weights under an hc plan")
2797            })?;
2798
2799            crate::hyper::pre_t1_ws(e, topology, &hyper.attn, x, ws, n_embd)?;
2800            // MEMRA_GLM5_Q8_FUSE_ATTN: see the eager walk above; same fusion, workspace form.
2801            let attn_q8 = if crate::glm5_q8_fuse_attn_on()
2802                && matches!(&layer.mixer, Mixer::Kda(la) if la.tp.is_none())
2803            {
2804                let pair = e.rms_norm_zq8_f32(
2805                    &ws.y,
2806                    layer.attn_norm.float_data(),
2807                    &mut ws.h,
2808                    n_embd,
2809                    1,
2810                    eps,
2811                )?;
2812                if crate::GLM5_Q8_FUSE_ATTN_DISPATCHES
2813                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
2814                    == 0
2815                {
2816                    eprintln!(
2817                        "[glm5-q8-fuse-attn] engaged (rms_norm_zq8_f32 -> kda6, hyper_range_decode_ws_body)"
2818                    );
2819                }
2820                Some(pair)
2821            } else {
2822                e.rms_norm(
2823                    &ws.y,
2824                    layer.attn_norm.float_data(),
2825                    &mut ws.h,
2826                    n_embd,
2827                    1,
2828                    eps,
2829                )?;
2830                None
2831            };
2832            let mixed = match &layer.mixer {
2833                Mixer::Full(fa) => self.full_attn_decode(e, fa, &ws.h, pos_d, pos, cache, il)?,
2834                Mixer::Linear(la) => self.linear_attn_decode(e, la, &ws.h, cache, il)?,
2835                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &ws.h, pos_d, 1, il, cache)?,
2836                Mixer::Kda(la) => crate::kda::kda_decode_cached_q8(
2837                    e,
2838                    la,
2839                    &ws.h,
2840                    attn_q8.as_ref().map(|(q, d)| (q, d)),
2841                    eps,
2842                    cache,
2843                    il,
2844                )?,
2845            };
2846            crate::hyper::post_t1_ws(e, topology, &mixed, x, ws, n_embd)?;
2847            std::mem::swap(x, &mut ws.xb);
2848
2849            crate::hyper::pre_t1_ws(e, topology, &hyper.mlp, x, ws, n_embd)?;
2850            // Door MEMRA_GLM5_Q8_FUSE — the workspace twin of the fusion above (same fused
2851            // kernel, same byte-identity argument, `ws.z` in place of a fresh allocation).
2852            let zq8 = if crate::glm5_q8_fuse_on() {
2853                let pair = e.rms_norm_zq8_f32(
2854                    &ws.y,
2855                    layer.post_attn_norm.float_data(),
2856                    &mut ws.z,
2857                    n_embd,
2858                    1,
2859                    eps,
2860                )?;
2861                if GLM5_Q8_FUSE_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2862                    eprintln!(
2863                        "[glm5-q8-fuse] engaged (rms_norm_zq8_f32, hyper_range_decode_ws_body)"
2864                    );
2865                }
2866                Some(pair)
2867            } else {
2868                e.rms_norm(
2869                    &ws.y,
2870                    layer.post_attn_norm.float_data(),
2871                    &mut ws.z,
2872                    n_embd,
2873                    1,
2874                    eps,
2875                )?;
2876                None
2877            };
2878            let ffn_out = self.hyper_ffn_branch(e, layer, &ws.z, 1, il, false, zq8.as_ref())?;
2879            crate::hyper::post_t1_ws(e, topology, &ffn_out, x, ws, n_embd)?;
2880            std::mem::swap(x, &mut ws.xb);
2881        }
2882        Ok(())
2883    }
2884
2885    /// One hc layer RANGE `[lo, hi)` of the BATCHED T=1 decode step: B independent sessions
2886    /// share one walk over the `[B, streams, n_embd]` stream state. The batched twin of
2887    /// `hyper_range_decode`, and the trunk of `decode_step_batch_hyper` (decode_batch.rs).
2888    ///
2889    /// SHAPE — batched where the arithmetic is row-independent, per-session where the state
2890    /// is, decode-exact where a reduction is width-dependent:
2891    ///
2892    ///   * The hc glue (`expand`/`pre_finish` kernels/`post`) is block-per-token by
2893    ///     construction (grid over t), so t=B batches it with per-row bytes unchanged.
2894    ///   * The hc mixing GEMM is the ONE width-dependent reduction in the glue
2895    ///     (cuBLASLt's n-dependent split — the lt_ndep probe), so this walk calls
2896    ///     `hyper::pre_exact`, which runs each row through the m=1 program the serial step
2897    ///     runs. rms_norm at m=B is a per-row program.
2898    ///   * The MIXERS (KDA conv ring + delta rule, MLA latent rows + kpool indexer plane,
2899    ///     and the Full/Linear classes for completeness) hold per-session recurrent or
2900    ///     latent state, so each session's row is routed to ITS OWN cache through the SAME
2901    ///     t=1 call its solo step makes — the per-seq loop is the v1 exactness doctrine
2902    ///     from this module's sibling (`decode_batch.rs` header), and the row copies in and
2903    ///     out are arithmetic-free materializations.
2904    ///   * The FFN batches at t=B: the MoE body's router is the fixed per-row program at
2905    ///     t < PRIME_MIN_T, expert dispatch is per-token, and the shexp trio rides the
2906    ///     per-column decode-exact arm at decode widths; the dense branch runs per-row so
2907    ///     each row executes the serial `hyper_ffn_branch` program verbatim.
2908    ///
2909    /// EXACTNESS BAR: row b of a B-row step must be BIT-IDENTICAL to session b decoding
2910    /// alone through `decode_step_hyper` — full-logit compare, per step. Gate:
2911    /// `glm5-hyper-batch-gate` (fixture-driven, red-armed with a swapped-row and a
2912    /// wrong-cache-slot mutation; receipts in
2913    /// `research/glm53-flash-bringup-20260827/batched-decode-gate/`).
2914    ///
2915    /// `pos_rows[bi]` is session bi's single-position device buffer, uploaded by the caller
2916    /// through THIS range's engine (the per-stage pos_d law under a pp split). `caches[bi]`
2917    /// advances exactly as its solo step would; `cache.pos` itself is bumped by the caller's
2918    /// epilogue, never here.
2919    #[allow(clippy::too_many_arguments)]
2920    pub(crate) fn hyper_batch_range_decode(
2921        &self,
2922        e: &Engine,
2923        topology: &crate::hyper::HyperTopology,
2924        mut x: CudaSlice<f32>,
2925        lo: usize,
2926        hi: usize,
2927        pos_rows: &[CudaSlice<i32>],
2928        caches: &mut [&mut Cache],
2929    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2930        let b_n = caches.len();
2931        assert_eq!(
2932            pos_rows.len(),
2933            b_n,
2934            "hyper_batch_range_decode: pos_rows built for a different batch width"
2935        );
2936        // MEMRA_HYPER_BATCH_SOLO=1 (default OFF, read per call — the rollback seam is its
2937        // absence). At B=1 this walk allocates `h` and `mixed`, then allocates a per-row
2938        // `h_row` and `dtod_copy_view`s the single row out of `h`, and then calls the SAME
2939        // t=1 mixer program `hyper_range_decode` calls. At one row the copy and the extra
2940        // allocation are pure overhead.
2941        //
2942        // THE REASON THIS MATTERS MORE THAN THAT OVERHEAD: this walk is the only hc decode
2943        // walk glm5 PP-N serving reaches (`decode_step_batch_hyper` -> `_ppn` -> here), and
2944        // it is the one hc walk with NO allocation-workspace door and NO decode-graph door.
2945        // Both of those guard `hyper_range_decode_eager`, which is reachable only through
2946        // `hyper_range_decode`. Measured 2026-09-03 on the 2x B200 pair: `MEMRA_HC_DECODE_WS=1`
2947        // moved serving 55.85 -> 55.96 tok/s (noise) and its once-per-boot `[hc-decode-ws]
2948        // engaged` line printed ZERO times, because the walk was never entered. Delegating at
2949        // B=1 puts serving on the walk both diets already gate.
2950        //
2951        // BYTE IDENTITY IS THIS WALK'S OWN CONTRACT, not a new claim: the header above states
2952        // "row b of a B-row step must be BIT-IDENTICAL to session b decoding alone through
2953        // `decode_step_hyper` — full-logit compare, per step", red-armed by
2954        // `glm5-hyper-batch-gate` with swapped-row and wrong-cache-slot mutations. At B=1 the
2955        // delegation IS that solo path, so the two arms are equal by the contract the batch
2956        // walk is already gated against.
2957        //
2958        // `cache.pos` is untouched here exactly as before ("`cache.pos` itself is bumped by the
2959        // caller's epilogue, never here"): the solo walk does not advance it either.
2960        if b_n == 1 && crate::hyper_batch_solo_on() {
2961            let pos = caches[0].pos;
2962            return self.hyper_range_decode(e, topology, x, lo, hi, &pos_rows[0], pos, caches[0]);
2963        }
2964        let n_embd = self.cfg.n_embd as usize;
2965        let eps = self.cfg.rms_eps;
2966        for il in lo..hi {
2967            let layer = &self.layers[il];
2968            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2969                format!("layer {il} carries no hyper-connection weights under an hc plan")
2970            })?;
2971
2972            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, b_n, n_embd)?;
2973            let mut h = e.uninit(b_n * n_embd)?;
2974            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, b_n, eps)?;
2975            // ---- mixers: per-session, each row through its OWN cache (the t=1 serial
2976            // program — cross-session contamination here is the failure mode the gate's
2977            // swapped-row mutation exists to catch) ----
2978            let mut mixed = e.uninit(b_n * n_embd)?;
2979            for bi in 0..b_n {
2980                let mut h_row = e.uninit(n_embd)?;
2981                e.dtod_copy_view(&h.slice(bi * n_embd..(bi + 1) * n_embd), &mut h_row)?;
2982                let cache: &mut Cache = &mut *caches[bi];
2983                let pos = cache.pos;
2984                let out_row = match &layer.mixer {
2985                    Mixer::Full(fa) => {
2986                        self.full_attn_decode(e, fa, &h_row, &pos_rows[bi], pos, cache, il)?
2987                    }
2988                    Mixer::Linear(la) => self.linear_attn_decode(e, la, &h_row, cache, il)?,
2989                    Mixer::Mla(mla) => {
2990                        self.mla_attn_cached(e, mla, &h_row, &pos_rows[bi], 1, il, cache)?
2991                    }
2992                    Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?,
2993                };
2994                e.copy_into(&mut mixed, bi * n_embd, &out_row, n_embd)?;
2995            }
2996            x = crate::hyper::post(e, topology, &mixed, &x, &mix, b_n, n_embd)?;
2997
2998            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, b_n, n_embd)?;
2999            let mut z = e.uninit(b_n * n_embd)?;
3000            e.rms_norm(
3001                &y,
3002                layer.post_attn_norm.float_data(),
3003                &mut z,
3004                n_embd,
3005                b_n,
3006                eps,
3007            )?;
3008            let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, b_n, il, false)?;
3009            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, b_n, n_embd)?;
3010        }
3011        Ok(x)
3012    }
3013
3014    /// The FFN branch of the batched hc decode walk (see `hyper_batch_range_decode`).
3015    ///
3016    /// MoE batches at t=B — that is the weight win this walk exists for (one stream of the
3017    /// routed experts serves B rows), and every stage of `moe_ffn_il_zq8` is per-row exact
3018    /// at decode widths (router: fixed per-row program at t < PRIME_MIN_T; experts:
3019    /// per-(token,expert) programs; shexp: the per-column decode-exact arm). The DENSE
3020    /// branch runs PER ROW through the serial `hyper_ffn_branch` instead: its
3021    /// `matmul_group` dispatch carries no per-row bit-identity contract across widths for
3022    /// every weight class this walk must serve, and a first-k-dense plan carries one such
3023    /// layer — per-row costs nothing and each row executes the solo step's program verbatim.
3024    /// `vrows` (lane/glm5-vrest): the VERIFY walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`)
3025    /// passes `true`, which lets the MoE body take the pairs-shaped batched routed-expert
3026    /// program across the t rows (`moe_vrows_pairs_q8` — bit-identical per row, fail-closed
3027    /// to the sequential loop for every unqualified shape). The batched DECODE walk
3028    /// (`decode_step_batch_hyper`) passes `false` — its priced dispatch class stays
3029    /// byte-stable; porting it is a named follow-up with its own re-price. The DENSE branch
3030    /// is per-row in both arms (its `matmul_group` dispatch carries no cross-width per-row
3031    /// bit-identity contract for every weight class this walk must serve; ~3 layers, named
3032    /// out of scope in the vrest attribution).
3033    pub(crate) fn hyper_ffn_branch_batch(
3034        &self,
3035        e: &Engine,
3036        layer: &crate::hybrid::HybridLayer,
3037        z: &CudaSlice<f32>,
3038        b_n: usize,
3039        il: usize,
3040        vrows: bool,
3041    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3042        let n_embd = self.cfg.n_embd as usize;
3043        match &layer.ffn {
3044            crate::hybrid::Ffn::Dense { .. } => {
3045                let mut out = e.uninit(b_n * n_embd)?;
3046                for bi in 0..b_n {
3047                    let mut z_row = e.uninit(n_embd)?;
3048                    e.dtod_copy_view(&z.slice(bi * n_embd..(bi + 1) * n_embd), &mut z_row)?;
3049                    let row = self.hyper_ffn_branch(e, layer, &z_row, 1, il, false, None)?;
3050                    e.copy_into(&mut out, bi * n_embd, &row, n_embd)?;
3051                }
3052                Ok(out)
3053            }
3054            crate::hybrid::Ffn::Moe(m) => {
3055                if vrows {
3056                    self.moe_ffn_il_zq8_vrows(e, m, z, b_n, il as u16)
3057                } else {
3058                    self.moe_ffn_il_zq8(e, m, z, None, b_n, il as u16)
3059                }
3060            }
3061        }
3062    }
3063
3064    // =============================== M2 ppN, mHC arm ===============================
3065    //
3066    // The three walks below are the hc twins of `decode_step_h_ppn`. They exist because the
3067    // GLM-5.3-Flash residency arithmetic does not close on one card: 171.2 GB of routed
3068    // experts against 2x96 GB means the second card is the only route to full residency, and
3069    // the pp door is how weights get there. Until these landed, all three hc walks refused
3070    // the door outright ("the sharded stage handoff is unwired for this residual topology"),
3071    // which is a loud refusal and was the right behaviour — a single-engine walk over
3072    // stage-sharded weights dereferences another device's pointers.
3073    //
3074    // WHAT IS DIFFERENT FROM THE GENERIC ARM, and it is exactly one thing: the payload on the
3075    // wire. The serial trunk hands `[n_embd]` (decode) or `[t, n_embd]` (prime) across a
3076    // boundary; the mHC trunk carries `streams` residual streams between layers, so the
3077    // boundary payload is `[streams, n_embd]` / `[t, streams, n_embd]`. `pp.rs`'s BoundarySlot
3078    // buffers are lazily sized from the caller's `n` and grow to the high-water mark, so no
3079    // slot-sizing change was needed for that — the wider payload just makes them wider.
3080    // `hyper::expand` runs on stage 0 (it takes no weights) and `hyper::collapse` +
3081    // output_norm + lm head on the last stage, which is where the loader already put the head
3082    // (`pp::layer_engine(e, n_trunk, n_trunk - 1)` in `hybrid.rs`) and, under
3083    // `HcCollapse::GatedHead`, the head trio.
3084    //
3085    // Per-layer state placement needed NO new contract: `pp::new_cache` already picks the
3086    // owning stage's `KvDev` per layer for all three of glm5_next's state classes
3087    // (`Recurrent` = KDA conv ring + delta-rule state, `LatentKvCache` = MLA rows + the kpool
3088    // indexer plane, `KvCache` = full attention), and the kpool `index_pool_keys` plane is
3089    // lazily allocated through the engine the mixer is called with, which under these walks is
3090    // the stage's engine. `glm5-hyper-ppn-gate` asserts the fence actually separates those
3091    // classes across stages, so that is a tested property rather than an argued one.
3092    //
3093    // NOT WIRED, and refused rather than approximated: the deferred-readback (pipelined) arm.
3094    // `decode_step_h_ppn_deferred` calls `refuse_hyper`, and this lane did not change that.
3095    //
3096    // Gate: `glm5-hyper-ppn-gate` (bit-identical logits vs the unsplit hc walk, decode and
3097    // prime, at every N/knob combination), receipts in
3098    // `research/glm53-flash-bringup-20260827/ppn-hyper-gate/`.
3099
3100    /// ppN twin of `forward_hyper`: the stateless prefill as N stage subgraphs.
3101    fn forward_hyper_ppn(
3102        &self,
3103        e: &Engine,
3104        tokens: &[u32],
3105        last_only: bool,
3106        topology: &crate::hyper::HyperTopology,
3107        fence: &[usize],
3108    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3109        let n_embd = self.cfg.n_embd as usize;
3110        let eps = self.cfg.rms_eps;
3111        let t = tokens.len();
3112        let width = topology.streams * n_embd;
3113        let trace = memra_reference::hidden_trace::enabled();
3114        if trace {
3115            memra_reference::hidden_trace::emit_tokens(tokens);
3116        }
3117        let pos: Vec<i32> = (0..t as i32).collect();
3118
3119        if crate::pp::pp2_streams_off() {
3120            // Same-stream rollback seam: one engine, one stream, an explicit copy pair per
3121            // boundary. Structurally identical to the split walk, which is the point — it is
3122            // the arm that says "the split is the split, not the streams".
3123            let pos_d = e.htod_i32(&pos)?;
3124            let embedded = self.embed(e, tokens)?;
3125            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
3126            x = self.hyper_range_forward(e, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
3127            for s in 1..fence.len() - 1 {
3128                let boundary_tx = e.clone_dtod(&x)?;
3129                let boundary_rx = e.clone_dtod(&boundary_tx)?;
3130                x = self.hyper_range_forward(
3131                    e,
3132                    topology,
3133                    boundary_rx,
3134                    fence[s],
3135                    fence[s + 1],
3136                    &pos_d,
3137                    t,
3138                    trace,
3139                )?;
3140            }
3141            return self.hyper_head_logits(e, topology, &x, t, n_embd, eps, last_only);
3142        }
3143
3144        let rt = crate::pp::PpNRt::get(e)?;
3145        let n_st = fence.len() - 1;
3146        assert_eq!(
3147            rt.n_stages(),
3148            n_st,
3149            "PpNRt stage count {} != fence stages {n_st}",
3150            rt.n_stages()
3151        );
3152        // #87 REVERSE PUBLICATION: order every stage stream behind the caller's stream before
3153        // the first stage allocation (anatomy: `PpNRt::fence_stages_behind`).
3154        rt.fence_stages_behind(&e.stream())?;
3155
3156        let mut slot = {
3157            let _st0 = rt.enter(0);
3158            let e0 = rt.engine(0, e);
3159            // PER-STAGE pos_d (M2 pipelining law): each stage uploads its OWN copy on ITS
3160            // stream, so the buffer is allocated, consumed and freed on one stream.
3161            let pos_d = e0.htod_i32(&pos)?;
3162            let embedded = self.embed(e0, tokens)?;
3163            let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
3164            let x =
3165                self.hyper_range_forward(e0, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
3166            rt.tx(0, &x, t * width)?
3167        };
3168        for s in 1..n_st - 1 {
3169            let _st = rt.enter(s);
3170            let es = rt.engine(s, e);
3171            let pos_d = es.htod_i32(&pos)?;
3172            let x = rt.rx(s - 1, slot, t * width)?;
3173            let x = self.hyper_range_forward(
3174                es,
3175                topology,
3176                x,
3177                fence[s],
3178                fence[s + 1],
3179                &pos_d,
3180                t,
3181                trace,
3182            )?;
3183            slot = rt.tx(s, &x, t * width)?;
3184        }
3185        let _stl = rt.enter(n_st - 1);
3186        let el = rt.engine(n_st - 1, e);
3187        let pos_d = el.htod_i32(&pos)?;
3188        let x = rt.rx(n_st - 2, slot, t * width)?;
3189        let x = self.hyper_range_forward(
3190            el,
3191            topology,
3192            x,
3193            fence[n_st - 1],
3194            fence[n_st],
3195            &pos_d,
3196            t,
3197            trace,
3198        )?;
3199        self.hyper_head_logits(el, topology, &x, t, n_embd, eps, last_only)
3200    }
3201
3202    /// Trunk exit shared by `forward_hyper` and its ppN twin: collapse the stream state, apply
3203    /// `output_norm`, and project. Runs on the LAST stage's engine under the pp door, which is
3204    /// where `hybrid.rs` uploaded `output_norm`, the lm head and (under `HcCollapse::GatedHead`)
3205    /// the head trio.
3206    #[allow(clippy::too_many_arguments)]
3207    fn hyper_head_logits(
3208        &self,
3209        e: &Engine,
3210        topology: &crate::hyper::HyperTopology,
3211        x: &CudaSlice<f32>,
3212        t: usize,
3213        n_embd: usize,
3214        eps: f32,
3215        last_only: bool,
3216    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3217        let collapsed =
3218            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
3219        if memra_reference::hidden_trace::enabled() {
3220            memra_reference::hidden_trace::emit_last_row(
3221                "collapse",
3222                -1,
3223                t,
3224                n_embd,
3225                &e.dtoh(&collapsed)?,
3226            );
3227        }
3228        let mut hn = e.uninit(t * n_embd)?;
3229        e.rms_norm(
3230            &collapsed,
3231            self.output_norm.float_data(),
3232            &mut hn,
3233            n_embd,
3234            t,
3235            eps,
3236        )?;
3237        let logits = if last_only {
3238            let last = e.view(&hn, t * n_embd);
3239            let last_row = last.slice((t - 1) * n_embd..t * n_embd);
3240            let mut hlast = e.uninit(n_embd)?;
3241            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3242            e.matmul(&self.output, &hlast, 1)?
3243        } else {
3244            e.matmul(&self.output, &hn, t)?
3245        };
3246        e.dtoh(&logits)
3247    }
3248
3249    /// ppN twin of `prime_cache_hyper`: the monolithic stateful prime as N stage subgraphs.
3250    /// The returned device buffers (`h_seed`, `hiddens`) are owned by the LAST stage's engine,
3251    /// the same contract `decode_step_h_ppn` publishes for its `h_seed`.
3252    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3253    #[allow(clippy::too_many_arguments)]
3254    fn prime_cache_hyper_ppn(
3255        &self,
3256        e: &Engine,
3257        tokens: &[u32],
3258        cache: &mut Cache,
3259        queued_after: usize,
3260        topology: &crate::hyper::HyperTopology,
3261        fence: &[usize],
3262        overlay: Option<&crate::vision::EmbedOverlay>,
3263    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3264        let n_embd = self.cfg.n_embd as usize;
3265        let eps = self.cfg.rms_eps;
3266        let t = tokens.len();
3267        let width = topology.streams * n_embd;
3268        if cache.pos + t > cache.max_ctx {
3269            return Err("prime_cache: prompt exceeds cache max_ctx".into());
3270        }
3271        let seq_end = cache.pos + t + queued_after;
3272        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
3273        // glm5 DFlash2 feature tap: set ONCE per call, before any stage walk — every stage
3274        // range of this chunk writes the same rows at the chunk's absolute offset.
3275        if let Some(sink) = cache.hc_taps.as_mut() {
3276            sink.base = cache.pos;
3277        }
3278
3279        if crate::pp::pp2_streams_off() {
3280            let pos_d = e.htod_i32(&pos)?;
3281            let mut embedded = self.embed(e, tokens)?;
3282            if let Some(ov) = overlay {
3283                // Mixed-embedding splice at embedding intake (the same point the streams-on
3284                // arm and prime_chunk_hyper use). The caller's overlay is already windowed
3285                // to THIS call (prefill_tick rebases spans call-relative), so chunk_off = 0.
3286                ov.splice_into(e, &mut embedded, 0, t, n_embd)?;
3287            }
3288            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
3289            x = self.hyper_range_prime(
3290                e,
3291                topology,
3292                x,
3293                fence[0],
3294                fence[1],
3295                &pos_d,
3296                t,
3297                cache,
3298                seq_end,
3299                HcTapArm::FromCache,
3300            )?;
3301            for s in 1..fence.len() - 1 {
3302                let boundary_tx = e.clone_dtod(&x)?;
3303                let boundary_rx = e.clone_dtod(&boundary_tx)?;
3304                x = self.hyper_range_prime(
3305                    e,
3306                    topology,
3307                    boundary_rx,
3308                    fence[s],
3309                    fence[s + 1],
3310                    &pos_d,
3311                    t,
3312                    cache,
3313                    seq_end,
3314                    HcTapArm::FromCache,
3315                )?;
3316            }
3317            return self.hyper_prime_tail(e, topology, &x, t, n_embd, eps, cache);
3318        }
3319
3320        {
3321            let rt = crate::pp::PpNRt::get(e)?;
3322            let n_st = fence.len() - 1;
3323            assert_eq!(
3324                rt.n_stages(),
3325                n_st,
3326                "PpNRt stage count {} != fence stages {n_st}",
3327                rt.n_stages()
3328            );
3329            let caller_stream = e.stream();
3330            rt.fence_stages_behind(&caller_stream)?;
3331            // OVERLAY RESIDENCY LAW (rewritten in lane/glm53-vision-ppn, 2026-09-01; was the
3332            // OVERLAY DEVICE LAW). The splice reads `overlay.rows` through stage 0's engine,
3333            // so those rows must live in STAGE 0's CUDA context. That is the real invariant,
3334            // and it is what is checked.
3335            //
3336            // The pre-lane check was `!std::ptr::eq(rt.engine(0, e), e)` — "stage 0 must BE
3337            // the primary engine". It refused the deployed 3-card shape outright: the worker's
3338            // primary engine follows the LAST pp stage (`worker::worker_device`, and that is
3339            // load-bearing — pinning the primary to stage 0 was the v0.72 tag-blocker-2
3340            // regressor, 112.5 -> 17.5 tok/s on spec+PP), so with MEMRA_PP_DEVICES=0,1,2 the
3341            // primary is dev2 while stage 0 owns dev0 and `PpNRt::build` hands stage 0 its
3342            // own Engine. Vision was therefore unservable on the ppN shape with the only
3343            // in-tag rollback costing ~3x decode (MEMRA_PP_STREAMS=0). The fix is to PUBLISH
3344            // the overlay into stage 0's context at construction
3345            // (`EmbedOverlay::new_published`, driven by `vision_intake_engine` below), never
3346            // to relax the check: the identity now passes because the pointers ARE in the
3347            // right domain.
3348            //
3349            // ORDERING, the seam the accrace lane taught (MEMRA_PP_EXIT_PUBLISH): rows built
3350            // on the caller's stream are covered by `fence_stages_behind` above, which orders
3351            // every stage stream behind the caller before stage 0 issues its splice; rows
3352            // published onto the intake engine were host-synchronized when they were uploaded.
3353            // Both producers are ordered before this body's first stage-0 kernel.
3354            if let Some(ov) = overlay
3355                && !ov.resident_in(rt.engine(0, e))
3356            {
3357                return Err(format!(
3358                    "vision embedding overlay rows are resident on dev{} but pp stage 0's \
3359                     embedding intake runs on dev{}: the overlay must be published into the \
3360                     intake engine's context (build it with EmbedOverlay::new_published; \
3361                     MEMRA_VISION_OVERLAY_PUBLISH=0 pins the pre-publication program, whose \
3362                     only vision-capable shape is MEMRA_PP_STREAMS=0)",
3363                    ov.ctx().ordinal(),
3364                    rt.engine(0, e).ctx().ordinal(),
3365                )
3366                .into());
3367            }
3368            let mut slot = {
3369                let _st0 = rt.enter(0);
3370                let e0 = rt.engine(0, e);
3371                let pos_d = e0.htod_i32(&pos)?;
3372                let mut embedded = self.embed(e0, tokens)?;
3373                if let Some(ov) = overlay {
3374                    // Mixed-embedding splice at stage-0 embedding intake, BEFORE stream
3375                    // expansion — the reference's execute_multimodal splice point. Stages
3376                    // s > 0 only ever see the [t, streams, hidden] boundary payload, so no
3377                    // other stage carries overlay arithmetic. chunk_off = 0: the caller's
3378                    // overlay is already windowed to this call.
3379                    ov.splice_into(e0, &mut embedded, 0, t, n_embd)?;
3380                }
3381                let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
3382                let x = self.hyper_range_prime(
3383                    e0,
3384                    topology,
3385                    x,
3386                    fence[0],
3387                    fence[1],
3388                    &pos_d,
3389                    t,
3390                    cache,
3391                    seq_end,
3392                    HcTapArm::FromCache,
3393                )?;
3394                rt.tx(0, &x, t * width)?
3395            };
3396            for s in 1..n_st - 1 {
3397                let _st = rt.enter(s);
3398                let es = rt.engine(s, e);
3399                let pos_d = es.htod_i32(&pos)?;
3400                let x = rt.rx(s - 1, slot, t * width)?;
3401                let x = self.hyper_range_prime(
3402                    es,
3403                    topology,
3404                    x,
3405                    fence[s],
3406                    fence[s + 1],
3407                    &pos_d,
3408                    t,
3409                    cache,
3410                    seq_end,
3411                    HcTapArm::FromCache,
3412                )?;
3413                slot = rt.tx(s, &x, t * width)?;
3414            }
3415            let out = {
3416                let _stl = rt.enter(n_st - 1);
3417                let el = rt.engine(n_st - 1, e);
3418                let pos_d = el.htod_i32(&pos)?;
3419                let x = rt.rx(n_st - 2, slot, t * width)?;
3420                let x = self.hyper_range_prime(
3421                    el,
3422                    topology,
3423                    x,
3424                    fence[n_st - 1],
3425                    fence[n_st],
3426                    &pos_d,
3427                    t,
3428                    cache,
3429                    seq_end,
3430                    HcTapArm::FromCache,
3431                )?;
3432                self.hyper_prime_tail(el, topology, &x, t, n_embd, eps, cache)?
3433            };
3434            // EXIT PUBLICATION (lane/glm5-accrace 2026-09-01, the same law the batched and
3435            // spec ppN bodies carry): `hyper_prime_tail`'s dtoh drains the LAST stage only,
3436            // and the TX-wait chain reaches each earlier stage just as far as its `ev_tx`.
3437            // Every earlier stage's stream still holds the tail its stage-scope locals
3438            // enqueue on drop, and the caller resumes here to allocate (the glm5 spec
3439            // session's MTP plane warm, the worker's next step). Anatomy + receipts:
3440            // `PpNRt::publish_all_to`.
3441            rt.publish_all_to(&caller_stream)?;
3442            Ok(out)
3443        }
3444    }
3445
3446    /// DOOR `MEMRA_B200_PRIME_V2` arm 2: the PIPELINED mHC PP-2 prime. Stage 0 of chunk k+1
3447    /// runs on device 0 while stage 1 of chunk k runs on device 1, over disjoint per-stage
3448    /// caches ([`PrimeCacheStages`]), on two scoped host threads.
3449    ///
3450    /// TWO HOST THREADS ARE REQUIRED, not a style choice, and this is the whole reason the
3451    /// serial chunk loop in `prime_cache_hyper` gets no overlap for free: the grouped MoE
3452    /// prefill's sigmoid host oracle drains its stage's stream ONCE PER LAYER
3453    /// (`Engine::moe_router_sigmoid_topk_host`), 21 times per stage per chunk on the deployed
3454    /// 24/21 cut. Two CUDA streams driven from one host thread therefore serialize however the
3455    /// calls are ordered — the host is the serializer, not the device. `prime_cache_pp2_pipelined`
3456    /// records exactly this finding for the serial trunk; this is its mHC twin, and the two
3457    /// schedulers are deliberately kept in the same shape so a fix to one reads as a fix to the
3458    /// other.
3459    ///
3460    /// BIT-IDENTICAL to the serial chunk loop at the SAME ranges: every chunk runs the same
3461    /// `hyper_range_prime` calls on the same operand bytes through the same stage engines, and
3462    /// the boundary payload is an exact copy of the `[t, streams, hidden]` state. Only which
3463    /// host thread issues a stage, and when, changes. That is arm 2's bar in
3464    /// `glm5-prime-v2-gate`; arm 1 (the schedule) carries the near-tie band instead, and the two
3465    /// axes are separated so a scheduling bug cannot hide inside a tolerance.
3466    ///
3467    /// REFUSALS, by name rather than by silence. A vision overlay declines (the splice is a
3468    /// stage-0 embedding-intake transform whose gate ran on the serial ppN body), and an ARMED
3469    /// hc tap sink declines because [`PrimeCacheStages`] gives each stage shell
3470    /// `hc_taps: None` — the DFlash2 draft-source rows would go quietly missing, which is the
3471    /// exact shape of a loud failure failing quietly. Both fall back to the serial chunk loop.
3472    #[allow(clippy::too_many_arguments)] // allow: mirrors prime_cache_pp2_pipelined's contract
3473    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3474    fn prime_cache_hyper_pp2_pipelined(
3475        &self,
3476        e: &Engine,
3477        tokens: &[u32],
3478        cache: &mut Cache,
3479        seq_end: usize,
3480        topology: &crate::hyper::HyperTopology,
3481        ranges: &[(usize, usize)],
3482        fence: &[usize],
3483        taps: Option<&std::sync::Mutex<&mut crate::cache::HcTapSink>>,
3484    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3485        debug_assert_eq!(fence.len(), 3);
3486        debug_assert!(ranges.len() >= 2);
3487        // Each stage call gets its OWN chunk base; the sink itself is shared. See `HcTapArm`
3488        // for why that is sound (disjoint slot columns AND disjoint row windows).
3489        let arm = |base: usize| match taps {
3490            Some(sink) => HcTapArm::Shared(sink, base),
3491            None => HcTapArm::FromCache,
3492        };
3493        let rt = crate::pp::PpNRt::get(e)?;
3494        assert_eq!(
3495            rt.n_stages(),
3496            2,
3497            "the pipelined mHC prime requires exactly two PP stages"
3498        );
3499        let n_embd = self.cfg.n_embd as usize;
3500        let eps = self.cfg.rms_eps;
3501        let width = topology.streams * n_embd;
3502        let t = tokens.len();
3503        let initial_base = cache.pos;
3504        let caller_stream = e.stream();
3505
3506        // #87 reverse publication before any new stage allocation, then prewarm both boundary
3507        // slots while the stage streams are empty — a lazily grown slot B after stage 1(N) is
3508        // queued would synchronize that stream and erase the first overlap.
3509        rt.fence_stages_behind(&caller_stream)?;
3510        let max_payload = ranges.iter().map(|(s, x)| (x - s) * width).max().unwrap();
3511        rt.prepare_overlap_slots(0, max_payload)?;
3512
3513        // ENGAGEMENT RECEIPT, once per TAP STATE rather than once per process. The armed case
3514        // is the one the product route runs and the one this arm used to refuse, so a log that
3515        // only ever showed the first state reached could show `taps=none` forever while the
3516        // interesting half went unrecorded.
3517        static SAID: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3518        let bit = 1u8 << u8::from(taps.is_some());
3519        if SAID.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
3520            eprintln!(
3521                "[prime-v2] arm2 pipelined stages=2 chunks={} overlap={} t={t} devices={:?} \
3522                 taps={} (stage 0 of chunk k+1 overlaps stage 1 of chunk k on two host \
3523                 threads; MEMRA_B200_PRIME_V2, logged once per process)",
3524                ranges.len(),
3525                ranges.len().saturating_sub(1),
3526                (0..2)
3527                    .map(|s| rt.engine(s, e).ctx().ordinal())
3528                    .collect::<Vec<_>>(),
3529                if taps.is_some() { "armed" } else { "none" },
3530            );
3531        }
3532
3533        let mut hiddens = e.uninit(t * n_embd)?;
3534        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3535        let mut stage_caches = PrimeCacheStages::new(cache, fence);
3536        let (cache0, cache1) = stage_caches.pp2_parts();
3537        let (first_start, first_end) = ranges[0];
3538        let mut slot = self.prime_hyper_pp2_stage0_enqueue(
3539            e,
3540            rt,
3541            topology,
3542            &tokens[first_start..first_end],
3543            cache0,
3544            seq_end,
3545            fence,
3546            initial_base + first_start,
3547            arm(initial_base + first_start),
3548        )?;
3549        cache0.pos = initial_base + first_end;
3550
3551        for (i, &(start, end)) in ranges.iter().enumerate() {
3552            let base = initial_base + start;
3553            debug_assert_eq!(
3554                cache1.pos, base,
3555                "stage 1 must drain chunks in original position order"
3556            );
3557            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
3558                let next_base = initial_base + next_start;
3559                debug_assert_eq!(
3560                    cache0.pos, next_base,
3561                    "stage 0 must issue chunks in original position order"
3562                );
3563                let cache0_stage = &mut *cache0;
3564                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
3565                    let stage0 = scope.spawn(move || -> Result<usize, String> {
3566                        let next = self
3567                            .prime_hyper_pp2_stage0_enqueue(
3568                                e,
3569                                rt,
3570                                topology,
3571                                &tokens[next_start..next_end],
3572                                cache0_stage,
3573                                seq_end,
3574                                fence,
3575                                next_base,
3576                                arm(next_base),
3577                            )
3578                            .map_err(|err| err.to_string())?;
3579                        cache0_stage.pos = initial_base + next_end;
3580                        Ok(next)
3581                    });
3582                    let x = self.prime_hyper_pp2_stage1_enqueue(
3583                        e,
3584                        rt,
3585                        topology,
3586                        slot,
3587                        end - start,
3588                        cache1,
3589                        seq_end,
3590                        fence,
3591                        base,
3592                        arm(base),
3593                    )?;
3594                    let out = {
3595                        rt.bind_stage(1)?;
3596                        let _st1 = rt.enter(1);
3597                        let e1 = rt.engine(1, e);
3598                        self.hyper_prime_tail(e1, topology, &x, end - start, n_embd, eps, cache1)?
3599                    };
3600                    let next = match stage0.join() {
3601                        Ok(result) => {
3602                            result.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?
3603                        }
3604                        Err(payload) => std::panic::resume_unwind(payload),
3605                    };
3606                    Ok((out, Some(next)))
3607                })?
3608            } else {
3609                let x = self.prime_hyper_pp2_stage1_enqueue(
3610                    e,
3611                    rt,
3612                    topology,
3613                    slot,
3614                    end - start,
3615                    cache1,
3616                    seq_end,
3617                    fence,
3618                    base,
3619                    arm(base),
3620                )?;
3621                let out = {
3622                    rt.bind_stage(1)?;
3623                    let _st1 = rt.enter(1);
3624                    let e1 = rt.engine(1, e);
3625                    self.hyper_prime_tail(e1, topology, &x, end - start, n_embd, eps, cache1)?
3626                };
3627                (out, None)
3628            };
3629
3630            rt.publish_to(1, &caller_stream)?;
3631            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
3632            last = Some((out.0, out.1));
3633            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3634            // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
3635            // so the device finished it. Stamp the odometer /health reads, so a BUSY
3636            // worker mid-long-prefill is never mistaken for a wedged one.
3637            crate::progress::note_prime_rows(end - start);
3638            HYPER_PRIME_PIPELINED_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3639
3640            if let Some(next) = next_slot {
3641                // The caller copy above reads a stage-1 allocation. Before stage 1 of the next
3642                // chunk can allocate or reuse blocks, mirror #87's body-entry fence. Stage
3643                // 0(N+1) is already queued before this wait is appended, so its overlap with
3644                // stage 1(N) is preserved.
3645                rt.fence_stages_behind(&caller_stream)?;
3646                slot = next;
3647            }
3648        }
3649
3650        debug_assert_eq!(cache0.pos, initial_base + t);
3651        debug_assert_eq!(cache1.pos, initial_base + t);
3652        let (logits, h_seed) = last.expect("the pipelined mHC prime ran at least one chunk");
3653        stage_caches.commit();
3654        Ok((logits, h_seed, hiddens))
3655    }
3656
3657    /// Stage 0 of one pipelined mHC prime chunk: embed, expand to the `[t, streams, hidden]`
3658    /// stream state, walk `fence[0]..fence[1]`, hand the state to the boundary. The mHC twin of
3659    /// [`Self::prime_pp2_stage0_enqueue`]; the ONLY differences are the stream expansion and the
3660    /// `streams * n_embd` payload width.
3661    #[allow(clippy::too_many_arguments)] // allow: mirrors prime_pp2_stage0_enqueue's contract
3662    fn prime_hyper_pp2_stage0_enqueue(
3663        &self,
3664        e: &Engine,
3665        rt: &crate::pp::PpNRt,
3666        topology: &crate::hyper::HyperTopology,
3667        tokens: &[u32],
3668        cache: &mut Cache,
3669        seq_end: usize,
3670        fence: &[usize],
3671        base: usize,
3672        taps: HcTapArm<'_, '_>,
3673    ) -> Result<usize, Box<dyn std::error::Error>> {
3674        let t = tokens.len();
3675        let n_embd = self.cfg.n_embd as usize;
3676        let width = topology.streams * n_embd;
3677        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
3678        rt.bind_stage(0)?;
3679        let _st0 = rt.enter(0);
3680        let e0 = rt.engine(0, e);
3681        let pos_d = e0.htod_i32(&pos)?;
3682        let embedded = self.embed(e0, tokens)?;
3683        let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
3684        let _overlap = crate::pp::enter_prime_pipe_stage();
3685        let x = self.hyper_range_prime(
3686            e0, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end, taps,
3687        )?;
3688        rt.tx_pipelined(0, &x, t * width)
3689    }
3690
3691    /// Stage 1 of one pipelined mHC prime chunk: receive the stream state, walk
3692    /// `fence[1]..fence[2]`. The caller runs [`Self::hyper_prime_tail`] on the result, which is
3693    /// where `cache.pos` advances — exactly as the serial ppN body orders it.
3694    #[allow(clippy::too_many_arguments)] // allow: mirrors prime_pp2_stage1_enqueue's contract
3695    fn prime_hyper_pp2_stage1_enqueue(
3696        &self,
3697        e: &Engine,
3698        rt: &crate::pp::PpNRt,
3699        topology: &crate::hyper::HyperTopology,
3700        slot: usize,
3701        t: usize,
3702        cache: &mut Cache,
3703        seq_end: usize,
3704        fence: &[usize],
3705        base: usize,
3706        taps: HcTapArm<'_, '_>,
3707    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3708        let n_embd = self.cfg.n_embd as usize;
3709        let width = topology.streams * n_embd;
3710        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
3711        rt.bind_stage(1)?;
3712        let _st1 = rt.enter(1);
3713        let e1 = rt.engine(1, e);
3714        let pos_d = e1.htod_i32(&pos)?;
3715        let x = rt.rx(0, slot, t * width)?;
3716        let _overlap = crate::pp::enter_prime_pipe_stage();
3717        self.hyper_range_prime(
3718            e1, topology, x, fence[1], fence[2], &pos_d, t, cache, seq_end, taps,
3719        )
3720    }
3721
3722    /// ONE call of the mHC prime walk. Carries `tokens.len()` rows of stream state and of every
3723    /// per-layer transient, appends this call's rows to the mixers' own state, and advances
3724    /// `cache.pos`. `seq_end` is the REQUEST's absolute end, passed in rather than recomputed,
3725    /// so no arithmetic here is a function of how the prompt was split.
3726    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3727    fn prime_chunk_hyper(
3728        &self,
3729        e: &Engine,
3730        tokens: &[u32],
3731        cache: &mut Cache,
3732        seq_end: usize,
3733        chunk_off: usize,
3734        overlay: Option<&crate::vision::EmbedOverlay>,
3735    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3736        let topology = *self
3737            .hyper
3738            .as_ref()
3739            .ok_or("prime_chunk_hyper on a model with no HyperConnections topology")?;
3740        let n_embd = self.cfg.n_embd as usize;
3741        let t = tokens.len();
3742        let eps = self.cfg.rms_eps;
3743        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
3744        let pos_d = e.htod_i32(&pos)?;
3745        // glm5 DFlash2 feature tap: chunked primes write their rows at the chunk's absolute
3746        // offset (cache.pos advances per chunk — the dflash_taps.base precedent).
3747        if let Some(sink) = cache.hc_taps.as_mut() {
3748            sink.base = cache.pos;
3749        }
3750
3751        let mut embedded = self.embed(e, tokens)?;
3752        if let Some(ov) = overlay {
3753            // Mixed-embedding splice (shared with the ppN twin — EmbedOverlay::splice_into):
3754            // image rows overwrite placeholder-token embeddings inside this chunk's
3755            // prompt-relative window [chunk_off, chunk_off+t), BEFORE stream expansion —
3756            // the reference's splice point (execute_multimodal replaces rows before
3757            // hc_expand).
3758            ov.splice_into(e, &mut embedded, chunk_off, t, n_embd)?;
3759        }
3760        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
3761
3762        for (il, layer) in self.layers.iter().enumerate() {
3763            let hyper = layer.hyper.as_ref().ok_or_else(|| {
3764                format!("layer {il} carries no hyper-connection weights under an hc plan")
3765            })?;
3766
3767            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.attn, &x, t, n_embd)?;
3768            let mut h = e.uninit(t * n_embd)?;
3769            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3770            let mixed = match &layer.mixer {
3771                Mixer::Full(fa) => {
3772                    self.full_attn_prime(e, fa, &h, None, &pos_d, t, cache, il, seq_end)?
3773                }
3774                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
3775                Mixer::Mla(mla) if mla.tp.is_some() => {
3776                    self.mla_tp_attn_cached(e, mla, &h, &pos_d, t, il, cache, false)?
3777                }
3778                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, &pos_d, t, il, cache)?,
3779                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
3780                    e,
3781                    la,
3782                    &h,
3783                    t,
3784                    eps,
3785                    cache,
3786                    il,
3787                    crate::kda::ConvArm::Prefill,
3788                )?,
3789                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
3790            };
3791            x = crate::hyper::post(e, &topology, &mixed, &x, &mix, t, n_embd)?;
3792
3793            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.mlp, &x, t, n_embd)?;
3794            let mut z = e.uninit(t * n_embd)?;
3795            e.rms_norm(
3796                &y,
3797                layer.post_attn_norm.float_data(),
3798                &mut z,
3799                n_embd,
3800                t,
3801                eps,
3802            )?;
3803            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true, None)?;
3804            x = crate::hyper::post(e, &topology, &ffn_out, &x, &mix, t, n_embd)?;
3805            // glm5 DFlash2 feature tap (see hyper_range_prime — the unsplit chunk walk
3806            // taps the same completed-layer-output contraction).
3807            self.glm5_hc_tap(e, cache, &topology, il, &x, t)?;
3808        }
3809
3810        let hiddens =
3811            crate::hyper::collapse(e, &topology, self.hyper_head.as_ref(), &x, t, n_embd)?;
3812        let mut hn = e.uninit(t * n_embd)?;
3813        e.rms_norm(
3814            &hiddens,
3815            self.output_norm.float_data(),
3816            &mut hn,
3817            n_embd,
3818            t,
3819            eps,
3820        )?;
3821        let last = e.view(&hn, t * n_embd);
3822        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
3823        let mut hlast = e.uninit(n_embd)?;
3824        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3825        let logits = e.matmul(&self.output, &hlast, 1)?;
3826        let host = e.dtoh(&logits)?;
3827
3828        // h_seed is the PRE-output_norm hidden of the last row (MTP-PLAN §A), taken from the
3829        // collapsed stack so it means the same thing it does on the serial path.
3830        let stack = e.view(&hiddens, t * n_embd);
3831        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
3832        let mut h_seed = e.uninit(n_embd)?;
3833        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
3834        cache.pos += t;
3835        Ok((host, h_seed, hiddens))
3836    }
3837
3838    /// Prime exit shared by `prime_cache_hyper` and its ppN twin: collapse, output_norm, last
3839    /// row logits, and the pre-output_norm hidden seed taken from the collapsed stack (MTP-PLAN
3840    /// §A) so it means the same thing it does on the serial path.
3841    #[allow(clippy::too_many_arguments)]
3842    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3843    fn hyper_prime_tail(
3844        &self,
3845        e: &Engine,
3846        topology: &crate::hyper::HyperTopology,
3847        x: &CudaSlice<f32>,
3848        t: usize,
3849        n_embd: usize,
3850        eps: f32,
3851        cache: &mut Cache,
3852    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3853        let hiddens = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
3854        let mut hn = e.uninit(t * n_embd)?;
3855        e.rms_norm(
3856            &hiddens,
3857            self.output_norm.float_data(),
3858            &mut hn,
3859            n_embd,
3860            t,
3861            eps,
3862        )?;
3863        let last = e.view(&hn, t * n_embd);
3864        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
3865        let mut hlast = e.uninit(n_embd)?;
3866        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3867        let logits = e.matmul(&self.output, &hlast, 1)?;
3868        let host = e.dtoh(&logits)?;
3869        let stack = e.view(&hiddens, t * n_embd);
3870        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
3871        let mut h_seed = e.uninit(n_embd)?;
3872        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
3873        cache.pos += t;
3874        Ok((host, h_seed, hiddens))
3875    }
3876
3877    /// ppN twin of `decode_step_hyper`: the T=1 step as N stage subgraphs, each on its own
3878    /// stream (and, under `MEMRA_PP_DEVICES`, its own device/engine), with the
3879    /// transport-selected boundary handoff of the `[streams, n_embd]` state at each fence cut.
3880    /// `cache.pos` is snapshotted once and advanced once.
3881    fn decode_step_hyper_ppn(
3882        &self,
3883        e: &Engine,
3884        token: u32,
3885        cache: &mut Cache,
3886        topology: &crate::hyper::HyperTopology,
3887        fence: &[usize],
3888    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3889        let n_embd = self.cfg.n_embd as usize;
3890        let eps = self.cfg.rms_eps;
3891        let pos = cache.pos;
3892        let width = topology.streams * n_embd;
3893
3894        if crate::pp::pp2_streams_off() {
3895            let pos_d = e.htod_i32(&[pos as i32])?;
3896            let embedded = e.htod(&self.embd.try_gather(n_embd, &[token])?)?;
3897            let mut x = crate::hyper::expand(e, topology, &embedded, 1, n_embd)?;
3898            x = self.hyper_range_decode(e, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
3899            for s in 1..fence.len() - 1 {
3900                let boundary_tx = e.clone_dtod(&x)?;
3901                let boundary_rx = e.clone_dtod(&boundary_tx)?;
3902                x = self.hyper_range_decode(
3903                    e,
3904                    topology,
3905                    boundary_rx,
3906                    fence[s],
3907                    fence[s + 1],
3908                    &pos_d,
3909                    pos,
3910                    cache,
3911                )?;
3912            }
3913            return self.hyper_decode_tail(e, topology, &x, n_embd, eps, cache);
3914        }
3915
3916        let rt = crate::pp::PpNRt::get(e)?;
3917        let n_st = fence.len() - 1;
3918        assert_eq!(
3919            rt.n_stages(),
3920            n_st,
3921            "PpNRt stage count {} != fence stages {n_st}",
3922            rt.n_stages()
3923        );
3924        rt.fence_stages_behind(&e.stream())?;
3925
3926        let mut slot = {
3927            let _st0 = rt.enter(0);
3928            let e0 = rt.engine(0, e);
3929            let pos_d = e0.htod_i32(&[pos as i32])?;
3930            let embedded = e0.htod(&self.embd.try_gather(n_embd, &[token])?)?;
3931            let x = crate::hyper::expand(e0, topology, &embedded, 1, n_embd)?;
3932            let x =
3933                self.hyper_range_decode(e0, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
3934            rt.tx(0, &x, width)?
3935        };
3936        for s in 1..n_st - 1 {
3937            let _st = rt.enter(s);
3938            let es = rt.engine(s, e);
3939            let pos_d = es.htod_i32(&[pos as i32])?;
3940            let x = rt.rx(s - 1, slot, width)?;
3941            let x = self.hyper_range_decode(
3942                es,
3943                topology,
3944                x,
3945                fence[s],
3946                fence[s + 1],
3947                &pos_d,
3948                pos,
3949                cache,
3950            )?;
3951            slot = rt.tx(s, &x, width)?;
3952        }
3953        let _stl = rt.enter(n_st - 1);
3954        let el = rt.engine(n_st - 1, e);
3955        let pos_d = el.htod_i32(&[pos as i32])?;
3956        let x = rt.rx(n_st - 2, slot, width)?;
3957        let x = self.hyper_range_decode(
3958            el,
3959            topology,
3960            x,
3961            fence[n_st - 1],
3962            fence[n_st],
3963            &pos_d,
3964            pos,
3965            cache,
3966        )?;
3967        self.hyper_decode_tail(el, topology, &x, n_embd, eps, cache)
3968    }
3969
3970    /// Decode exit shared by `decode_step_hyper` and its ppN twin. `h_seed` is the COLLAPSED
3971    /// hidden (not the pre-collapse stream state and not the post-norm row): that is what the
3972    /// serial hc step publishes, and the two must not drift.
3973    fn hyper_decode_tail(
3974        &self,
3975        e: &Engine,
3976        topology: &crate::hyper::HyperTopology,
3977        x: &CudaSlice<f32>,
3978        n_embd: usize,
3979        eps: f32,
3980        cache: &mut Cache,
3981    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3982        let h_seed = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, 1, n_embd)?;
3983        let mut hn = e.uninit(n_embd)?;
3984        e.rms_norm(
3985            &h_seed,
3986            self.output_norm.float_data(),
3987            &mut hn,
3988            n_embd,
3989            1,
3990            eps,
3991        )?;
3992        let logits = e.matmul(&self.output, &hn, 1)?;
3993        let host = e.dtoh(&logits)?;
3994        cache.pos += 1;
3995        Ok((host, h_seed))
3996    }
3997
3998    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
3999    pub fn forward(
4000        &self,
4001        e: &Engine,
4002        tokens: &[u32],
4003    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4004        if self.hyper.is_some() {
4005            return self.forward_hyper(e, tokens, false);
4006        }
4007        if self.is_gemma4_e4b() {
4008            return self.gemma4_e4b_forward(e, tokens, false);
4009        }
4010        if self.uses_gemma_program() {
4011            return self.gemma4_forward(e, tokens, false);
4012        }
4013        let cfg = &self.cfg;
4014        let n_embd = cfg.n_embd as usize;
4015        let t = tokens.len();
4016        let eps = cfg.rms_eps;
4017        let pos: Vec<i32> = (0..t as i32).collect();
4018        let pos_d = e.htod_i32(&pos)?;
4019
4020        let mut x = self.embed(e, tokens)?; // [T, n_embd]
4021
4022        for (il, layer) in self.layers.iter().enumerate() {
4023            // attn_norm
4024            let mut h = e.uninit(t * n_embd)?;
4025            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4026
4027            let mixed = match &layer.mixer {
4028                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
4029                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
4030                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
4031                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
4032            };
4033
4034            // residual 1
4035            let mut x1 = e.uninit(t * n_embd)?;
4036            e.add(&x, &mixed, &mut x1, t * n_embd)?;
4037
4038            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
4039            let mut z = e.uninit(t * n_embd)?;
4040            e.rms_norm(
4041                &x1,
4042                layer.post_attn_norm.float_data(),
4043                &mut z,
4044                n_embd,
4045                t,
4046                eps,
4047            )?;
4048            let ffn_out = match &layer.ffn {
4049                crate::hybrid::Ffn::Dense {
4050                    ffn_gate,
4051                    ffn_up,
4052                    ffn_down,
4053                } => {
4054                    let n_ff = ffn_gate.out_features();
4055                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
4056                    let up = g2.pop().unwrap();
4057                    let gate = g2.pop().unwrap();
4058                    let mut act = e.uninit(t * n_ff)?;
4059                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
4060                    // both the dense MLP and the shared expert, and its limit is
4061                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
4062                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
4063                    Self::ffn_act_lim(
4064                        e,
4065                        &self.cfg,
4066                        &gate,
4067                        &up,
4068                        1.0,
4069                        1.0,
4070                        self.cfg.clamp_shexp_at(il as u32),
4071                        &mut act,
4072                        t * n_ff,
4073                    )?;
4074                    e.matmul(ffn_down, &act, t)?
4075                }
4076                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
4077            };
4078            let mut x2 = e.uninit(t * n_embd)?;
4079            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4080            x = x2;
4081        }
4082
4083        let mut hn = e.uninit(t * n_embd)?;
4084        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4085        let logits = e.matmul(&self.output, &hn, t)?;
4086        e.dtoh(&logits)
4087    }
4088
4089    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
4090    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
4091    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
4092    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
4093    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
4094    pub fn forward_last(
4095        &self,
4096        e: &Engine,
4097        tokens: &[u32],
4098    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4099        if self.hyper.is_some() {
4100            return self.forward_hyper(e, tokens, true);
4101        }
4102        if self.uses_gemma_program() {
4103            return self.gemma4_forward(e, tokens, true);
4104        }
4105        let cfg = &self.cfg;
4106        let n_embd = cfg.n_embd as usize;
4107        let t = tokens.len();
4108        let eps = cfg.rms_eps;
4109        let pos: Vec<i32> = (0..t as i32).collect();
4110        let pos_d = e.htod_i32(&pos)?;
4111
4112        let mut x = self.embed(e, tokens)?; // [T, n_embd]
4113        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
4114        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
4115        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
4116        let anat = Self::prime_anatomy_on();
4117        let mut anat_last = if anat {
4118            e.stream().synchronize()?;
4119            Some(std::time::Instant::now())
4120        } else {
4121            None
4122        };
4123        macro_rules! anat_mark {
4124            ($slot:expr) => {
4125                if let Some(ts) = anat_last.as_mut() {
4126                    e.stream().synchronize()?;
4127                    Self::prime_anatomy_slots()[$slot].fetch_add(
4128                        ts.elapsed().as_nanos() as u64,
4129                        std::sync::atomic::Ordering::Relaxed,
4130                    );
4131                    *ts = std::time::Instant::now();
4132                }
4133            };
4134        }
4135        for (il, layer) in self.layers.iter().enumerate() {
4136            let mut h = e.uninit(t * n_embd)?;
4137            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4138            if probe {
4139                e.stream().synchronize()?;
4140                eprintln!("[probe] L{il} norm ok");
4141            }
4142            anat_mark!(4);
4143            let mixed = match &layer.mixer {
4144                Mixer::Full(fa) => {
4145                    let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
4146                    anat_mark!(0);
4147                    y
4148                }
4149                Mixer::Linear(la) => {
4150                    let y = self.linear_attn(e, la, &h, t)?;
4151                    anat_mark!(1);
4152                    y
4153                }
4154                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
4155                Mixer::Kda(la) => {
4156                    let y = crate::kda::kda_attn(e, la, &h, t, eps)?;
4157                    // KDA shares the linear-mixer anatomy slot: same mixer class, one bucket.
4158                    anat_mark!(1);
4159                    y
4160                }
4161            };
4162            if probe {
4163                e.stream().synchronize()?;
4164                eprintln!("[probe] L{il} mixer ok");
4165            }
4166            let mut x1 = e.uninit(t * n_embd)?;
4167            e.add(&x, &mixed, &mut x1, t * n_embd)?;
4168            let mut z = e.uninit(t * n_embd)?;
4169            e.rms_norm(
4170                &x1,
4171                layer.post_attn_norm.float_data(),
4172                &mut z,
4173                n_embd,
4174                t,
4175                eps,
4176            )?;
4177            anat_mark!(4);
4178            let ffn_out = match &layer.ffn {
4179                crate::hybrid::Ffn::Dense {
4180                    ffn_gate,
4181                    ffn_up,
4182                    ffn_down,
4183                } => {
4184                    let n_ff = ffn_gate.out_features();
4185                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
4186                    let up = g2.pop().unwrap();
4187                    let gate = g2.pop().unwrap();
4188                    let mut act = e.uninit(t * n_ff)?;
4189                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
4190                    Self::ffn_act_lim(
4191                        e,
4192                        &self.cfg,
4193                        &gate,
4194                        &up,
4195                        1.0,
4196                        1.0,
4197                        self.cfg.clamp_shexp_at(il as u32),
4198                        &mut act,
4199                        t * n_ff,
4200                    )?;
4201                    let y = e.matmul(ffn_down, &act, t)?;
4202                    anat_mark!(3);
4203                    y
4204                }
4205                crate::hybrid::Ffn::Moe(m) => {
4206                    let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
4207                    anat_mark!(2);
4208                    y
4209                }
4210            };
4211            if probe {
4212                e.stream().synchronize()?;
4213                eprintln!("[probe] L{il} ffn ok");
4214            }
4215            let mut x2 = e.uninit(t * n_embd)?;
4216            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4217            x = x2;
4218        }
4219        if anat {
4220            let s = Self::prime_anatomy_slots();
4221            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
4222            eprintln!(
4223                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
4224                 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
4225                ms(0),
4226                ms(1),
4227                ms(2),
4228                ms(3),
4229                ms(4)
4230            );
4231        }
4232        // norm over all T, then slice the LAST row and run lm_head on that single row.
4233        let mut hn = e.uninit(t * n_embd)?;
4234        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4235        let last = e.view(&hn, t * n_embd); // [T, n_embd]
4236        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
4237        let mut hlast = e.uninit(n_embd)?;
4238        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
4239        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
4240        e.dtoh(&logits)
4241    }
4242
4243    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
4244    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
4245    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
4246    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
4247    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
4248    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
4249    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
4250    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
4251    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
4252    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
4253    ///       argmax gate is the accuracy authority, exactly as for forward_last);
4254    ///   (c) `cache.pos`/KV len/len_d advance by T.
4255    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
4256    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
4257    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
4258    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
4259    ///
4260    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
4261    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
4262    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
4263    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
4264    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
4265    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
4266    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
4267    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
4268    /// differently under load — research/tick-seg-20260807, receipt in
4269    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
4270    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
4271    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
4272    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
4273    /// caller that SPLITS one request across calls passes the remainder.
4274    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4275    pub fn prime_cache(
4276        &self,
4277        e: &Engine,
4278        tokens: &[u32],
4279        cache: &mut Cache,
4280        queued_after: usize,
4281    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4282        self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
4283    }
4284
4285    /// The engine that owns EMBEDDING INTAKE for the current placement — the engine whose
4286    /// context a vision overlay's rows must live in (`EmbedOverlay::new_published`).
4287    ///
4288    /// It is NOT always the primary engine. Under a per-stage-stream ppN split the embedding
4289    /// happens on stage 0 (`prime_cache_hyper_ppn` embeds inside the stage-0 scope and every
4290    /// later stage sees only the expanded stream state), and stage 0 gets its own Engine
4291    /// whenever its device differs from the primary's — which is the deployed 3-card shape,
4292    /// because the worker's primary engine follows the LAST stage. On a single-device
4293    /// placement, with the door shut, or on the `MEMRA_PP_STREAMS=0` seam, intake is the
4294    /// primary engine and this returns `e` unchanged (byte-identical to the pre-lane path).
4295    ///
4296    /// This mirrors `prime_cache_hyper`'s own door test deliberately, and the ppN prime's
4297    /// residency refusal is the enforcement: if this ever picks the wrong engine, the prime
4298    /// fails CLOSED with a named error instead of peer-reading an overlay.
4299    pub fn vision_intake_engine<'a>(
4300        &self,
4301        e: &'a Engine,
4302    ) -> Result<&'a Engine, Box<dyn std::error::Error>> {
4303        if self.hyper.is_some()
4304            && !crate::pp::pp2_streams_off()
4305            && crate::pp::pp_cuts(self.layers.len()).is_some()
4306        {
4307            let rt = crate::pp::PpNRt::get(e)?;
4308            return Ok(rt.engine(0, e));
4309        }
4310        Ok(e)
4311    }
4312
4313    /// `prime_cache` with a vision embedding overlay (lane/vision): image merger outputs
4314    /// replace the `<|image_pad|>` token embeddings at prompt-relative positions before the
4315    /// trunk walk — the mixed-embedding prime. Text-only callers use `prime_cache` (overlay
4316    /// None, byte-identical path). Scope: the serial chunk walk, the single-engine hyper walk,
4317    /// and the hyper ppN twin (splice at stage-0 embedding intake; lane/glm5-vision-default-on,
4318    /// gated by glm5-hyper-ppn-gate's overlay arm). The serial PP-2 pipelined prime and
4319    /// gemma4 E4B still refuse loudly.
4320    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4321    pub fn prime_cache_overlaid(
4322        &self,
4323        e: &Engine,
4324        tokens: &[u32],
4325        cache: &mut Cache,
4326        queued_after: usize,
4327        overlay: Option<&crate::vision::EmbedOverlay>,
4328    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4329        // FORWARD PROGRESS (memra#50), the CALL-granularity half. Every chunked walk below
4330        // stamps `crate::progress` per chunk; a MONOLITHIC walk (a prompt at or under one
4331        // chunk, `MEMRA_PRIME_CHUNK=0`, gemma4 v0, the E4B arm) stamps nothing on its way
4332        // through and would leave the odometer frozen for the whole prime. Comparing the
4333        // event count across the call is what tells the two apart WITHOUT trusting any one
4334        // walk to remember: if nothing was stamped, the call's own completion is the only
4335        // honest progress point there is, and it is stamped here. The scheduler primes one
4336        // session per call, so under the wave shape that produced memra#50 this alone already
4337        // beats between sessions.
4338        let events_before = crate::progress::events();
4339        let out = self.prime_cache_overlaid_inner(e, tokens, cache, queued_after, overlay);
4340        if out.is_ok() && crate::progress::events() == events_before {
4341            crate::progress::note_prime_rows(tokens.len());
4342        }
4343        out
4344    }
4345
4346    #[allow(clippy::type_complexity)] // allow: mirrors `prime_cache_overlaid`'s signature
4347    fn prime_cache_overlaid_inner(
4348        &self,
4349        e: &Engine,
4350        tokens: &[u32],
4351        cache: &mut Cache,
4352        queued_after: usize,
4353        overlay: Option<&crate::vision::EmbedOverlay>,
4354    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4355        cache.ensure_usable("prime_cache")?;
4356        if self.hyper.is_some() {
4357            // The mixed-embedding splice lands BEFORE stream expansion (the same point
4358            // the reference's execute_multimodal replaces rows — before hc_expand), so
4359            // the hyper walk needs no overlay-specific arithmetic (lane/glm5-vision).
4360            return self.prime_cache_hyper(e, tokens, cache, queued_after, overlay);
4361        }
4362        let _pp_walk =
4363            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
4364                let rt = crate::pp::PpNRt::get(e)?;
4365                Some(rt.acquire_walk("prime_cache")?)
4366            } else {
4367                None
4368            };
4369        let n_embd = self.cfg.n_embd as usize;
4370        let t = tokens.len();
4371        // MEMRA_PRIME_TROWS=1: prefill through the same-session t-row walk (per-row t=1
4372        // program = the tokenwise-prime ORACLE class) — replaces the host-canonical
4373        // per-token step-TP prime. Text-only fresh primes; anything else falls through.
4374        // MEMRA_STEP_GEMM_PRIME: prime the prompt through the batched GEMM path, CHUNKED.
4375        // The batch entry supplies both halves of the fast prime — the GEMM trunk at m = chunk
4376        // and the grouped NVFP4 MoE — which is why routing only the MoE through the ordinary
4377        // chunk loop measured 26.9 s against 3.7 s here. Chunking keeps the transients bounded:
4378        // a whole 32k prompt in one call would build a 262144-pair CSR and ~4.3 GB of partials
4379        // per rank, the blow-up the chunked prime exists to prevent.
4380        //
4381        // CONTINUATION (lane/gemm-suffix, 2026-08-28): the entry NO LONGER requires
4382        // `cache.pos == 0`. The batch core has been continuation-capable since 7700e0b6
4383        // (positions carry each sequence's base; the fresh-prompt guard narrowed to B > 1),
4384        // and d99b2ea3 named this outer guard as the remaining blocker in its own message.
4385        // Every multi-turn suffix and every tick remainder was paying the walk's measured
4386        // ~7.2 ms/token against this path's ~1.0 ms/token, which is why session-affinity
4387        // reuse measured a 1.012x wash on a growing conversation.
4388        // ONE DEFECT HAD TO BE FIXED FIRST, and it was LIVE before this lift:
4389        // `step35_prime_batch_layers` passed `ts[s]` — the CHUNK's length — as `seq_end`.
4390        // `seq_end` is the REQUEST's absolute end position and it steers step35's SWA arm
4391        // (`seq_end > win`, win = 512 on step37). A chunk SHORTER than the window at a
4392        // NONZERO base therefore selected the UNWINDOWED FA arm over a view that the `off`
4393        // trim leaves at ~win-1+t rows: it attended OUTSIDE the sliding window. That was
4394        // already reachable with no continuation at all — a fresh prompt of 4096+k for k in
4395        // [PRIME_MIN_T, 512) ends in a trailing chunk of exactly that shape. `seq_end` is now
4396        // threaded from here (request-absolute, `+ queued_after`, computed ONCE before the
4397        // chunk loop, chunk-size-invariant exactly as on the walk), which is what makes the
4398        // suffix arm expressible at all rather than merely reachable.
4399        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
4400        let seq_end = if legacy_calllocal {
4401            cache.pos + t
4402        } else {
4403            cache.pos + t + queued_after
4404        };
4405        // MEMRA_STEP_GEMM_PRIME_SUFFIX is the SUFFIX-ONLY seam: off (its default in this
4406        // commit) leaves continuations on the walk while fresh primes stay on the fast path;
4407        // MEMRA_STEP_GEMM_PRIME=0 is the whole-path seam. The `seq_end` threading above is
4408        // deliberately NOT behind either door — it is a correctness fix for the fresh path too.
4409        if overlay.is_none()
4410            && (cache.pos == 0 || step_gemm_prime_suffix_on())
4411            && t >= PRIME_MIN_T
4412            && crate::step_gemm_prime_on()
4413            && self.uses_sliding_gated_moe_program()
4414        {
4415            let n_embd = self.cfg.n_embd as usize;
4416            let base = cache.pos;
4417            let width = crate::cache::PRIME_CHUNK_MAX_TOKENS;
4418            let mut hiddens = e.uninit(t * n_embd)?;
4419            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
4420            let mut start = 0usize;
4421            while start < t {
4422                // A trailing chunk below the walk floor folds into the previous one; every chunk
4423                // this entry sees must clear PRIME_MIN_T on its own.
4424                let mut end = (start + width).min(t);
4425                if t - end > 0 && t - end < PRIME_MIN_T {
4426                    end = t;
4427                }
4428                let mut out = self.step35_prime_cache_batch(
4429                    e,
4430                    &[&tokens[start..end]],
4431                    &mut [cache],
4432                    &[seq_end],
4433                )?;
4434                if out.len() != 1 {
4435                    return Err("B=1 batched prime returned a non-singleton".into());
4436                }
4437                let (logits, h_seed, hidden) = out.remove(0);
4438                e.copy_into(
4439                    &mut hiddens,
4440                    start * n_embd,
4441                    &hidden,
4442                    (end - start) * n_embd,
4443                )?;
4444                last = Some((logits, h_seed));
4445                // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
4446                // so the device finished it. Stamp the odometer /health reads, so a BUSY
4447                // worker mid-long-prefill is never mistaken for a wedged one.
4448                crate::progress::note_prime_rows(end - start);
4449                start = end;
4450            }
4451            let (logits, h_seed) = last.expect("prime produced no chunk");
4452            // ENGAGEMENT RECEIPT, both directions. `base` is the discriminator: base=0 is a
4453            // fresh prime (this line existed before the lift), base>0 is a SUFFIX riding the
4454            // GEMM trunk — the arm this lane added. The declining twin below counts the other
4455            // direction, so a log that shows neither line is an instrument fault, not a pass.
4456            eprintln!(
4457                "[gemm-prime] ENGAGED t={t} base={base} seq_end={seq_end} chunks<={width} (GEMM trunk + grouped MoE)"
4458            );
4459            return Ok((logits, h_seed, hiddens));
4460        }
4461        if self.uses_sliding_gated_moe_program() {
4462            eprintln!(
4463                "[gemm-prime] WALK t={t} base={} seq_end={seq_end} (batched prime declined)",
4464                cache.pos
4465            );
4466        }
4467        if overlay.is_none()
4468            && let Some(out) = self.step35_prime_trows(e, tokens, cache)?
4469        {
4470            return Ok(out);
4471        }
4472        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
4473        // session cache — every chunk (including the first) takes the continuation arm
4474        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
4475        assert!(
4476            t >= PRIME_MIN_T,
4477            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
4478        );
4479        assert!(
4480            cache.pos + t <= cache.max_ctx,
4481            "prime_cache: prompt exceeds cache max_ctx"
4482        );
4483
4484        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
4485        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
4486        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
4487        // each chunk runs the full layer stack with transients sized to the chunk, appending its
4488        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
4489        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
4490        // exactly the state carry it was built for). Full-attn chunks after the first attend to
4491        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
4492        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
4493        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
4494        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
4495        if self.is_gemma4_e4b() || self.uses_gemma_program() {
4496            if self.is_gemma4_e4b() {
4497                if overlay.is_some() {
4498                    return Err(
4499                        "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
4500                    );
4501                }
4502                return self.gemma4_e4b_prime(e, tokens, cache);
4503            }
4504            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
4505            // An overlay takes the masked-prefill arm: image rows splice in unscaled
4506            // (gemma4.cpp:182 — embd batches skip the sqrt(n_embd) scale) and the image
4507            // spans become bidirectional attention islands (lane/gemma-vision).
4508            return self.gemma4_prime(e, tokens, cache, overlay);
4509        }
4510        if crate::pp::prime_pipe_on()
4511            && crate::pp::prime_pp_on()
4512            && !crate::pp::pp2_streams_off()
4513            && crate::pp::pp_cuts(self.layers.len())
4514                .is_some_and(|fence| matches!(fence.len(), 4 | 5))
4515        {
4516            crate::pp::pp_wave_on()
4517                .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
4518        }
4519        let ranges = prime_chunk_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
4520        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
4521        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
4522        // the prefill's ARITHMETIC, so two rigs with different values produced different
4523        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
4524        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
4525        // (VERDICT.md) — and it is NOT what docs originally said:
4526        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
4527        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
4528        //     output head), so growing a chunk cannot move an existing row's value.
4529        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
4530        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
4531        //     not describe our leak.
4532        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
4533        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
4534        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
4535        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
4536        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
4537        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
4538        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
4539        // the source — every row is in one numeric class, so the chunk size no longer steers
4540        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
4541        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
4542        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
4543        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
4544        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
4545        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
4546        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
4547        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
4548        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
4549        // across calls, the request still ends at the same absolute position, whatever the tick
4550        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
4551        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
4552        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
4553        // default. Read per call, not cached (the probe flips it in-process between arms). Never
4554        // on in a measured default run.
4555        // `seq_end` (and its MEMRA_PRIME_CALLLOCAL seam) is computed ONCE above the batched
4556        // entry so both prime arms read the identical request-absolute value.
4557        if ranges.len() == 1 {
4558            return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
4559        }
4560        // PIPELINED PP PRIME. PP-2 retains its independently-qualified two-stage schedule;
4561        // PP-3/4 require the explicit MEMRA_PP_WAVE=1 persistent-stage wavefront. The serial
4562        // split stays reachable through MEMRA_PRIME_PIPE=0 and is the exactness oracle.
4563        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
4564            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
4565                if overlay.is_some() {
4566                    return Err(
4567                        "vision embedding overlay + pipelined PP prime unsupported (v1); \
4568                         run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
4569                            .into(),
4570                    );
4571                }
4572                if crate::pp::pp_multi_stream_same_device() {
4573                    return Err(
4574                        "prime chunk pipeline refused with 2 stage streams on one device — \
4575                         that concurrent-stream placement remains quarantined by the deferred \
4576                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
4577                         the serial split."
4578                            .into(),
4579                    );
4580                }
4581                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
4582            }
4583            if let Some(fence) =
4584                crate::pp::pp_cuts(self.layers.len()).filter(|f| matches!(f.len(), 4 | 5))
4585            {
4586                let wave_on = crate::pp::pp_wave_on()
4587                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
4588                let stages = fence.len() - 1;
4589                if crate::pp::pp_wave_route_enabled(
4590                    wave_on,
4591                    crate::pp::pp2_overlap(),
4592                    stages,
4593                    ranges.len(),
4594                ) {
4595                    if overlay.is_some() {
4596                        return Err(
4597                            "vision embedding overlay + pipelined PP prime unsupported; \
4598                             run the serial prime (MEMRA_PP_WAVE=0 or MEMRA_PRIME_PIPE=0)"
4599                                .into(),
4600                        );
4601                    }
4602                    let rt = crate::pp::PpNRt::get(e)?;
4603                    let double_slot = crate::pp::pp2_overlap();
4604                    crate::pp::pp_wave_eligibility(
4605                        stages,
4606                        double_slot,
4607                        rt.host_bounce_active(),
4608                        rt.repeated_stage_device(),
4609                    )
4610                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
4611                    return self
4612                        .prime_cache_ppn_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
4613                }
4614            }
4615        }
4616        let mut hiddens = e.uninit(t * n_embd)?;
4617        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
4618        for &(start, end) in &ranges {
4619            // chunked prime writes tap rows at the chunk's absolute offset
4620            if let Some(taps) = cache.dflash_taps.as_mut() {
4621                taps.base = start;
4622            }
4623            let (l, hs, x) =
4624                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
4625            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
4626            last = Some((l, hs));
4627            // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
4628            // so the device finished it. Stamp the odometer /health reads, so a BUSY
4629            // worker mid-long-prefill is never mistaken for a wedged one.
4630            crate::progress::note_prime_rows(end - start);
4631        }
4632        let (logits, h_seed) = last.unwrap();
4633        Ok((logits, h_seed, hiddens))
4634    }
4635
4636    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
4637    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
4638    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
4639    /// norm, lm head, and caller hidden-stack copy as the serial split.
4640    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4641    fn prime_cache_pp2_pipelined(
4642        &self,
4643        e: &Engine,
4644        tokens: &[u32],
4645        cache: &mut Cache,
4646        seq_end: usize,
4647        ranges: &[(usize, usize)],
4648        fence: &[usize],
4649    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4650        debug_assert_eq!(fence.len(), 3);
4651        debug_assert!(ranges.len() >= 2);
4652        let rt = crate::pp::PpNRt::get(e)?;
4653        assert_eq!(
4654            rt.n_stages(),
4655            2,
4656            "prime pipeline requires exactly two PP stages"
4657        );
4658        let n_embd = self.cfg.n_embd as usize;
4659        let t = tokens.len();
4660        let initial_base = cache.pos;
4661        let caller_stream = e.stream();
4662
4663        // #87 reverse publication before any new stage allocation, then prewarm both
4664        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
4665        // after stage 1(N) is queued would synchronize that stream and erase the first
4666        // overlap on a two-chunk prompt.
4667        rt.fence_stages_behind(&caller_stream)?;
4668        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
4669        rt.prepare_overlap_slots(0, max_payload)?;
4670
4671        let mut hiddens = e.uninit(t * n_embd)?;
4672        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
4673        let mut stage_caches = PrimeCacheStages::new(cache, fence);
4674        let (cache0, cache1) = stage_caches.pp2_parts();
4675        let (first_start, first_end) = ranges[0];
4676        let mut slot = self.prime_pp2_stage0_enqueue(
4677            e,
4678            rt,
4679            &tokens[first_start..first_end],
4680            cache0,
4681            seq_end,
4682            fence,
4683            initial_base + first_start,
4684            true,
4685        )?;
4686        cache0.pos = initial_base + first_end;
4687
4688        for (i, &(start, end)) in ranges.iter().enumerate() {
4689            let base = initial_base + start;
4690            debug_assert_eq!(
4691                cache1.pos, base,
4692                "stage 1 must drain chunks in original position order"
4693            );
4694            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
4695                let next_base = initial_base + next_start;
4696                debug_assert_eq!(
4697                    cache0.pos, next_base,
4698                    "stage 0 must issue chunks in original position order"
4699                );
4700                let cache0_stage = &mut *cache0;
4701                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
4702                // on one host thread therefore serialize even if the calls are ordered as
4703                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
4704                // stage 1 consumes slot N while stage 0 produces slot N+1.
4705                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
4706                    let stage0 = scope.spawn(move || -> Result<usize, String> {
4707                        let next = self
4708                            .prime_pp2_stage0_enqueue(
4709                                e,
4710                                rt,
4711                                &tokens[next_start..next_end],
4712                                cache0_stage,
4713                                seq_end,
4714                                fence,
4715                                next_base,
4716                                true,
4717                            )
4718                            .map_err(|err| err.to_string())?;
4719                        cache0_stage.pos = initial_base + next_end;
4720                        Ok(next)
4721                    });
4722                    let x = self.prime_pp2_stage1_enqueue(
4723                        e,
4724                        rt,
4725                        slot,
4726                        end - start,
4727                        cache1,
4728                        seq_end,
4729                        fence,
4730                        base,
4731                        true,
4732                    )?;
4733                    let out = {
4734                        rt.bind_stage(1)?;
4735                        let _st1 = rt.enter(1);
4736                        let e1 = rt.engine(1, e);
4737                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
4738                    };
4739                    let next = match stage0.join() {
4740                        Ok(result) => {
4741                            result.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?
4742                        }
4743                        Err(payload) => std::panic::resume_unwind(payload),
4744                    };
4745                    Ok((out, Some(next)))
4746                })?
4747            } else {
4748                let x = self.prime_pp2_stage1_enqueue(
4749                    e,
4750                    rt,
4751                    slot,
4752                    end - start,
4753                    cache1,
4754                    seq_end,
4755                    fence,
4756                    base,
4757                    true,
4758                )?;
4759                let out = {
4760                    rt.bind_stage(1)?;
4761                    let _st1 = rt.enter(1);
4762                    let e1 = rt.engine(1, e);
4763                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
4764                };
4765                (out, None)
4766            };
4767
4768            rt.publish_to(1, &caller_stream)?;
4769            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
4770            last = Some((out.0, out.1));
4771            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4772            // FORWARD PROGRESS (memra#50): this chunk's logits are already host-side,
4773            // so the device finished it. Stamp the odometer /health reads, so a BUSY
4774            // worker mid-long-prefill is never mistaken for a wedged one.
4775            crate::progress::note_prime_rows(end - start);
4776
4777            if let Some(next) = next_slot {
4778                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
4779                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
4780                // Stage 0(N+1) is already queued before this wait is appended, so its
4781                // overlap with stage 1(N) is preserved.
4782                rt.fence_stages_behind(&caller_stream)?;
4783                slot = next;
4784            }
4785        }
4786
4787        debug_assert_eq!(cache0.pos, initial_base + t);
4788        debug_assert_eq!(cache1.pos, initial_base + t);
4789        let (logits, h_seed) = last.unwrap();
4790        stage_caches.commit();
4791        Ok((logits, h_seed, hiddens))
4792    }
4793
4794    /// PP-3/4 prime wavefront: one prompt microchunk is one wave. One scoped host worker owns each
4795    /// non-head stage for the whole walk; the caller thread owns the head stage. Forward boundary
4796    /// messages preserve wave order, while reverse exact-wave acknowledgements are sent only after
4797    /// downstream `rx` has recorded `ev_rx`. An upstream stage therefore cannot cycle back to either
4798    /// shared slot before that slot's current generation has a host-observed release point.
4799    ///
4800    /// PP-2 remains on its independently qualified scheduler above. This path is reachable only
4801    /// through MEMRA_PP_WAVE=1 and the topology gate.
4802    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4803    fn prime_cache_ppn_pipelined(
4804        &self,
4805        e: &Engine,
4806        tokens: &[u32],
4807        cache: &mut Cache,
4808        seq_end: usize,
4809        ranges: &[(usize, usize)],
4810        fence: &[usize],
4811    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4812        let stages = fence.len().saturating_sub(1);
4813        debug_assert!((3..=4).contains(&stages));
4814        debug_assert!(ranges.len() >= 2);
4815        let rt = crate::pp::PpNRt::get(e)?;
4816        assert_eq!(
4817            rt.n_stages(),
4818            stages,
4819            "prime wavefront PpNRt/fence stage mismatch"
4820        );
4821        let n_embd = self.cfg.n_embd as usize;
4822        let initial_base = cache.pos;
4823        let caller_stream = e.stream();
4824        let primary_context = crate::pp::PrimaryContextRestore::new(e);
4825
4826        rt.fence_stages_behind(&caller_stream)?;
4827        let max_payload = ranges
4828            .iter()
4829            .map(|(start, end)| (end - start) * n_embd)
4830            .max()
4831            .unwrap_or(0);
4832        for boundary in 0..stages - 1 {
4833            rt.prepare_overlap_slots(boundary, max_payload)?;
4834        }
4835
4836        let mut stage_caches = PrimeCacheStages::new(cache, fence);
4837        let waves: Vec<_> = ranges
4838            .iter()
4839            .map(|&(start, end)| PrimePpWave {
4840                start,
4841                end,
4842                tokens: &tokens[start..end],
4843            })
4844            .collect();
4845        let mut forward_senders = Vec::with_capacity(stages - 1);
4846        let mut forward_receivers = Vec::with_capacity(stages - 1);
4847        let mut release_senders = Vec::with_capacity(stages - 1);
4848        let mut release_receivers = Vec::with_capacity(stages - 1);
4849        for _ in 0..stages - 1 {
4850            let (forward_sender, forward_receiver) = std::sync::mpsc::channel();
4851            let (release_sender, release_receiver) = std::sync::mpsc::channel();
4852            forward_senders.push(Some(forward_sender));
4853            forward_receivers.push(Some(forward_receiver));
4854            release_senders.push(Some(release_sender));
4855            release_receivers.push(Some(release_receiver));
4856        }
4857        let mut stage_channels = Vec::with_capacity(stages - 1);
4858        for stage in 0..stages - 1 {
4859            stage_channels.push(Some(PrimePpStageChannels {
4860                incoming: (stage > 0).then(|| forward_receivers[stage - 1].take().unwrap()),
4861                release_upstream: (stage > 0).then(|| release_senders[stage - 1].take().unwrap()),
4862                outgoing: forward_senders[stage].take().unwrap(),
4863                released_downstream: release_receivers[stage].take().unwrap(),
4864            }));
4865        }
4866        let head_incoming = forward_receivers[stages - 2].take().unwrap();
4867        let head_release = release_senders[stages - 2].take().unwrap();
4868        // allow: one-shot composite type; naming it would hide the shape that matters here —
4869        // this is the head stage's per-wave output, indexed by wave.
4870        #[allow(clippy::type_complexity)]
4871        let mut results: Vec<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>> =
4872            std::iter::repeat_with(|| None).take(waves.len()).collect();
4873        let walk_result = std::thread::scope(|scope| -> Result<(), Box<dyn std::error::Error>> {
4874            let waves_ref = &waves;
4875            let mut handles = Vec::with_capacity(stages - 1);
4876            // `stage` is a STAGE ID, not merely an index: it indexes two different containers
4877            // (`stage_channels`, `stage_caches.stages()`) and is passed to the worker as the
4878            // stage it owns. An iterator form would keep only one of the three uses.
4879            #[allow(clippy::needless_range_loop)]
4880            for stage in 0..stages - 1 {
4881                let channels = stage_channels[stage].take().unwrap();
4882                let cache_state = &stage_caches.stages()[stage];
4883                handles.push(scope.spawn(move || -> Result<(), String> {
4884                    let result = self.prime_ppn_wave_worker(
4885                        e,
4886                        rt,
4887                        waves_ref,
4888                        cache_state,
4889                        channels.incoming.as_ref(),
4890                        channels.release_upstream.as_ref(),
4891                        &channels.outgoing,
4892                        &channels.released_downstream,
4893                        stage,
4894                        seq_end,
4895                        fence,
4896                        initial_base,
4897                    );
4898                    if let Err(error) = &result {
4899                        channels.notify_failure(&error.to_string());
4900                    }
4901                    result.map_err(|error| error.to_string())
4902                }));
4903            }
4904
4905            let mut head_panic = None;
4906            let head_result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(
4907                || -> Result<(), Box<dyn std::error::Error>> {
4908                    let mut head_cache = stage_caches.stages()[stages - 1]
4909                        .lock()
4910                        .map_err(|_| "prime PP head cache lock poisoned")?;
4911                    for (wave_index, wave) in waves_ref.iter().enumerate() {
4912                        let incoming = recv_prime_pp_signal(
4913                            &head_incoming,
4914                            PrimePpWaveSlot {
4915                                wave: wave_index,
4916                                slot: 0,
4917                            },
4918                            false,
4919                            "prime PP head input",
4920                        )?;
4921                        results[wave_index] = Some(self.prime_ppn_wave_final(
4922                            e,
4923                            rt,
4924                            wave,
4925                            &mut head_cache,
4926                            incoming,
4927                            &head_release,
4928                            seq_end,
4929                            fence,
4930                            initial_base,
4931                        )?);
4932                        // FORWARD PROGRESS (memra#50): the head stage's host-side result for
4933                        // THIS wave exists here, which is the only point inside a multi-wave
4934                        // PP prime where progress is observable while the walk is still
4935                        // running. Stamping in the post-join copy loop instead would leave
4936                        // /health blind for the whole walk and then fire N stamps at once.
4937                        crate::progress::note_prime_rows(wave.end - wave.start);
4938                    }
4939                    Ok(())
4940                },
4941            )) {
4942                Ok(result) => result,
4943                Err(payload) => {
4944                    head_panic = Some(payload);
4945                    Err("prime PP head-stage host walker panicked".into())
4946                }
4947            };
4948            if let Err(error) = &head_result {
4949                let _ = head_release.send(PrimePpSignal::Error(error.to_string()));
4950            }
4951            let mut first_error = head_result.err().map(|error| error.to_string());
4952            let mut worker_panic = None;
4953            for handle in handles {
4954                match handle.join() {
4955                    Ok(Ok(())) => {}
4956                    Ok(Err(error)) => {
4957                        first_error.get_or_insert(error);
4958                    }
4959                    Err(payload) => {
4960                        if worker_panic.is_none() {
4961                            worker_panic = Some(payload);
4962                        }
4963                    }
4964                }
4965            }
4966            if let Some(payload) = head_panic {
4967                std::panic::resume_unwind(payload);
4968            }
4969            if let Some(payload) = worker_panic {
4970                std::panic::resume_unwind(payload);
4971            }
4972            if let Some(error) = first_error {
4973                return Err(error.into());
4974            }
4975            Ok(())
4976        });
4977        let publish_result = if walk_result.is_ok() {
4978            Some(rt.publish_to(stages - 1, &caller_stream))
4979        } else {
4980            None
4981        };
4982        let restore_result = primary_context.restore();
4983        walk_result?;
4984        if let Some(result) = publish_result {
4985            result?;
4986        }
4987        restore_result?;
4988        let mut hiddens = e.uninit(tokens.len() * n_embd)?;
4989        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
4990        for (wave_index, (wave, result)) in waves.iter().zip(results).enumerate() {
4991            debug_assert_eq!((wave.start, wave.end), ranges[wave_index]);
4992            let out = result.ok_or("prime PP wavefront completed without a head-stage result")?;
4993            e.copy_into(
4994                &mut hiddens,
4995                wave.start * n_embd,
4996                &out.2,
4997                (wave.end - wave.start) * n_embd,
4998            )?;
4999            last = Some((out.0, out.1));
5000            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5001            // The forward-progress stamp is NOT here: this loop runs after the wavefront has
5002            // been joined, so stamping it would fire every wave's progress at once at the END
5003            // of the walk and report nothing at all while the walk is running. It lives on the
5004            // head walker instead (memra#50, review of #106).
5005        }
5006        stage_caches.commit();
5007        drop(stage_caches);
5008
5009        static LOGGED: std::sync::Once = std::sync::Once::new();
5010        LOGGED.call_once(|| {
5011            eprintln!(
5012                "[pp-wave] PP{stages} prime wavefront engaged: microchunks={} \
5013                 (experimental, MEMRA_PP_WAVE=1)",
5014                ranges.len(),
5015            );
5016        });
5017        let (logits, h_seed) = last.expect("prime PP wavefront produced no microchunk");
5018        crate::pp::record_pp_wave_tick();
5019        Ok((logits, h_seed, hiddens))
5020    }
5021
5022    #[allow(clippy::too_many_arguments)]
5023    fn prime_ppn_wave_worker(
5024        &self,
5025        e: &Engine,
5026        rt: &crate::pp::PpNRt,
5027        waves: &[PrimePpWave<'_>],
5028        cache: &std::sync::Mutex<Cache>,
5029        incoming: Option<&std::sync::mpsc::Receiver<PrimePpSignal>>,
5030        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
5031        outgoing: &std::sync::mpsc::Sender<PrimePpSignal>,
5032        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
5033        stage: usize,
5034        seq_end: usize,
5035        fence: &[usize],
5036        initial_base: usize,
5037    ) -> Result<(), Box<dyn std::error::Error>> {
5038        debug_assert_eq!(incoming.is_some(), stage > 0);
5039        debug_assert_eq!(release_upstream.is_some(), stage > 0);
5040        let mut cache = cache
5041            .lock()
5042            .map_err(|_| "prime PP cache stage lock poisoned")?;
5043        let mut credits = PrimePpWaveCredits::default();
5044        for (wave_index, wave) in waves.iter().enumerate() {
5045            let incoming = match incoming {
5046                Some(receiver) => Some(recv_prime_pp_signal(
5047                    receiver,
5048                    PrimePpWaveSlot {
5049                        wave: wave_index,
5050                        slot: 0,
5051                    },
5052                    false,
5053                    "prime PP stage input",
5054                )?),
5055                None => None,
5056            };
5057            let sent = self.prime_ppn_wave_stage(
5058                e,
5059                rt,
5060                wave,
5061                &mut cache,
5062                stage,
5063                incoming,
5064                release_upstream,
5065                &mut credits,
5066                released_downstream,
5067                seq_end,
5068                fence,
5069                initial_base,
5070            )?;
5071            send_prime_pp_signal(outgoing, PrimePpSignal::Slot(sent), "prime PP stage output")?;
5072        }
5073        while let Some(expected) = credits.pending.front().copied() {
5074            let released = recv_prime_pp_signal(
5075                released_downstream,
5076                expected,
5077                true,
5078                "prime PP final slot release",
5079            )?;
5080            credits.record_release(released)?;
5081        }
5082        Ok(())
5083    }
5084
5085    #[allow(clippy::too_many_arguments)]
5086    fn prime_ppn_wave_stage(
5087        &self,
5088        e: &Engine,
5089        rt: &crate::pp::PpNRt,
5090        wave: &PrimePpWave<'_>,
5091        cache: &mut Cache,
5092        stage: usize,
5093        incoming: Option<PrimePpWaveSlot>,
5094        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
5095        credits: &mut PrimePpWaveCredits,
5096        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
5097        seq_end: usize,
5098        fence: &[usize],
5099        initial_base: usize,
5100    ) -> Result<PrimePpWaveSlot, Box<dyn std::error::Error>> {
5101        debug_assert!(stage + 1 < fence.len() - 1);
5102        let t = wave.end - wave.start;
5103        let base = initial_base + wave.start;
5104        debug_assert_eq!(cache.pos, base, "prime PP stage advanced out of order");
5105        let n_embd = self.cfg.n_embd as usize;
5106        let payload = t * n_embd;
5107        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
5108        rt.bind_stage(stage)?;
5109        let _stage = rt.enter(stage);
5110        let engine = rt.engine(stage, e);
5111        let positions_d = engine.htod_i32(&positions)?;
5112        let x = if stage == 0 {
5113            debug_assert!(incoming.is_none());
5114            self.embed(engine, wave.tokens)?
5115        } else {
5116            let incoming = incoming.ok_or("prime PP stage has no incoming boundary slot")?;
5117            let x = rt.rx(stage - 1, incoming.slot, payload)?;
5118            send_prime_pp_signal(
5119                release_upstream.ok_or("prime PP stage has no upstream release channel")?,
5120                PrimePpSignal::Slot(incoming),
5121                "prime PP upstream slot release",
5122            )?;
5123            x
5124        };
5125        let x = {
5126            let _wave_cell = crate::pp::enter_pp_wave_cell();
5127            let _overlap = crate::pp::enter_prime_pipe_stage();
5128            self.prime_layers(
5129                engine,
5130                x,
5131                fence[stage],
5132                fence[stage + 1],
5133                &positions_d,
5134                t,
5135                base,
5136                cache,
5137                seq_end,
5138            )?
5139        };
5140        if let Some(expected) = credits.release_required() {
5141            let released =
5142                recv_prime_pp_signal(released_downstream, expected, true, "prime PP slot credit")?;
5143            credits.record_release(released)?;
5144        }
5145        let sent = PrimePpWaveSlot {
5146            wave: credits.next_wave,
5147            slot: rt.tx_pipelined(stage, &x, payload)?,
5148        };
5149        credits.record_send(sent)?;
5150        cache.pos = base + t;
5151        Ok(sent)
5152    }
5153
5154    #[allow(clippy::too_many_arguments)]
5155    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5156    fn prime_ppn_wave_final(
5157        &self,
5158        e: &Engine,
5159        rt: &crate::pp::PpNRt,
5160        wave: &PrimePpWave<'_>,
5161        cache: &mut Cache,
5162        incoming: PrimePpWaveSlot,
5163        release_upstream: &std::sync::mpsc::Sender<PrimePpSignal>,
5164        seq_end: usize,
5165        fence: &[usize],
5166        initial_base: usize,
5167    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5168        let stage = fence.len() - 2;
5169        let t = wave.end - wave.start;
5170        let base = initial_base + wave.start;
5171        debug_assert_eq!(cache.pos, base, "prime PP head stage advanced out of order");
5172        let n_embd = self.cfg.n_embd as usize;
5173        let payload = t * n_embd;
5174        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
5175        rt.bind_stage(stage)?;
5176        let _stage = rt.enter(stage);
5177        let engine = rt.engine(stage, e);
5178        let positions_d = engine.htod_i32(&positions)?;
5179        let x = rt.rx(stage - 1, incoming.slot, payload)?;
5180        send_prime_pp_signal(
5181            release_upstream,
5182            PrimePpSignal::Slot(incoming),
5183            "prime PP head slot release",
5184        )?;
5185        let _wave_cell = crate::pp::enter_pp_wave_cell();
5186        let _overlap = crate::pp::enter_prime_pipe_stage();
5187        let x = self.prime_layers(
5188            engine,
5189            x,
5190            fence[stage],
5191            fence[stage + 1],
5192            &positions_d,
5193            t,
5194            base,
5195            cache,
5196            seq_end,
5197        )?;
5198        self.prime_chunk_epilogue(engine, x, t, cache)
5199    }
5200
5201    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
5202    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
5203    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
5204    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
5205    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
5206    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
5207    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
5208        if Engine::gdn_db_on()
5209            && Engine::gdn_chunked_enabled()
5210            && t >= 16
5211            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
5212            && num_k * 2 == num_v
5213        {
5214            num_k
5215        } else {
5216            num_v
5217        }
5218    }
5219
5220    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
5221    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
5222    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
5223    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
5224    fn f16out_on(e: &Engine, t: usize) -> bool {
5225        crate::f16_ffi::pp_f16_enabled()
5226            && t >= 16
5227            && !e.verify_exact_on()
5228            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
5229    }
5230
5231    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
5232    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
5233    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
5234    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
5235    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
5236    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
5237    /// see one entry, byte-identical behavior.
5238    pub fn prime_slabs_get(
5239        &self,
5240        e: &Engine,
5241        t: usize,
5242        n_embd: usize,
5243        n_ff_max: usize,
5244    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
5245        let mut slabs = self.prime_slabs.lock().unwrap();
5246        let dev = e.ctx().ordinal();
5247        let need_new = match slabs.get(&dev) {
5248            None => true,
5249            Some(sl) => sl.lock().unwrap().t_cap < t,
5250        };
5251        if need_new {
5252            slabs.insert(
5253                dev,
5254                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
5255                    t_cap: t,
5256                    h: e.uninit(t * n_embd)?,
5257                    x1: e.uninit(t * n_embd)?,
5258                    z: e.uninit(t * n_embd)?,
5259                    act: e.uninit(t * n_ff_max)?,
5260                    xa: e.uninit(t * n_embd)?,
5261                    xb: e.uninit(t * n_embd)?,
5262                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
5263                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
5264                    gate: e.uninit(t * n_ff_max)?,
5265                    up: e.uninit(t * n_ff_max)?,
5266                    ffn_out: e.uninit(t * n_embd)?,
5267                    seg_glue: Vec::new(),
5268                    mixed: e.uninit(t * n_embd)?,
5269                    seg_mid: Vec::new(),
5270                    seg_t: 0,
5271                })),
5272            );
5273        }
5274        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
5275    }
5276
5277    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
5278    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
5279    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
5280    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5281    fn prime_chunk(
5282        &self,
5283        e: &Engine,
5284        tokens: &[u32],
5285        cache: &mut Cache,
5286        seq_end: usize,
5287        chunk_off: usize,
5288        overlay: Option<&crate::vision::EmbedOverlay>,
5289    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5290        if crate::pp::pp_host_bounce_active()
5291            && (self.uses_gemma_program() || !crate::pp::prime_pp_on())
5292        {
5293            return Err(
5294                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
5295                 has no active prime stage split and would peer-read remote weights; keep \
5296                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
5297                    .into(),
5298            );
5299        }
5300        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
5301        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
5302        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
5303        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
5304        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
5305        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
5306        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
5307        // loader is off and there is nothing remote to split for.
5308        if !self.uses_gemma_program()
5309            && !crate::pp::pp2_streams_off()
5310            && crate::pp::prime_pp_on()
5311            && let Some(fence) = crate::pp::pp_cuts(self.layers.len())
5312        {
5313            if overlay.is_some() {
5314                return Err("vision embedding overlay + PP prime unsupported (v1); \
5315                         run single-device or MEMRA_PRIME_PP=0"
5316                    .into());
5317            }
5318            return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
5319        }
5320        if crate::pp::pp_host_bounce_active() {
5321            return Err(
5322                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
5323                 refusing an unsplit remote-weight walk"
5324                    .into(),
5325            );
5326        }
5327        let t = tokens.len();
5328        let base = cache.pos;
5329        debug_assert!(
5330            seq_end >= base + t,
5331            "prime_chunk: seq_end must cover this chunk"
5332        );
5333        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5334        let pos_d = e.htod_i32(&pos)?;
5335
5336        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
5337        if let Some(ov) = overlay {
5338            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
5339            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
5340            // Images larger than one prime chunk straddle boundaries, hence the clipping.
5341            //
5342            // Through the SHARED `EmbedOverlay::splice_into` (lane/glm53-vision-ppn): this
5343            // site used to carry its own byte-identical copy of that loop, which meant the
5344            // overlay residency law had to be re-stated per call site. One implementation,
5345            // one law, and the splice point cannot drift between arms.
5346            ov.splice_into(e, &mut x_embed, chunk_off, t, self.cfg.n_embd as usize)?;
5347        }
5348        let x = self.prime_layers(
5349            e,
5350            x_embed,
5351            0,
5352            self.layers.len(),
5353            &pos_d,
5354            t,
5355            base,
5356            cache,
5357            seq_end,
5358        )?;
5359        self.prime_chunk_epilogue(e, x, t, cache)
5360    }
5361
5362    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
5363    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
5364    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
5365    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
5366    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
5367    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
5368    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
5369    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
5370    ///     the plain add (materialize) and the next stage hoists its own first norm — the
5371    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
5372    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
5373    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
5374    ///     each stage walks through its own resident transients;
5375    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
5376    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
5377    #[allow(clippy::too_many_arguments)]
5378    fn prime_layers(
5379        &self,
5380        e: &Engine,
5381        x_in: CudaSlice<f32>,
5382        lo: usize,
5383        hi: usize,
5384        pos_d: &CudaSlice<i32>,
5385        t: usize,
5386        base: usize,
5387        cache: &mut Cache,
5388        seq_end: usize,
5389    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5390        let cfg = &self.cfg;
5391        let n_embd = cfg.n_embd as usize;
5392        let eps = cfg.rms_eps;
5393        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
5394        // standalone convert launches). Only when the f16 lane serves and T reaches the
5395        // GEMM tier; bit-identical either way.
5396        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
5397        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
5398        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
5399        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
5400        // live prefix is fully overwritten before use; x ping-pongs xa<->xb; the inactive
5401        // capacity tail must stay behind checked views. The hidden-stack return clones the
5402        // final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
5403        let n_ff_max = self
5404            .layers
5405            .iter()
5406            .map(|l| match &l.ffn {
5407                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
5408                _ => n_embd,
5409            })
5410            .max()
5411            .unwrap_or(n_embd)
5412            .max(n_embd);
5413        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
5414        let slab = if use_slabs {
5415            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
5416        } else {
5417            None
5418        };
5419        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
5420        let mut x_own; // fallback storage when slabs are off
5421        type SlabRefs<'a> = (
5422            &'a mut CudaSlice<f32>,
5423            &'a mut CudaSlice<f32>,
5424            &'a mut CudaSlice<f32>,
5425            &'a mut CudaSlice<f32>,
5426            &'a mut CudaSlice<u8>,
5427            &'a mut CudaSlice<u8>,
5428            &'a mut CudaSlice<f32>,
5429            &'a mut CudaSlice<f32>,
5430            &'a mut CudaSlice<f32>,
5431        );
5432        let (mut x_cur, mut x_nxt, sl): (
5433            &mut CudaSlice<f32>,
5434            &mut CudaSlice<f32>,
5435            Option<SlabRefs>,
5436        );
5437        #[allow(clippy::type_complexity)]
5438        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5439        let mut seg: Option<(
5440            &mut Vec<Option<cudarc::driver::CudaGraph>>,
5441            &mut Vec<Option<cudarc::driver::CudaGraph>>,
5442            &mut CudaSlice<f32>,
5443            &mut usize,
5444        )> = None;
5445        let mut x_own2;
5446        match slab_guard.as_mut() {
5447            Some(g) => {
5448                let slabs = &mut **g;
5449                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
5450                let PrimeSlabs {
5451                    xa,
5452                    xb,
5453                    h,
5454                    x1,
5455                    z,
5456                    act,
5457                    h16,
5458                    z16,
5459                    gate,
5460                    up,
5461                    ffn_out,
5462                    seg_glue,
5463                    mixed,
5464                    seg_mid,
5465                    seg_t,
5466                    ..
5467                } = slabs;
5468                x_cur = xa;
5469                x_nxt = xb;
5470                seg = Some((seg_glue, seg_mid, mixed, seg_t));
5471                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
5472            }
5473            None => {
5474                x_own = x_in;
5475                x_own2 = e.uninit(t * n_embd)?;
5476                x_cur = &mut x_own;
5477                x_nxt = &mut x_own2;
5478                sl = None;
5479            }
5480        }
5481        let mut alloc_h;
5482        let mut alloc_x1;
5483        let mut alloc_z;
5484        let mut alloc_act;
5485        let mut alloc_h16;
5486        let mut alloc_z16;
5487        let mut alloc_gate;
5488        let mut alloc_up;
5489        let mut alloc_fo;
5490        let (h, x1, z, act): (
5491            &mut CudaSlice<f32>,
5492            &mut CudaSlice<f32>,
5493            &mut CudaSlice<f32>,
5494            &mut CudaSlice<f32>,
5495        );
5496        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
5497        let (sl_gate, sl_up, sl_fo): (
5498            &mut CudaSlice<f32>,
5499            &mut CudaSlice<f32>,
5500            &mut CudaSlice<f32>,
5501        );
5502        match sl {
5503            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
5504                h = a;
5505                x1 = b;
5506                z = c;
5507                act = d;
5508                h16 = e16;
5509                z16 = f16b;
5510                sl_gate = g;
5511                sl_up = u;
5512                sl_fo = fo;
5513            }
5514            None => {
5515                alloc_h = e.uninit(t * n_embd)?;
5516                alloc_x1 = e.uninit(t * n_embd)?;
5517                alloc_z = e.uninit(t * n_embd)?;
5518                alloc_act = e.uninit(t * n_ff_max)?;
5519                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5520                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5521                alloc_gate = e.uninit(t * n_ff_max)?;
5522                alloc_up = e.uninit(t * n_ff_max)?;
5523                alloc_fo = e.uninit(t * n_embd)?;
5524                h = &mut alloc_h;
5525                x1 = &mut alloc_x1;
5526                z = &mut alloc_z;
5527                act = &mut alloc_act;
5528                h16 = &mut alloc_h16;
5529                z16 = &mut alloc_z16;
5530                sl_gate = &mut alloc_gate;
5531                sl_up = &mut alloc_up;
5532                sl_fo = &mut alloc_fo;
5533            }
5534        }
5535        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
5536        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
5537        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
5538        // first prime at this t (capture does not execute -> launch right after).
5539        let n_layers = self.layers.len();
5540        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
5541        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
5542        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
5543        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
5544        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
5545        // machinery stays (byte-identical) as their foundation.
5546        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
5547        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
5548        // step35 rides its own mixer through the normal per-layer arm below.
5549        let use_seg = f16fuse
5550            && seg.is_some()
5551            && !self.uses_sliding_gated_moe_program()
5552            && lo == 0
5553            && hi == n_layers
5554            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1")
5555            // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): below the
5556            // driver-free floor this prime call takes the eager fused else-arm (the
5557            // byte-identical twin the opt-in was gated against) instead of replaying
5558            // the S-mid/S-glue segment graphs into an exhausted card. Probe runs only
5559            // when the opt-in flag is armed (short-circuit order).
5560            && {
5561                let ok = crate::spec::graph_launch_headroom_ok(e);
5562                if !ok {
5563                    static NOTED: std::sync::Once = std::sync::Once::new();
5564                    NOTED.call_once(|| crate::spec::graph_replay_suspended_note("prime-seg"));
5565                }
5566                ok
5567            };
5568        if let Some((sg, sm, _, st)) = seg.as_mut()
5569            && **st != t
5570        {
5571            sg.clear();
5572            sg.extend((0..n_layers).map(|_| None));
5573            sm.clear();
5574            sm.extend((0..n_layers).map(|_| None));
5575            **st = t;
5576        }
5577        {
5578            let layer_lo = &self.layers[lo];
5579            if f16fuse {
5580                e.rms_norm_f16out(
5581                    x_cur,
5582                    layer_lo.attn_norm.float_data(),
5583                    h,
5584                    h16,
5585                    n_embd,
5586                    t,
5587                    eps,
5588                )?;
5589            } else {
5590                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
5591            }
5592        }
5593        let anat = Self::prime_anatomy_on();
5594        let mut anat_last = if anat {
5595            e.stream().synchronize()?;
5596            Some(std::time::Instant::now())
5597        } else {
5598            None
5599        };
5600        // Closes the region that just ENDED into `slot`, restarting the clock.
5601        macro_rules! anat_mark {
5602            ($slot:expr) => {
5603                if let Some(ts) = anat_last.as_mut() {
5604                    e.stream().synchronize()?;
5605                    Self::prime_anatomy_slots()[$slot].fetch_add(
5606                        ts.elapsed().as_nanos() as u64,
5607                        std::sync::atomic::Ordering::Relaxed,
5608                    );
5609                    *ts = std::time::Instant::now();
5610                }
5611            };
5612        }
5613        for il in lo..hi {
5614            let layer = &self.layers[il];
5615            let hx16 = if f16fuse { Some(&*h16) } else { None };
5616            if use_seg {
5617                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
5618                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
5619                let (pre, pre16, w_out) = match &layer.mixer {
5620                    Mixer::Full(fa) => {
5621                        let g3 = match hx16 {
5622                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
5623                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
5624                        };
5625                        let (pre, pre16) =
5626                            self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
5627                        (pre, pre16, &fa.wo)
5628                    }
5629                    Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("core-split prime"),
5630                    Mixer::Kda(_) => {
5631                        crate::hybrid::kda_path_unimplemented("core-split captured prime")
5632                    }
5633                    Mixer::Linear(la) => {
5634                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
5635                        let g4 = match hx16 {
5636                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
5637                            None => e.matmul_group(&ws, h, t)?,
5638                        };
5639                        let (pre, pre16) =
5640                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
5641                        (pre, pre16, &la.ssm_out)
5642                    }
5643                };
5644                {
5645                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
5646                    let pre_n = pre.len() / t;
5647                    let xh_pre = match pre16 {
5648                        Some(x) => x,
5649                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
5650                    };
5651                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
5652                        let y = e.matmul(w_out, &pre, t)?;
5653                        e.copy_into(mslab, 0, &y, t * n_embd)?;
5654                    }
5655                    if sm[il].is_none() {
5656                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
5657                        let w_post = layer.post_attn_norm.float_data();
5658                        e.stream().synchronize()?;
5659                        e.stream()
5660                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
5661                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
5662                            e.add(x_cur, mslab, x1, t * n_embd)?;
5663                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
5664                            Ok(())
5665                        })();
5666                        let g = e.stream().end_capture(
5667                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
5668                        r?;
5669                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
5670                    }
5671                    sm[il].as_ref().unwrap().launch()?;
5672                }
5673            } else {
5674                let mixed = match &layer.mixer {
5675                    Mixer::Full(fa) => {
5676                        let y =
5677                            self.full_attn_prime(e, fa, h, hx16, pos_d, t, cache, il, seq_end)?;
5678                        anat_mark!(0);
5679                        y
5680                    }
5681                    Mixer::Linear(la) => {
5682                        let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
5683                        anat_mark!(1);
5684                        y
5685                    }
5686                    Mixer::Mla(mla) => self.mla_attn_cached(e, mla, h, pos_d, t, il, cache)?,
5687                    Mixer::Kda(la) => {
5688                        let y = crate::kda::kda_prime_cached(e, la, h, t, eps, cache, il)?;
5689                        anat_mark!(1);
5690                        y
5691                    }
5692                };
5693                if f16fuse {
5694                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
5695                    // bit-identical) — the standalone add pass disappears.
5696                    e.add_rms_norm_f16out(
5697                        x_cur,
5698                        &mixed,
5699                        layer.post_attn_norm.float_data(),
5700                        x1,
5701                        z,
5702                        z16,
5703                        n_embd,
5704                        t,
5705                        eps,
5706                    )?;
5707                } else {
5708                    e.add(x_cur, &mixed, x1, t * n_embd)?;
5709                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
5710                }
5711                anat_mark!(4);
5712            }
5713            let zx16 = if f16fuse { Some(&*z16) } else { None };
5714            match &layer.ffn {
5715                crate::hybrid::Ffn::Dense {
5716                    ffn_gate,
5717                    ffn_up,
5718                    ffn_down,
5719                } => {
5720                    let n_ff = ffn_gate.out_features();
5721                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
5722                    // the allocating group + copy when a mirror is missing.
5723                    let mut into_ok = false;
5724                    if let Some(xh) = zx16 {
5725                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
5726                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
5727                    }
5728                    if !into_ok {
5729                        let mut g2 = match zx16 {
5730                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
5731                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
5732                        };
5733                        let up_y = g2.pop().unwrap();
5734                        let gate_y = g2.pop().unwrap();
5735                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
5736                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
5737                    }
5738                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
5739                    // operand in-epilogue; non-silu activations keep the standalone convert.
5740                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
5741                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
5742                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
5743                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
5744                    {
5745                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
5746                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
5747                        Some(a16)
5748                    } else {
5749                        Self::ffn_act_lim(
5750                            e,
5751                            &self.cfg,
5752                            sl_gate,
5753                            sl_up,
5754                            1.0,
5755                            1.0,
5756                            d_lim,
5757                            act,
5758                            t * n_ff,
5759                        )?;
5760                        None
5761                    };
5762                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
5763                    let xh_act = match act16 {
5764                        Some(x) => x,
5765                        None => e.f16_act(act, t * n_ff, n_ff)?,
5766                    };
5767                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
5768                        let y = e.matmul(ffn_down, &*act, t)?;
5769                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
5770                    }
5771                }
5772                crate::hybrid::Ffn::Moe(m) => {
5773                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
5774                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
5775                    anat_mark!(2);
5776                }
5777            }
5778            if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
5779                anat_mark!(3);
5780            }
5781            if use_seg && il + 1 < hi {
5782                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
5783                let w_next = self.layers[il + 1].attn_norm.float_data();
5784                let (sg, _, _, _) = seg.as_mut().unwrap();
5785                if sg[il].is_none() {
5786                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
5787                    e.stream().synchronize()?;
5788                    e.stream()
5789                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
5790                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
5791                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
5792                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
5793                        Ok(())
5794                    })();
5795                    let g = e.stream().end_capture(
5796                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
5797                    );
5798                    r?;
5799                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
5800                }
5801                sg[il].as_ref().unwrap().launch()?;
5802            } else {
5803                if il + 1 < hi {
5804                    let w_next = self.layers[il + 1].attn_norm.float_data();
5805                    if f16fuse {
5806                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
5807                    } else {
5808                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
5809                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
5810                    }
5811                } else {
5812                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
5813                }
5814            }
5815            anat_mark!(4);
5816            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
5817            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
5818            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
5819            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
5820            // unset (the default) costs one OnceLock read per layer.
5821            if let Some(path) = Self::prime_trace_path() {
5822                let row = base + t - 1;
5823                let host = e.dtoh(x_nxt)?;
5824                let last = &host[(t - 1) * n_embd..t * n_embd];
5825                use std::io::Write as _;
5826                let mut f = std::fs::OpenOptions::new()
5827                    .create(true)
5828                    .append(true)
5829                    .open(path)?;
5830                let mut h64: u64 = 0xcbf29ce484222325;
5831                for v in last {
5832                    h64 ^= v.to_bits() as u64;
5833                    h64 = h64.wrapping_mul(0x100000001b3);
5834                }
5835                writeln!(
5836                    f,
5837                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
5838                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
5839                    last[0], last[1], last[2]
5840                )?;
5841            }
5842            // dflash/dspark tap (no-op when no sink armed): post-layer residual rows for
5843            // drafter conditioning — the qwen twin of the gemma4 tap sites.
5844            self.dflash_tap(e, cache, il, x_nxt, t)?;
5845            std::mem::swap(&mut x_cur, &mut x_nxt);
5846        }
5847        if anat {
5848            let s = Self::prime_anatomy_slots();
5849            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
5850            eprintln!(
5851                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
5852                 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
5853                ms(0),
5854                ms(1),
5855                ms(2),
5856                ms(3),
5857                ms(4)
5858            );
5859        }
5860        // hidden-stack return: clone the final x out of the slab
5861        let mut x = e.uninit(t * n_embd)?;
5862        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
5863        drop(slab_guard);
5864        Ok(x)
5865    }
5866
5867    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
5868    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
5869    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
5870    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
5871    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5872    fn prime_chunk_epilogue(
5873        &self,
5874        e: &Engine,
5875        x: CudaSlice<f32>,
5876        t: usize,
5877        cache: &mut Cache,
5878    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5879        let n_embd = self.cfg.n_embd as usize;
5880        let eps = self.cfg.rms_eps;
5881        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
5882        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
5883        // the post-norm copy happens after hn exists).
5884        let mut h_seed = e.uninit(n_embd)?;
5885        if !crate::spec::spec_hpost() {
5886            e.copy_view_into(
5887                &mut h_seed,
5888                0,
5889                &x.slice((t - 1) * n_embd..t * n_embd),
5890                n_embd,
5891            )?;
5892        }
5893        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
5894        let mut hn = e.uninit(t * n_embd)?;
5895        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5896        if crate::spec::spec_hpost() {
5897            e.copy_view_into(
5898                &mut h_seed,
5899                0,
5900                &hn.slice((t - 1) * n_embd..t * n_embd),
5901                n_embd,
5902            )?;
5903        }
5904        let last = e.view(&hn, t * n_embd);
5905        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
5906        let mut hlast = e.uninit(n_embd)?;
5907        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
5908        let logits = e.matmul(&self.output, &hlast, 1)?;
5909        cache.pos += t;
5910        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
5911        // post-norm stack hn (MEMRA_SPEC_HPOST).
5912        Ok((
5913            e.dtoh(&logits)?,
5914            h_seed,
5915            if crate::spec::spec_hpost() { hn } else { x },
5916        ))
5917    }
5918
5919    /// Post-final-norm hidden state of one row of a prime-returned hidden stack — the
5920    /// embedding-pooling read (lane/embed-serve). `hiddens` is `prime_cache`'s third
5921    /// return: the pre-norm stack by default, but ALREADY post-norm under
5922    /// MEMRA_SPEC_HPOST (see `prime_chunk_epilogue`), so the norm is applied only in
5923    /// the default shape. Returns the host f32 row (`n_embd` wide).
5924    pub fn hidden_postnorm_row(
5925        &self,
5926        e: &Engine,
5927        hiddens: &CudaSlice<f32>,
5928        row: usize,
5929    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5930        let n_embd = self.cfg.n_embd as usize;
5931        let mut x1 = e.uninit(n_embd)?;
5932        e.copy_view_into(
5933            &mut x1,
5934            0,
5935            &hiddens.slice(row * n_embd..(row + 1) * n_embd),
5936            n_embd,
5937        )?;
5938        if crate::spec::spec_hpost() {
5939            return e.dtoh(&x1);
5940        }
5941        let mut hn = e.uninit(n_embd)?;
5942        e.rms_norm(
5943            &x1,
5944            self.output_norm.float_data(),
5945            &mut hn,
5946            n_embd,
5947            1,
5948            self.cfg.rms_eps,
5949        )?;
5950        e.dtoh(&hn)
5951    }
5952
5953    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
5954    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
5955    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
5956    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
5957    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
5958    /// prefill kernels. Structure mirrors the verify split exactly:
5959    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
5960    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
5961    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
5962    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
5963    ///                  there via the sharded loader) → `publish_to`
5964    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
5965    /// round's stage-freed buffers must not be reused under the caller's queued reads);
5966    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
5967    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
5968    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
5969    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
5970    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
5971    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
5972    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
5973    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
5974    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
5975    /// and its liveness counter is bumped here — the gate goes green with this function.
5976    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5977    fn prime_chunk_ppn(
5978        &self,
5979        e: &Engine,
5980        tokens: &[u32],
5981        cache: &mut Cache,
5982        seq_end: usize,
5983        fence: &[usize],
5984    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5985        let rt = crate::pp::PpNRt::get(e)?;
5986        let n_st = fence.len() - 1;
5987        assert_eq!(
5988            rt.n_stages(),
5989            n_st,
5990            "PpNRt stage count {} != fence stages {n_st}",
5991            rt.n_stages()
5992        );
5993        let n_embd = self.cfg.n_embd as usize;
5994        let t = tokens.len();
5995        let base = cache.pos;
5996        debug_assert!(
5997            seq_end >= base + t,
5998            "prime_chunk_ppn: seq_end must cover this chunk"
5999        );
6000        let payload = t * n_embd;
6001        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
6002        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
6003        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
6004        let caller_stream = e.stream();
6005        rt.fence_stages_behind(&caller_stream)?;
6006
6007        if n_st == 2 {
6008            let slot =
6009                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
6010            let x =
6011                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
6012            let out = {
6013                rt.bind_stage(1)?;
6014                let _st1 = rt.enter(1);
6015                let e1 = rt.engine(1, e);
6016                self.prime_chunk_epilogue(e1, x, t, cache)?
6017            };
6018            rt.publish_to(1, &caller_stream)?;
6019            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6020            // NO forward-progress stamp here (memra#50, review of #106): this is a per-CHUNK
6021            // body whose callers already stamp per chunk, and whose single-range case is
6022            // stamped by `prime_cache_overlaid`'s call-granularity shim. Stamping here too
6023            // would double-count `prime_progress.rows`/`chunks`, and those counts are
6024            // published as operator receipts and read by the hardware cell's PASS criterion.
6025            return Ok(out);
6026        }
6027
6028        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
6029
6030        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
6031        let mut slot = {
6032            let _st0 = rt.enter(0);
6033            let e0 = rt.engine(0, e);
6034            let pos_d = e0.htod_i32(&pos)?;
6035            let x = self.embed(e0, tokens)?;
6036            let x =
6037                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
6038            rt.tx(0, &x, payload)?
6039            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6040        };
6041
6042        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6043        for s in 1..n_st - 1 {
6044            let _st = rt.enter(s);
6045            let es = rt.engine(s, e);
6046            let pos_d = es.htod_i32(&pos)?;
6047            let x = rt.rx(s - 1, slot, payload)?;
6048            let x = self.prime_layers(
6049                es,
6050                x,
6051                fence[s],
6052                fence[s + 1],
6053                &pos_d,
6054                t,
6055                base,
6056                cache,
6057                seq_end,
6058            )?;
6059            slot = rt.tx(s, &x, payload)?;
6060        }
6061
6062        // ---- LAST STAGE: RX + final range + the shared epilogue ----
6063        let _stl = rt.enter(n_st - 1);
6064        let el = rt.engine(n_st - 1, e);
6065        let pos_d = el.htod_i32(&pos)?;
6066        let x = rt.rx(n_st - 2, slot, payload)?;
6067        let x = self.prime_layers(
6068            el,
6069            x,
6070            fence[n_st - 1],
6071            fence[n_st],
6072            &pos_d,
6073            t,
6074            base,
6075            cache,
6076            seq_end,
6077        )?;
6078        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
6079        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
6080        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
6081        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
6082        // stage stream host-side, but the law is stated in events, not in a dtoh side
6083        // effect a later deferred form would remove.
6084        rt.publish_to(n_st - 1, &caller_stream)?;
6085        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6086        Ok(out)
6087    }
6088
6089    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6090    fn prime_pp2_stage0_enqueue(
6091        &self,
6092        e: &Engine,
6093        rt: &crate::pp::PpNRt,
6094        tokens: &[u32],
6095        cache: &mut Cache,
6096        seq_end: usize,
6097        fence: &[usize],
6098        base: usize,
6099        pipelined: bool,
6100    ) -> Result<usize, Box<dyn std::error::Error>> {
6101        let t = tokens.len();
6102        let n_embd = self.cfg.n_embd as usize;
6103        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
6104        rt.bind_stage(0)?;
6105        let _st0 = rt.enter(0);
6106        let e0 = rt.engine(0, e);
6107        let pos_d = e0.htod_i32(&pos)?;
6108        let x = self.embed(e0, tokens)?;
6109        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
6110        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
6111        if pipelined {
6112            rt.tx_pipelined(0, &x, t * n_embd)
6113        } else {
6114            rt.tx(0, &x, t * n_embd)
6115        }
6116    }
6117
6118    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6119    fn prime_pp2_stage1_enqueue(
6120        &self,
6121        e: &Engine,
6122        rt: &crate::pp::PpNRt,
6123        slot: usize,
6124        t: usize,
6125        cache: &mut Cache,
6126        seq_end: usize,
6127        fence: &[usize],
6128        base: usize,
6129        pipelined: bool,
6130    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6131        let n_embd = self.cfg.n_embd as usize;
6132        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
6133        rt.bind_stage(1)?;
6134        let _st1 = rt.enter(1);
6135        let e1 = rt.engine(1, e);
6136        let pos_d = e1.htod_i32(&pos)?;
6137        let x = rt.rx(0, slot, t * n_embd)?;
6138        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
6139        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
6140    }
6141
6142    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
6143    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
6144    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
6145    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
6146    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
6147    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
6148    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
6149    /// bookkeeping still runs on the host per call — the real replay path moves the write
6150    /// slot to the len_d device counter (increment 3).
6151    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
6152    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
6153    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
6154    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
6155    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
6156    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
6157    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
6158    pub fn prime_chunk_captured(
6159        &self,
6160        e: &Engine,
6161        x_in: &CudaSlice<f32>,
6162        pos_d: &CudaSlice<i32>,
6163        t: usize,
6164        cache: &mut Cache,
6165        len_d: &CudaSlice<i32>,
6166        logits_out: &mut CudaSlice<f32>,
6167        h_seed_out: &mut CudaSlice<f32>,
6168    ) -> Result<(), Box<dyn std::error::Error>> {
6169        self.refuse_hyper("prime_chunk_captured")?;
6170        cache.ensure_usable("prime_chunk_captured")?;
6171        let cfg = &self.cfg;
6172        let n_embd = cfg.n_embd as usize;
6173        let eps = cfg.rms_eps;
6174        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
6175        let mut x = e.uninit(t * n_embd)?;
6176        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
6177        for (il, layer) in self.layers.iter().enumerate() {
6178            let mut h = e.uninit(t * n_embd)?;
6179            let mut hx16: Option<CudaSlice<u8>> = None;
6180            if f16fuse {
6181                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
6182                e.rms_norm_f16out(
6183                    &x,
6184                    layer.attn_norm.float_data(),
6185                    &mut h,
6186                    &mut b16,
6187                    n_embd,
6188                    t,
6189                    eps,
6190                )?;
6191                hx16 = Some(b16);
6192            } else {
6193                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6194            }
6195            let mixed = match &layer.mixer {
6196                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
6197                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
6198                // come from the caller (see step35_attn_pre_wo's doc note).
6199                Mixer::Full(fa) => {
6200                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
6201                }
6202                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("captured-graph prime"),
6203                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("captured prime chunk"),
6204                Mixer::Linear(la) => {
6205                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6206                    let g4 = match hx16.as_ref() {
6207                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
6208                        None => e.matmul_group(&ws, &h, t)?,
6209                    };
6210                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
6211                }
6212            };
6213            let mut x1 = e.uninit(t * n_embd)?;
6214            e.add(&x, &mixed, &mut x1, t * n_embd)?;
6215            let mut z = e.uninit(t * n_embd)?;
6216            let mut zx16: Option<CudaSlice<u8>> = None;
6217            if f16fuse {
6218                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
6219                e.rms_norm_f16out(
6220                    &x1,
6221                    layer.post_attn_norm.float_data(),
6222                    &mut z,
6223                    &mut b16,
6224                    n_embd,
6225                    t,
6226                    eps,
6227                )?;
6228                zx16 = Some(b16);
6229            } else {
6230                e.rms_norm(
6231                    &x1,
6232                    layer.post_attn_norm.float_data(),
6233                    &mut z,
6234                    n_embd,
6235                    t,
6236                    eps,
6237                )?;
6238            }
6239            let ffn_out = match &layer.ffn {
6240                crate::hybrid::Ffn::Dense {
6241                    ffn_gate,
6242                    ffn_up,
6243                    ffn_down,
6244                } => {
6245                    let n_ff = ffn_gate.out_features();
6246                    let mut g2 = match &zx16 {
6247                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
6248                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
6249                    };
6250                    let up = g2.pop().unwrap();
6251                    let gate = g2.pop().unwrap();
6252                    let mut act = e.uninit(t * n_ff)?;
6253                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
6254                    Self::ffn_act_lim(
6255                        e,
6256                        &self.cfg,
6257                        &gate,
6258                        &up,
6259                        1.0,
6260                        1.0,
6261                        self.cfg.clamp_shexp_at(il as u32),
6262                        &mut act,
6263                        t * n_ff,
6264                    )?;
6265                    e.matmul(ffn_down, &act, t)?
6266                }
6267                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
6268            };
6269            let mut x2 = e.uninit(t * n_embd)?;
6270            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6271            x = x2;
6272        }
6273        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
6274        if !crate::spec::spec_hpost() {
6275            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
6276        }
6277        let mut hn = e.uninit(t * n_embd)?;
6278        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6279        if crate::spec::spec_hpost() {
6280            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
6281        }
6282        let mut hlast = e.uninit(n_embd)?;
6283        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
6284        let logits = e.matmul(&self.output, &hlast, 1)?;
6285        let nv = logits.len();
6286        e.copy_into(logits_out, 0, &logits, nv)?;
6287        Ok(())
6288    }
6289
6290    fn step35_prime_batch_on() -> bool {
6291        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
6292    }
6293
6294    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
6295    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
6296    #[allow(clippy::too_many_arguments)]
6297    /// `seq_ends[s]`: sequence s's REQUEST-absolute end position — NOT `ts[s]`. It is the
6298    /// only thing step35's SWA arm keys on, so a chunk-local value here decides the attention
6299    /// kernel from the chunk size (and, below the 512-row window at a nonzero base, drops the
6300    /// window mask entirely). See the batched entry's note in `prime_cache_overlaid`.
6301    #[allow(clippy::too_many_arguments)]
6302    fn step35_prime_batch_layers(
6303        &self,
6304        e: &Engine,
6305        mut x: CudaSlice<f32>,
6306        lo: usize,
6307        hi: usize,
6308        ts: &[usize],
6309        offs: &[usize],
6310        seq_ends: &[usize],
6311        pos_ds: &[CudaSlice<i32>],
6312        caches: &mut [&mut Cache],
6313    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6314        let cfg = &self.cfg;
6315        let n_embd = cfg.n_embd as usize;
6316        let eps = cfg.rms_eps;
6317        let b = ts.len();
6318        let total: usize = ts.iter().sum();
6319        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
6320
6321        let split = |e: &Engine,
6322                     y: &CudaSlice<f32>,
6323                     dim: usize|
6324         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
6325            let mut out = Vec::with_capacity(b);
6326            for s in 0..b {
6327                let mut ys = e.uninit(ts[s] * dim)?;
6328                e.copy_view_into(
6329                    &mut ys,
6330                    0,
6331                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
6332                    ts[s] * dim,
6333                )?;
6334                out.push(ys);
6335            }
6336            Ok(out)
6337        };
6338
6339        // MEMRA_PRIME_PROF=1: per-phase wall inside the prime, sync-bounded (absolute time
6340        // inflates; the SPLIT is the signal). Two inspection passes failed to find where a
6341        // 3.8 s/4096-token chunk goes against a ~0.55 s compute budget, and nsys cannot capture
6342        // through the server's worker, so the walk measures itself.
6343        let prof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
6344        let mut ph = [0f64; 4]; // 0 norm+qkv, 1 attn, 2 o_proj+norm, 3 moe
6345        let mark = |e: &Engine, acc: usize, t0: &mut std::time::Instant, ph: &mut [f64; 4]| {
6346            if prof {
6347                let _ = e.stream().synchronize();
6348                ph[acc] += t0.elapsed().as_secs_f64() * 1e3;
6349                *t0 = std::time::Instant::now();
6350            }
6351        };
6352        let mut pt = std::time::Instant::now();
6353        for il in lo..hi {
6354            let layer = &self.layers[il];
6355            let Mixer::Full(fa) = &layer.mixer else {
6356                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
6357            };
6358
6359            let mut h = e.uninit(total * n_embd)?;
6360            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6361            if f16fuse {
6362                e.rms_norm_f16out(
6363                    &x,
6364                    layer.attn_norm.float_data(),
6365                    &mut h,
6366                    &mut hx16,
6367                    n_embd,
6368                    total,
6369                    eps,
6370                )?;
6371            } else {
6372                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
6373            }
6374
6375            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
6376            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
6377            // application stay verbatim.
6378            let gate_w = fa
6379                .attn_gate
6380                .as_ref()
6381                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
6382            let mut g4 = if f16fuse {
6383                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
6384            } else {
6385                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
6386            };
6387            let gate = g4.pop().unwrap();
6388            let mut parts: Vec<Vec<CudaSlice<f32>>> =
6389                (0..b).map(|_| Vec::with_capacity(3)).collect();
6390            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
6391                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
6392                    parts[s].push(ys);
6393                }
6394            }
6395            let gates = split(e, &gate, gate_w.out_features())?;
6396            let geometry = self.step35_geom(il);
6397            let hd = geometry.head_dim_k as usize;
6398            let nh = geometry.n_head as usize;
6399            let mut ag_cat = e.uninit(total * nh * hd)?;
6400            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
6401                mark(e, 0, &mut pt, &mut ph);
6402                let ag = self.step35_attn_pre_wo(
6403                    e,
6404                    fa,
6405                    g3s,
6406                    None,
6407                    Some(&gate),
6408                    &pos_ds[s],
6409                    ts[s],
6410                    Some(&mut *caches[s]),
6411                    il,
6412                    seq_ends[s],
6413                )?;
6414                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
6415            }
6416            mark(e, 1, &mut pt, &mut ph);
6417            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
6418
6419            let mut x1 = e.uninit(total * n_embd)?;
6420            let mut z = e.uninit(total * n_embd)?;
6421            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6422            if f16fuse {
6423                e.add_rms_norm_f16out(
6424                    &x,
6425                    &mixed,
6426                    layer.post_attn_norm.float_data(),
6427                    &mut x1,
6428                    &mut z,
6429                    &mut zx16,
6430                    n_embd,
6431                    total,
6432                    eps,
6433                )?;
6434            } else {
6435                e.add(&x, &mixed, &mut x1, total * n_embd)?;
6436                e.rms_norm(
6437                    &x1,
6438                    layer.post_attn_norm.float_data(),
6439                    &mut z,
6440                    n_embd,
6441                    total,
6442                    eps,
6443                )?;
6444            }
6445
6446            mark(e, 2, &mut pt, &mut ph);
6447            let ffn_out = match &layer.ffn {
6448                crate::hybrid::Ffn::Dense {
6449                    ffn_gate,
6450                    ffn_up,
6451                    ffn_down,
6452                } => {
6453                    let n_ff = ffn_gate.out_features();
6454                    let mut g2 = if f16fuse {
6455                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
6456                    } else {
6457                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
6458                    };
6459                    let up = g2.pop().unwrap();
6460                    let gate = g2.pop().unwrap();
6461                    let mut act = e.uninit(total * n_ff)?;
6462                    let d_lim = cfg.clamp_shexp_at(il as u32);
6463                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
6464                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
6465                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
6466                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
6467                            Some(y) => y,
6468                            None => e.matmul(ffn_down, &act, total)?,
6469                        }
6470                    } else {
6471                        Self::ffn_act_lim(
6472                            e,
6473                            cfg,
6474                            &gate,
6475                            &up,
6476                            1.0,
6477                            1.0,
6478                            d_lim,
6479                            &mut act,
6480                            total * n_ff,
6481                        )?;
6482                        e.matmul(ffn_down, &act, total)?
6483                    }
6484                }
6485                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
6486            };
6487            let mut x2 = e.uninit(total * n_embd)?;
6488            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
6489            x = x2;
6490            mark(e, 3, &mut pt, &mut ph);
6491        }
6492        if prof {
6493            eprintln!(
6494                "[prime-prof] t={total} layers={} norm+qkv={:.0}ms attn={:.0}ms o_proj={:.0}ms moe={:.0}ms",
6495                hi - lo,
6496                ph[0],
6497                ph[1],
6498                ph[2],
6499                ph[3]
6500            );
6501        }
6502        Ok(x)
6503    }
6504
6505    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6506    fn step35_prime_batch_epilogue(
6507        &self,
6508        e: &Engine,
6509        x: CudaSlice<f32>,
6510        ts: &[usize],
6511        offs: &[usize],
6512        caches: &mut [&mut Cache],
6513    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6514        let n_embd = self.cfg.n_embd as usize;
6515        let total: usize = ts.iter().sum();
6516        let mut hn = e.uninit(total * n_embd)?;
6517        e.rms_norm(
6518            &x,
6519            self.output_norm.float_data(),
6520            &mut hn,
6521            n_embd,
6522            total,
6523            self.cfg.rms_eps,
6524        )?;
6525
6526        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
6527        let mut out = Vec::with_capacity(ts.len());
6528        for s in 0..ts.len() {
6529            let mut hidden = e.uninit(ts[s] * n_embd)?;
6530            e.copy_view_into(
6531                &mut hidden,
6532                0,
6533                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
6534                ts[s] * n_embd,
6535            )?;
6536            let last0 = (offs[s] + ts[s] - 1) * n_embd;
6537            let mut h_seed = e.uninit(n_embd)?;
6538            e.copy_view_into(
6539                &mut h_seed,
6540                0,
6541                &hidden_src.slice(last0..last0 + n_embd),
6542                n_embd,
6543            )?;
6544            // Exactness-first: the serial reference runs the output head at m=1.
6545            let mut hlast = e.uninit(n_embd)?;
6546            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6547            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
6548            caches[s].pos += ts[s];
6549            out.push((logits, h_seed, hidden));
6550        }
6551        Ok(out)
6552    }
6553
6554    /// `seq_ends[s]` = sequence s's REQUEST-absolute end position (`cache.pos + prompt_len
6555    /// + queued_after`, computed once before any chunk loop). Only step35's SWA arm reads it,
6556    /// and it must NOT be this chunk's own length: see the note on the batched entry in
6557    /// `prime_cache_overlaid` for the window the chunk-local value opened.
6558    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6559    fn step35_prime_cache_batch(
6560        &self,
6561        e: &Engine,
6562        prompts: &[&[u32]],
6563        caches: &mut [&mut Cache],
6564        seq_ends: &[usize],
6565    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6566        assert_eq!(
6567            seq_ends.len(),
6568            prompts.len(),
6569            "step35 batched prime: one seq_end per sequence"
6570        );
6571        validate_step_prime_batch_modes(
6572            step_tp_prefill_enabled()?,
6573            step_ep_grouped_prefill_enabled()?,
6574        )?;
6575        if crate::pp::pp_host_bounce_active()
6576            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
6577        {
6578            return Err(
6579                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
6580                 stage split; refusing an unsplit remote-weight walk"
6581                    .into(),
6582            );
6583        }
6584        if !Self::step35_prime_batch_on() {
6585            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
6586        }
6587        // Continuation chunks are admitted (positions above carry each sequence's base). The
6588        // remaining restriction is genuine: a CROSS-REQUEST batch mixing sequences at different
6589        // positions still needs per-request queued_after to place its KV, so B > 1 keeps the
6590        // fresh-prompt rule.
6591        if prompts.len() > 1 && caches.iter().any(|c| c.pos != 0) {
6592            return Err(
6593                "step35 batched prime supports continuation only at B=1; a cross-request batch \
6594                 at mixed positions requires per-request queued_after"
6595                    .into(),
6596            );
6597        }
6598
6599        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
6600        for &t in &ts {
6601            assert!(
6602                t >= PRIME_MIN_T,
6603                "step35 batched prime needs T >= {PRIME_MIN_T}"
6604            );
6605        }
6606        for (s, c) in caches.iter().enumerate() {
6607            // POS-INCLUSIVE, like the walk's assert: a continuation chunk's rows land at
6608            // c.pos.., so the fresh-only `ts[s] <= max_ctx` form under-checked it.
6609            assert!(
6610                c.pos + ts[s] <= c.max_ctx,
6611                "step35 batched prime exceeds cache max_ctx"
6612            );
6613            assert!(
6614                seq_ends[s] >= c.pos + ts[s],
6615                "step35 batched prime: seq_end must cover this chunk"
6616            );
6617        }
6618        let mut transaction = CacheTaintGuard::arm(caches);
6619        // MEMRA_STEP35_PRIME_BATCH_TSEND=1: CANARY SEAM restoring the pre-fix chunk-local
6620        // `seq_end` (this chunk's own length, which `ts[s]` used to supply here). It is suffix-
6621        // and chunk-VARIANT by construction, so the suffix byte-identity gate MUST break under
6622        // it. That is how the defect is DEMONSTRATED rather than argued: one binary, one seam,
6623        // the legacy arm fails cold-vs-rewound identity and the default arm passes. Read per
6624        // call; never on in a measured default run.
6625        let legacy_tsend = std::env::var("MEMRA_STEP35_PRIME_BATCH_TSEND").as_deref() == Ok("1");
6626        let seq_ends_eff: Vec<usize> = if legacy_tsend {
6627            ts.clone()
6628        } else {
6629            seq_ends.to_vec()
6630        };
6631        let offs: Vec<usize> = ts
6632            .iter()
6633            .scan(0usize, |a, &t| {
6634                let o = *a;
6635                *a += t;
6636                Some(o)
6637            })
6638            .collect();
6639        let total: usize = ts.iter().sum();
6640        let payload = total * self.cfg.n_embd as usize;
6641        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
6642        // Positions start at each sequence's CURRENT cache position, not 0, so this entry can
6643        // prime a continuation chunk. The attention core already supports it: step35_attn_pre_wo
6644        // with Some(cache) is PRIME mode — it appends this chunk's post-rope K / raw V and
6645        // attends THROUGH the cache view — so only the hardcoded 0..t and the guard below ever
6646        // restricted it to fresh prompts.
6647        let positions: Vec<Vec<i32>> = ts
6648            .iter()
6649            .zip(caches.iter())
6650            .map(|(&t, c)| {
6651                let base = c.pos as i32;
6652                (0..t as i32).map(|i| base + i).collect()
6653            })
6654            .collect();
6655        let upload_positions =
6656            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
6657                positions
6658                    .iter()
6659                    .map(|p| e.htod_i32(p))
6660                    .collect::<Result<_, _>>()
6661            };
6662
6663        static ONCE: std::sync::Once = std::sync::Once::new();
6664        ONCE.call_once(|| {
6665            eprintln!(
6666                "[step35-prime-batch] first concat prime: B={} tokens={total}",
6667                prompts.len()
6668            );
6669        });
6670
6671        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
6672            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
6673                let rt = crate::pp::PpNRt::get(e)?;
6674                let n_st = fence.len() - 1;
6675                assert_eq!(
6676                    rt.n_stages(),
6677                    n_st,
6678                    "step35 prime batch stage count mismatch"
6679                );
6680                let caller_stream = e.stream();
6681                rt.fence_stages_behind(&caller_stream)?;
6682
6683                let mut slot = {
6684                    let _st0 = rt.enter(0);
6685                    let e0 = rt.engine(0, e);
6686                    let pos_ds = upload_positions(e0)?;
6687                    let x = self.embed(e0, &cat_tokens)?;
6688                    let x = self.step35_prime_batch_layers(
6689                        e0,
6690                        x,
6691                        fence[0],
6692                        fence[1],
6693                        &ts,
6694                        &offs,
6695                        &seq_ends_eff,
6696                        &pos_ds,
6697                        caches,
6698                    )?;
6699                    rt.tx(0, &x, payload)?
6700                };
6701                for s in 1..n_st - 1 {
6702                    let _st = rt.enter(s);
6703                    let es = rt.engine(s, e);
6704                    let pos_ds = upload_positions(es)?;
6705                    let x = rt.rx(s - 1, slot, payload)?;
6706                    let x = self.step35_prime_batch_layers(
6707                        es,
6708                        x,
6709                        fence[s],
6710                        fence[s + 1],
6711                        &ts,
6712                        &offs,
6713                        &seq_ends_eff,
6714                        &pos_ds,
6715                        caches,
6716                    )?;
6717                    slot = rt.tx(s, &x, payload)?;
6718                }
6719
6720                let _stl = rt.enter(n_st - 1);
6721                let el = rt.engine(n_st - 1, e);
6722                let pos_ds = upload_positions(el)?;
6723                let x = rt.rx(n_st - 2, slot, payload)?;
6724                let x = self.step35_prime_batch_layers(
6725                    el,
6726                    x,
6727                    fence[n_st - 1],
6728                    fence[n_st],
6729                    &ts,
6730                    &offs,
6731                    &seq_ends_eff,
6732                    &pos_ds,
6733                    caches,
6734                )?;
6735                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
6736                rt.publish_to(n_st - 1, &caller_stream)?;
6737                crate::pp::STEP35_PRIME_BATCH_SPLITS
6738                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6739                out
6740            } else {
6741                let pos_ds = upload_positions(e)?;
6742                let x = self.embed(e, &cat_tokens)?;
6743                let x = self.step35_prime_batch_layers(
6744                    e,
6745                    x,
6746                    0,
6747                    self.layers.len(),
6748                    &ts,
6749                    &offs,
6750                    &seq_ends_eff,
6751                    &pos_ds,
6752                    caches,
6753                )?;
6754                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
6755            }
6756        } else {
6757            let pos_ds = upload_positions(e)?;
6758            let x = self.embed(e, &cat_tokens)?;
6759            let x = self.step35_prime_batch_layers(
6760                e,
6761                x,
6762                0,
6763                self.layers.len(),
6764                &ts,
6765                &offs,
6766                &seq_ends_eff,
6767                &pos_ds,
6768                caches,
6769            )?;
6770            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
6771        };
6772        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6773        transaction.commit();
6774        Ok(out)
6775    }
6776
6777    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
6778    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
6779    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
6780    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
6781    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
6782    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
6783    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
6784    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
6785    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
6786    /// over the quantized past; Linear: the stateful pad_view twin — the same state
6787    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
6788    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
6789    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
6790    /// back to single-chunk serving).
6791    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
6792    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
6793    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6794    pub fn prime_cache_batch(
6795        &self,
6796        e: &Engine,
6797        prompts: &[&[u32]],
6798        caches: &mut [&mut Cache],
6799    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6800        // FORWARD PROGRESS (memra#50, review of #106). This entry does NOT route through
6801        // `prime_cache_overlaid`, so without this shim the multi-session batched wave prefill
6802        // (`worker.rs`'s interactive prefill tick) was the one prefill path with zero odometer
6803        // coverage: its unqualified-rewrite FALLBACK arm calls `prime_cache` per prompt and is
6804        // covered, while the fast batched arm was not. Same rule as the other entry: if
6805        // nothing below stamped, the call's own completion is the honest progress point.
6806        let events_before = crate::progress::events();
6807        let out = self.prime_cache_batch_inner(e, prompts, caches);
6808        if out.is_ok() && crate::progress::events() == events_before {
6809            crate::progress::note_prime_rows(prompts.iter().map(|p| p.len()).sum());
6810        }
6811        out
6812    }
6813
6814    #[allow(clippy::type_complexity)] // allow: mirrors `prime_cache_batch`'s signature
6815    fn prime_cache_batch_inner(
6816        &self,
6817        e: &Engine,
6818        prompts: &[&[u32]],
6819        caches: &mut [&mut Cache],
6820    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6821        self.refuse_hyper("prime_cache_batch")?;
6822        for cache in caches.iter() {
6823            cache.ensure_usable("prime_cache_batch")?;
6824        }
6825        if crate::pp::pp_cuts(self.layers.len()).is_some()
6826            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
6827        {
6828            return Err("pipeline rewrite is not qualified for batched prime".into());
6829        }
6830        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::CarriedPrime) {
6831            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
6832                return Err("neither batched-prime nor eager rewrite is qualified".into());
6833            }
6834            if prompts.len() != caches.len() {
6835                return Err("prime fallback prompt/cache shape mismatch".into());
6836            }
6837            static ONCE: std::sync::Once = std::sync::Once::new();
6838            ONCE.call_once(|| {
6839                eprintln!(
6840                    "[rewrite] carried-prime.v1 unqualified; using individual native eager primes"
6841                );
6842            });
6843            let mut transaction = CacheTaintGuard::arm(caches);
6844            let result: Result<Vec<_>, Box<dyn std::error::Error>> = prompts
6845                .iter()
6846                .copied()
6847                .zip(caches.iter_mut())
6848                .map(|(prompt, cache)| self.prime_cache(e, prompt, cache, 0))
6849                .collect();
6850            if result.is_ok() {
6851                transaction.commit();
6852            }
6853            return result;
6854        }
6855        let _pp_walk =
6856            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
6857                let rt = crate::pp::PpNRt::get(e)?;
6858                Some(rt.acquire_walk("prime_cache_batch")?)
6859            } else {
6860                None
6861            };
6862        let cfg = &self.cfg;
6863        let n_embd = cfg.n_embd as usize;
6864        let eps = cfg.rms_eps;
6865        let b = prompts.len();
6866        assert!(b >= 1 && b == caches.len());
6867        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
6868        let carried = pos0s.iter().any(|&p| p > 0);
6869        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
6870        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
6871        // generic concat attn core below (uniform geometry, no per-layer swa window, no
6872        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
6873        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
6874        if self.uses_gemma_program() {
6875            return Err(
6876                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
6877                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
6878                    .into(),
6879            );
6880        }
6881        // Step35 has a dedicated concat walk: the generic core below cannot express its
6882        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
6883        if self.uses_sliding_gated_moe_program() {
6884            // The cross-request driver hands whole requests (no chunk loop of its own), so each
6885            // sequence's request-absolute end IS its base plus its prompt length — the value
6886            // `ts[s]` happened to equal for the fresh B>=1 batches this caller admits, which is
6887            // why this arm is bit-for-bit unchanged by the seq_end threading.
6888            let seq_ends: Vec<usize> = caches
6889                .iter()
6890                .zip(prompts.iter())
6891                .map(|(c, p)| c.pos + p.len())
6892                .collect();
6893            return self.step35_prime_cache_batch(e, prompts, caches, &seq_ends);
6894        }
6895        if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
6896            let rt = crate::pp::PpNRt::get(e)?;
6897            if rt.cross_device() {
6898                return Err(
6899                    "prime_cache_batch: generic dense concat prime has no cross-device PP split; use individual prime_cache calls"
6900                        .into(),
6901                );
6902            }
6903        }
6904        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
6905        for &t in &ts {
6906            assert!(
6907                t >= PRIME_MIN_T,
6908                "prime_cache_batch needs T >= {PRIME_MIN_T}"
6909            );
6910        }
6911        for (s, c) in caches.iter().enumerate() {
6912            assert!(
6913                c.pos + ts[s] <= c.max_ctx,
6914                "prime_cache_batch: prompt exceeds cache max_ctx"
6915            );
6916        }
6917        let mut transaction = CacheTaintGuard::arm(caches);
6918        let total: usize = ts.iter().sum();
6919        let offs: Vec<usize> = ts
6920            .iter()
6921            .scan(0usize, |a, &t| {
6922                let o = *a;
6923                *a += t;
6924                Some(o)
6925            })
6926            .collect();
6927        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
6928        let pos_ds: Vec<CudaSlice<i32>> = ts
6929            .iter()
6930            .zip(&pos0s)
6931            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
6932            .collect::<Result<_, _>>()?;
6933        // split a concat [total, dim] buffer into per-seq copies
6934        let split = |e: &Engine,
6935                     y: &CudaSlice<f32>,
6936                     dim: usize|
6937         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
6938            let mut out = Vec::with_capacity(b);
6939            for s in 0..b {
6940                let mut ys = e.uninit(ts[s] * dim)?;
6941                e.copy_view_into(
6942                    &mut ys,
6943                    0,
6944                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
6945                    ts[s] * dim,
6946                )?;
6947                out.push(ys);
6948            }
6949            Ok(out)
6950        };
6951
6952        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
6953        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
6954        for (il, layer) in self.layers.iter().enumerate() {
6955            let mut h = e.uninit(total * n_embd)?;
6956            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6957            e.rms_norm_f16out(
6958                &x,
6959                layer.attn_norm.float_data(),
6960                &mut h,
6961                &mut hx16,
6962                n_embd,
6963                total,
6964                eps,
6965            )?;
6966            // mixer: projection GROUP on the concat (m = total), stateful core per seq
6967            let mut mixed = e.uninit(total * n_embd)?;
6968            match &layer.mixer {
6969                Mixer::Full(fa) => {
6970                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
6971                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
6972                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
6973                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
6974                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
6975                    // back to the per-seq dispatch.
6976                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
6977                    let (n_head, n_head_kv, head_dim) = (
6978                        geometry.n_head as usize,
6979                        geometry.n_head_kv as usize,
6980                        geometry.head_dim_k as usize,
6981                    );
6982                    let fa_scale = geometry.attention_scale();
6983                    let use_favl = !carried
6984                        && (2..=8).contains(&b)
6985                        && (head_dim == 256 || head_dim == 128)
6986                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
6987                        && std::env::var("MEMRA_NOFA").is_err()
6988                        && std::env::var("MEMRA_FA_FLOOR").is_err()
6989                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
6990                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
6991                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
6992                    if use_favl {
6993                        let (qf_w, kf_w, vf_w) = (
6994                            fa.wq.out_features(),
6995                            fa.wk.out_features(),
6996                            fa.wv.out_features(),
6997                        );
6998                        // Same bounds contract as `Engine::q_gate_split`, applied to the varlen
6999                        // twin's PER-TOKEN stride. `attn_pre_vl8` takes raw device pointers so it
7000                        // cannot check its own extents; `qf_w` is the wq out-features that set
7001                        // them, and `q_gate_split_vl` reads 2*head_dim per head out of it.
7002                        memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
7003                        struct APre {
7004                            q: CudaSlice<f32>,
7005                            gate: Option<CudaSlice<f32>>,
7006                            qn: CudaSlice<f32>,
7007                            kn: CudaSlice<f32>,
7008                        }
7009                        let mut aps = Vec::with_capacity(b);
7010                        for &t in ts.iter().take(b) {
7011                            aps.push(APre {
7012                                q: e.uninit(t * n_head * head_dim)?,
7013                                gate: Some(e.uninit(t * n_head * head_dim)?),
7014                                qn: e.uninit(t * n_head * head_dim)?,
7015                                kn: e.uninit(t * n_head_kv * head_dim)?,
7016                            });
7017                        }
7018                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
7019                            let kvl = caches[0].kv[il].as_ref().unwrap();
7020                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
7021                        };
7022                        let pargs: Vec<crate::AttnPreVl> = (0..b)
7023                            .map(|s| {
7024                                let (o, t) = (offs[s], ts[s]);
7025                                let kvl = caches[s].kv[il].as_ref().unwrap();
7026                                assert!(
7027                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
7028                                    "prime_cache_batch attn vl: fresh + capacity"
7029                                );
7030                                crate::AttnPreVl {
7031                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
7032                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
7033                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
7034                                    q: e.addr_f32(&aps[s].q),
7035                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
7036                                    qn: e.addr_f32(&aps[s].qn),
7037                                    kn: e.addr_f32(&aps[s].kn),
7038                                    kc: e.addr_u8(&kvl.k),
7039                                    vc: e.addr_u8(&kvl.v),
7040                                    t: t as i32,
7041                                    pad: 0,
7042                                }
7043                            })
7044                            .collect();
7045                        e.attn_pre_vl8(
7046                            &pargs,
7047                            fa.q_norm.float_data(),
7048                            fa.k_norm.float_data(),
7049                            head_dim,
7050                            geometry.n_rot as usize,
7051                            n_head,
7052                            n_head_kv,
7053                            self.cfg.rms_eps,
7054                            geometry.rope_base,
7055                            1.0,
7056                            kv_dim_k,
7057                            kv_dim_v,
7058                            ktb,
7059                            vtb,
7060                        )?;
7061                        for s in 0..b {
7062                            let kvl = caches[s].kv[il].as_mut().unwrap();
7063                            kvl.len += ts[s];
7064                            let new_len = kvl.len as i32;
7065                            e.set_i32_one(&mut kvl.len_d, new_len)?;
7066                        }
7067                        let mut attns = Vec::with_capacity(b);
7068                        let mut mirrors = Vec::with_capacity(b);
7069                        for &t in ts.iter().take(b) {
7070                            attns.push(e.uninit(t * n_head * head_dim)?);
7071                            let n = t * n_head_kv * head_dim;
7072                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
7073                        }
7074                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
7075                        // promoted single-seq config is on; else the mma favl.
7076                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
7077                            Ok("0") => false,
7078                            // Same refusal as the single-seq twin (lib.rs fa_prefill): the
7079                            // batched bf16 stage reaches func("f32_to_bf16_bulk"), absent on a
7080                            // portable build.
7081                            Ok("1") => {
7082                                crate::refuse_portable_force(
7083                                    "MEMRA_FA3=1",
7084                                    "the sm_90a fa3/bf16 kernels",
7085                                );
7086                                true
7087                            }
7088                            _ => cfg!(memra_hopper_mma),
7089                        };
7090                        if fa3_on {
7091                            let mut q16s = Vec::with_capacity(b);
7092                            let mut v16s = Vec::with_capacity(b);
7093                            for s in 0..b {
7094                                let t = ts[s];
7095                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
7096                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
7097                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
7098                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
7099                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
7100                                e.f32_to_bf16_v(
7101                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
7102                                    &mut v16,
7103                                    t * n_head_kv * head_dim,
7104                                )?;
7105                                q16s.push(q16);
7106                                v16s.push((k16, v16));
7107                            }
7108                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
7109                            let mut kp = qp;
7110                            let mut vp = qp;
7111                            let mut op = [core::ptr::null_mut::<f32>(); 8];
7112                            let mut tsv = [0i32; 8];
7113                            for s in 0..b {
7114                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
7115                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
7116                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
7117                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
7118                                tsv[s] = ts[s] as i32;
7119                            }
7120                            let rc = unsafe {
7121                                crate::fa3_vl_raw(
7122                                    qp.as_ptr(),
7123                                    kp.as_ptr(),
7124                                    vp.as_ptr(),
7125                                    op.as_ptr(),
7126                                    tsv.as_ptr(),
7127                                    b as i32,
7128                                    n_head as i32,
7129                                    n_head_kv as i32,
7130                                    head_dim as i32,
7131                                    fa_scale,
7132                                    e.stream().cu_stream() as *mut core::ffi::c_void,
7133                                )
7134                            };
7135                            if rc != 0 {
7136                                return Err(format!("memra_fa3_vl rc={rc}").into());
7137                            }
7138                        } else {
7139                            let fargs: Vec<crate::FaSeqVl> = (0..b)
7140                                .map(|s| crate::FaSeqVl {
7141                                    q: e.addr_f32(&aps[s].qn),
7142                                    k16: e.addr_u8(&mirrors[s].0),
7143                                    v16: e.addr_u8(&mirrors[s].1),
7144                                    o: e.addr_f32(&attns[s]),
7145                                    kf: e.addr_f32(&aps[s].kn),
7146                                    vf: e.addr_f32v(
7147                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
7148                                    ),
7149                                    t: ts[s] as i32,
7150                                    pad: 0,
7151                                })
7152                                .collect();
7153                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
7154                        }
7155                        for (s, attn) in attns.into_iter().enumerate() {
7156                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
7157                                e,
7158                                attn,
7159                                &aps[s].gate,
7160                                ts[s],
7161                                n_head,
7162                                head_dim,
7163                            )?;
7164                            let mut done = false;
7165                            if let Some(xh) = &ag16 {
7166                                done = e.try_f16_gemm_pre_into_off(
7167                                    &fa.wo,
7168                                    xh,
7169                                    ts[s],
7170                                    &mut mixed,
7171                                    offs[s] * n_embd,
7172                                )?;
7173                            }
7174                            if !done {
7175                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
7176                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
7177                            }
7178                        }
7179                    } else {
7180                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
7181                            (0..b).map(|_| Vec::new()).collect();
7182                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
7183                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
7184                                parts[s].push(ys);
7185                            }
7186                        }
7187                        for (s, g3s) in parts.into_iter().enumerate() {
7188                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
7189                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
7190                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
7191                            )?;
7192                            let mut done = false;
7193                            if let Some(xh) = &ag16 {
7194                                done = e.try_f16_gemm_pre_into_off(
7195                                    &fa.wo,
7196                                    xh,
7197                                    ts[s],
7198                                    &mut mixed,
7199                                    offs[s] * n_embd,
7200                                )?;
7201                            }
7202                            if !done {
7203                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
7204                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
7205                            }
7206                        }
7207                    }
7208                }
7209                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched cache prime"),
7210                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched prime"),
7211                Mixer::Linear(la) => {
7212                    // task #16: NO split copies (cores read row-offset views of the concat
7213                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
7214                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
7215                    // varlen K5 launch for all sequences.
7216                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
7217                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
7218                    let outs =
7219                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
7220                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
7221                        let (o, t) = (offs[s], ts[s]);
7222                        let mut done = false;
7223                        if let Some(xh) = &gn16 {
7224                            done = e.try_f16_gemm_pre_into_off(
7225                                &la.ssm_out,
7226                                xh,
7227                                t,
7228                                &mut mixed,
7229                                o * n_embd,
7230                            )?;
7231                        }
7232                        if !done {
7233                            let m = e.matmul(&la.ssm_out, &gn, t)?;
7234                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
7235                        }
7236                    }
7237                }
7238            }
7239            let mut x1 = e.uninit(total * n_embd)?;
7240            let mut z = e.uninit(total * n_embd)?;
7241            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
7242            e.add_rms_norm_f16out(
7243                &x,
7244                &mixed,
7245                layer.post_attn_norm.float_data(),
7246                &mut x1,
7247                &mut z,
7248                &mut zx16,
7249                n_embd,
7250                total,
7251                eps,
7252            )?;
7253            let ffn_out = match &layer.ffn {
7254                crate::hybrid::Ffn::Dense {
7255                    ffn_gate,
7256                    ffn_up,
7257                    ffn_down,
7258                } => {
7259                    let n_ff = ffn_gate.out_features();
7260                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
7261                    let up = g2.pop().unwrap();
7262                    let gate = g2.pop().unwrap();
7263                    let mut act = e.uninit(total * n_ff)?;
7264                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
7265                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
7266                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
7267                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
7268                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
7269                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
7270                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
7271                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
7272                            Some(y) => y,
7273                            None => e.matmul(ffn_down, &act, total)?,
7274                        }
7275                    } else {
7276                        Self::ffn_act_lim(
7277                            e,
7278                            &self.cfg,
7279                            &gate,
7280                            &up,
7281                            1.0,
7282                            1.0,
7283                            d_lim,
7284                            &mut act,
7285                            total * n_ff,
7286                        )?;
7287                        e.matmul(ffn_down, &act, total)?
7288                    }
7289                }
7290                crate::hybrid::Ffn::Moe(m) => {
7291                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
7292                }
7293            };
7294            let mut x2 = e.uninit(total * n_embd)?;
7295            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
7296            x = x2;
7297        }
7298        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
7299        let mut hn = e.uninit(total * n_embd)?;
7300        e.rms_norm(
7301            &x,
7302            self.output_norm.float_data(),
7303            &mut hn,
7304            n_embd,
7305            total,
7306            eps,
7307        )?;
7308        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
7309        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
7310        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
7311        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
7312        // argmax battery arbitrates, same as every other prefill GEMM change.
7313        let mut hcat = e.uninit(b * n_embd)?;
7314        for s in 0..b {
7315            let last0 = (offs[s] + ts[s] - 1) * n_embd;
7316            e.copy_view_into(
7317                &mut hcat,
7318                s * n_embd,
7319                &hn.slice(last0..last0 + n_embd),
7320                n_embd,
7321            )?;
7322        }
7323        let logits_cat = if b >= 2 {
7324            e.try_f16_gemm(&self.output, &hcat, b)?
7325        } else {
7326            None
7327        };
7328        let logits_host: Option<Vec<f32>> = match &logits_cat {
7329            Some(lc) => Some(e.dtoh(lc)?),
7330            None => None,
7331        };
7332        let n_vocab = self.output.out_features();
7333        let mut hidden_all = if crate::spec::spec_hpost() {
7334            split(e, &hn, n_embd)?
7335        } else {
7336            split(e, &x, n_embd)?
7337        };
7338        let mut out = Vec::with_capacity(b);
7339        for s in 0..b {
7340            let last0 = (offs[s] + ts[s] - 1) * n_embd;
7341            let mut h_seed = e.uninit(n_embd)?;
7342            if !crate::spec::spec_hpost() {
7343                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
7344            } else {
7345                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
7346            }
7347            let logits = match &logits_host {
7348                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
7349                None => {
7350                    let mut hlast = e.uninit(n_embd)?;
7351                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
7352                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
7353                }
7354            };
7355            caches[s].pos += ts[s];
7356            out.push((logits, h_seed, hidden_all.remove(0)));
7357        }
7358        transaction.commit();
7359        Ok(out)
7360    }
7361
7362    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
7363    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
7364    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
7365    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
7366    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
7367    ///
7368    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
7369    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
7370    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
7371    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
7372    #[allow(clippy::too_many_arguments)]
7373    fn full_attn_prime(
7374        &self,
7375        e: &Engine,
7376        fa: &FullAttnLayer,
7377        h: &CudaSlice<f32>,
7378        hx: Option<&CudaSlice<u8>>,
7379        pos_d: &CudaSlice<i32>,
7380        t: usize,
7381        cache: &mut Cache,
7382        il: usize,
7383        seq_end: usize,
7384    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7385        if self.uses_sliding_gated_moe_program() {
7386            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
7387        }
7388        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
7389        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
7390        // this single-seq path composes proj+core identically (byte-for-byte the old body).
7391        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
7392        let g3 = match hx {
7393            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
7394            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7395        };
7396        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
7397    }
7398
7399    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
7400    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
7401    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
7402    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7403    fn full_attn_prime_core(
7404        &self,
7405        e: &Engine,
7406        fa: &FullAttnLayer,
7407        g3: Vec<CudaSlice<f32>>,
7408        pos_d: &CudaSlice<i32>,
7409        t: usize,
7410        cache: &mut Cache,
7411        il: usize,
7412    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7413        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
7414        if let Some(xh) = &ag16
7415            && let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)?
7416        {
7417            return Ok(y);
7418        }
7419        e.matmul(&fa.wo, &attn_g, t)
7420    }
7421
7422    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7423    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7424    fn full_attn_prime_core_inner(
7425        &self,
7426        e: &Engine,
7427        fa: &FullAttnLayer,
7428        g3: Vec<CudaSlice<f32>>,
7429        pos_d: &CudaSlice<i32>,
7430        t: usize,
7431        cache: &mut Cache,
7432        il: usize,
7433    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7434        let cfg = &self.cfg;
7435        let geometry = cfg.full_attention_geometry_at(il as u32);
7436        let n_head = geometry.n_head as usize;
7437        let n_head_kv = geometry.n_head_kv as usize;
7438        let head_dim = geometry.head_dim_k as usize;
7439        let scale = geometry.attention_scale();
7440        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
7441        let AttnPre { q, k, v, gate } = pre;
7442        let mut attn = e.uninit(t * n_head * head_dim)?;
7443        self.full_attn_prime_fa_dispatch(
7444            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
7445        )?;
7446        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
7447    }
7448
7449    /// task #18 (attn side): projections tail through KV append — everything before the
7450    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
7451    /// present BEFORE this chunk's append (base_len; 0 == fresh).
7452    #[allow(clippy::type_complexity)]
7453    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7454    fn full_attn_prime_pre_fa(
7455        &self,
7456        e: &Engine,
7457        fa: &FullAttnLayer,
7458        mut g3: Vec<CudaSlice<f32>>,
7459        pos_d: &CudaSlice<i32>,
7460        t: usize,
7461        cache: &mut Cache,
7462        il: usize,
7463    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
7464        let cfg = &self.cfg;
7465        let geometry = cfg.full_attention_geometry_at(il as u32);
7466        let n_head = geometry.n_head as usize;
7467        let n_head_kv = geometry.n_head_kv as usize;
7468        let head_dim = geometry.head_dim_k as usize;
7469        let eps = cfg.rms_eps;
7470
7471        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
7472        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
7473        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
7474        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7475        let v = g3.pop().unwrap();
7476        let mut k = g3.pop().unwrap();
7477        let qf = g3.pop().unwrap();
7478        let (mut q, gate) = if gated {
7479            let mut q = e.uninit(t * n_head * head_dim)?;
7480            let mut gate = e.uninit(t * n_head * head_dim)?;
7481            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7482            (q, Some(gate))
7483        } else {
7484            (qf, None)
7485        };
7486
7487        let mut qn = e.uninit(t * n_head * head_dim)?;
7488        e.rms_norm(
7489            &q,
7490            fa.q_norm.float_data(),
7491            &mut qn,
7492            head_dim,
7493            n_head * t,
7494            eps,
7495        )?;
7496        q = qn;
7497        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7498        e.rms_norm(
7499            &k,
7500            fa.k_norm.float_data(),
7501            &mut kn,
7502            head_dim,
7503            n_head_kv * t,
7504            eps,
7505        )?;
7506        k = kn;
7507        let rope_dims = geometry.n_rot as usize;
7508        e.rope_neox(
7509            &mut q,
7510            pos_d,
7511            head_dim,
7512            rope_dims,
7513            n_head,
7514            t,
7515            geometry.rope_base,
7516            1.0,
7517        )?;
7518        e.rope_neox(
7519            &mut k,
7520            pos_d,
7521            head_dim,
7522            rope_dims,
7523            n_head_kv,
7524            t,
7525            geometry.rope_base,
7526            1.0,
7527        )?;
7528
7529        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
7530        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
7531        {
7532            let kvl = cache.kv[il].as_mut().unwrap();
7533            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
7534            e.append_kv_quantized_rows(
7535                &k,
7536                &v,
7537                &mut kvl.k,
7538                &mut kvl.v,
7539                kvl.len,
7540                t,
7541                kvl.kv_dim_k,
7542                kvl.kv_dim_v,
7543                kvl.k_tok_bytes,
7544                kvl.v_tok_bytes,
7545                crate::Engine::kv_fp8_on(),
7546            )?;
7547            kvl.len += t;
7548            let new_len = kvl.len as i32;
7549            e.set_i32_one(&mut kvl.len_d, new_len)?;
7550        }
7551
7552        let base_len = {
7553            let kvl = cache.kv[il].as_ref().unwrap();
7554            kvl.len - t // KV rows present BEFORE this chunk's append above
7555        };
7556        Ok((AttnPre { q, k, v, gate }, base_len))
7557    }
7558
7559    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
7560    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
7561    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
7562    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
7563    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
7564    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
7565    #[allow(clippy::too_many_arguments)]
7566    fn full_attn_prime_fa_dispatch(
7567        &self,
7568        e: &Engine,
7569        q: &CudaSlice<f32>,
7570        k: &CudaSlice<f32>,
7571        v: &CudaSlice<f32>,
7572        attn: &mut CudaSlice<f32>,
7573        base_len: usize,
7574        t: usize,
7575        cache: &mut Cache,
7576        il: usize,
7577        head_dim: usize,
7578        n_head: usize,
7579        n_head_kv: usize,
7580        scale: f32,
7581    ) -> Result<(), Box<dyn std::error::Error>> {
7582        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
7583        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
7584        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
7585        // attend through the quantized cache exactly like every later chunk (quantize-then-
7586        // attend). One numeric class for every row => the chunk size cannot decide where a
7587        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
7588        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
7589        // pin-the-boundary approach).
7590        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
7591        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
7592        // with the fix unconditional, only re-introducing the class edge can prove the gate
7593        // still detects the mechanism. Never on in a measured default run.
7594        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
7595            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
7596                e.sdpa_naive(
7597                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7598                )?;
7599            } else {
7600                e.fa_prefill(
7601                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7602                )?;
7603            }
7604            return Ok(());
7605        }
7606        let kvl = cache.kv[il].as_ref().unwrap();
7607        let t_kv = base_len + t;
7608        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7609        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7610        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
7611        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
7612        // same numeric class, so the uniform contract holds on the fallback too.
7613        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
7614            e.sdpa_naive_quantized_view(
7615                q,
7616                &k_view,
7617                &v_view,
7618                attn,
7619                head_dim,
7620                n_head,
7621                n_head_kv,
7622                t,
7623                t_kv,
7624                scale,
7625                true,
7626                kvl.k_tok_bytes,
7627                kvl.v_tok_bytes,
7628            )?;
7629            return Ok(());
7630        }
7631        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
7632        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
7633        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
7634        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
7635        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
7636        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
7637        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
7638        let deqw = std::env::var("MEMRA_PRIME_DEQW")
7639            .map(|v| v != "0")
7640            .unwrap_or(true);
7641        if deqw {
7642            e.fa_prefill_view_ws(
7643                q,
7644                &k_view,
7645                &v_view,
7646                attn,
7647                head_dim,
7648                n_head,
7649                n_head_kv,
7650                t,
7651                t_kv,
7652                scale,
7653                true,
7654                kvl.k_tok_bytes,
7655                kvl.v_tok_bytes,
7656                crate::Engine::kv_fp8_on(),
7657            )?;
7658        } else {
7659            e.fa_prefill_view(
7660                q,
7661                &k_view,
7662                &v_view,
7663                attn,
7664                head_dim,
7665                n_head,
7666                n_head_kv,
7667                t,
7668                t_kv,
7669                scale,
7670                true,
7671                kvl.k_tok_bytes,
7672                kvl.v_tok_bytes,
7673                crate::Engine::kv_fp8_on(),
7674            )?;
7675        }
7676        Ok(())
7677    }
7678
7679    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
7680    /// (bit-identical composition) and hands wo its fp16 operand directly.
7681    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7682    fn full_attn_prime_post_fa(
7683        &self,
7684        e: &Engine,
7685        attn: CudaSlice<f32>,
7686        gate: &Option<CudaSlice<f32>>,
7687        t: usize,
7688        n_head: usize,
7689        head_dim: usize,
7690    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7691        let (attn_g, ag16) = match gate {
7692            Some(gate) => {
7693                let n = t * n_head * head_dim;
7694                let mut ag = e.uninit(n)?;
7695                if Self::f16out_on(e, t) {
7696                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
7697                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
7698                    (ag, Some(a16))
7699                } else {
7700                    let mut gsig = e.uninit(n)?;
7701                    e.sigmoid(gate, &mut gsig, n)?;
7702                    e.mul(&attn, &gsig, &mut ag, n)?;
7703                    (ag, None)
7704                }
7705            }
7706            None => (attn, None),
7707        };
7708        Ok((attn_g, ag16))
7709    }
7710
7711    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
7712    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
7713    /// carried THROUGH the cache like the spec verify does: carried-ring conv
7714    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
7715    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
7716    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
7717    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
7718    fn linear_attn_prime(
7719        &self,
7720        e: &Engine,
7721        la: &LinearAttnLayer,
7722        h: &CudaSlice<f32>,
7723        hx: Option<&CudaSlice<u8>>,
7724        t: usize,
7725        cache: &mut Cache,
7726        il: usize,
7727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7728        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
7729        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
7730        let g4 = match hx {
7731            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
7732            None => e.matmul_group(&ws, h, t)?,
7733        };
7734        self.linear_attn_prime_core(e, la, g4, t, cache, il)
7735    }
7736
7737    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
7738    fn linear_attn_prime_core(
7739        &self,
7740        e: &Engine,
7741        la: &LinearAttnLayer,
7742        mut g4: Vec<CudaSlice<f32>>,
7743        t: usize,
7744        cache: &mut Cache,
7745        il: usize,
7746    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7747        self.linear_attn_prime_core_pad(e, la, std::mem::take(&mut g4), t, cache, il, None)
7748    }
7749
7750    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
7751    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
7752    /// conv ring writes back from the true tail. None = classic path, byte-identical.
7753    #[allow(clippy::too_many_arguments)]
7754    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7755    fn linear_attn_prime_core_pad_inner(
7756        &self,
7757        e: &Engine,
7758        la: &LinearAttnLayer,
7759        mut g4: Vec<CudaSlice<f32>>,
7760        t: usize,
7761        cache: &mut Cache,
7762        il: usize,
7763        pad_len: Option<&CudaSlice<i32>>,
7764    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7765        // shim over the view twin (task #16): full-range views of the owned buffers.
7766        let geometry = la.geometry;
7767        let d_state = geometry.key_head_dim as usize;
7768        let num_k = geometry.key_heads as usize;
7769        let num_v = geometry.value_heads as usize;
7770        let key_dim = d_state * num_k;
7771        let value_dim = geometry.value_head_dim as usize * num_v;
7772        let conv_dim = key_dim * 2 + value_dim;
7773        let alpha = g4.pop().unwrap(); // [T, num_v]
7774        let beta_raw = g4.pop().unwrap(); // [T, num_v]
7775        let z = g4.pop().unwrap(); // [T, value_dim]
7776        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
7777        self.linear_attn_prime_core_pad_view(
7778            e,
7779            la,
7780            &qkv_mixed.slice(0..t * conv_dim),
7781            &z.slice(0..t * value_dim),
7782            &beta_raw.slice(0..t * num_v),
7783            &alpha.slice(0..t * num_v),
7784            t,
7785            cache,
7786            il,
7787            pad_len,
7788        )
7789    }
7790
7791    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
7792    /// shared verbatim by the per-seq scan path and the varlen batched path.
7793    #[allow(clippy::too_many_arguments)]
7794    fn linear_attn_gdn_prep(
7795        &self,
7796        e: &Engine,
7797        la: &LinearAttnLayer,
7798        qkv_mixed: &cudarc::driver::CudaView<f32>,
7799        beta_raw: &cudarc::driver::CudaView<f32>,
7800        alpha: &cudarc::driver::CudaView<f32>,
7801        t: usize,
7802        cache: &mut Cache,
7803        il: usize,
7804        pad_len: Option<&CudaSlice<i32>>,
7805    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
7806        let cfg = &self.cfg;
7807        let geometry = la.geometry;
7808        let d_state = geometry.key_head_dim as usize;
7809        let num_k = geometry.key_heads as usize;
7810        let num_v = geometry.value_heads as usize;
7811        let d_conv = geometry.conv_kernel as usize;
7812        let key_dim = d_state * num_k; // 2048
7813        let value_dim = geometry.value_head_dim as usize * num_v;
7814        let conv_dim = key_dim * 2 + value_dim; // 8192
7815        let eps = cfg.rms_eps;
7816        debug_assert!(
7817            t >= d_conv - 1,
7818            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
7819        );
7820
7821        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
7822        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
7823        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
7824        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
7825        let rl = cache.recur[il].as_mut().unwrap();
7826        let hk = Self::gdn_hk(e, t, num_v, num_k);
7827        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
7828        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
7829        let mut q_g = e.uninit(d_state * hk * t)?;
7830        let mut k_g = e.uninit(d_state * hk * t)?;
7831        let mut v_g = e.uninit(d_state * num_v * t)?;
7832        if conv_fuse {
7833            e.ssm_conv1d_gdn_state_pad(
7834                qkv_mixed,
7835                &mut rl.conv_state,
7836                la.ssm_conv1d.float_data(),
7837                &mut q_g,
7838                &mut k_g,
7839                &mut v_g,
7840                conv_dim,
7841                t,
7842                d_conv,
7843                d_state,
7844                num_v,
7845                num_k,
7846                key_dim,
7847                hk,
7848                pad_len,
7849            )?;
7850        } else {
7851            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
7852            e.ssm_conv1d_tm_state_pad_v(
7853                qkv_mixed,
7854                &mut rl.conv_state,
7855                la.ssm_conv1d.float_data(),
7856                &mut conv_out,
7857                conv_dim,
7858                t,
7859                d_conv,
7860                pad_len,
7861            )?;
7862            e.qkv_to_gdn_repack(
7863                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7864            )?;
7865        }
7866        let mut q_l2 = e.uninit(d_state * hk * t)?;
7867        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
7868        // Emitted only where a consumer exists (the wgmma config) — on other arches the
7869        // alloc + epilogue stores would be pure waste.
7870        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
7871            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
7872            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
7873            Some(qb)
7874        } else {
7875            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
7876            None
7877        };
7878        let mut k_l2 = e.uninit(d_state * hk * t)?;
7879        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
7880        let kb16 = if Engine::l2_v2_on(d_state) {
7881            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
7882            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
7883            Some(kb)
7884        } else {
7885            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
7886            None
7887        };
7888        let mut beta = e.uninit(t * num_v)?;
7889        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
7890        let mut g_log = e.uninit(t * num_v)?;
7891        e.gdn_glog_v(
7892            alpha,
7893            la.ssm_dt.float_data(),
7894            la.ssm_a.float_data(),
7895            &mut g_log,
7896            num_v,
7897            t,
7898        )?;
7899        if let Some(len_d) = pad_len {
7900            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
7901        }
7902        Ok(GdnPrep {
7903            hk,
7904            q_l2,
7905            k_l2,
7906            v_g,
7907            beta,
7908            g_log,
7909            kb16,
7910            qb16,
7911        })
7912    }
7913
7914    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
7915    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
7916    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
7917    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
7918    #[allow(clippy::too_many_arguments)]
7919    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7920    fn linear_attn_prime_core_batch(
7921        &self,
7922        e: &Engine,
7923        la: &LinearAttnLayer,
7924        g4: &[CudaSlice<f32>],
7925        offs: &[usize],
7926        ts: &[usize],
7927        caches: &mut [&mut Cache],
7928        il: usize,
7929    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
7930        let geometry = la.geometry;
7931        let d_state = geometry.key_head_dim as usize;
7932        let num_k = geometry.key_heads as usize;
7933        let num_v = geometry.value_heads as usize;
7934        let d_conv = geometry.conv_kernel as usize;
7935        let key_dim = d_state * num_k;
7936        let value_dim = geometry.value_head_dim as usize * num_v;
7937        let conv_dim = key_dim * 2 + value_dim;
7938        let eps = self.cfg.rms_eps;
7939        let scale = 1.0 / (d_state as f32).sqrt();
7940        let b = ts.len();
7941        let c = Engine::gdn_chunk_size();
7942        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
7943        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
7944        let carried = caches.iter().any(|c| c.pos > 0);
7945        let use_vl = !carried
7946            && (2..=8).contains(&b)
7947            && Engine::gdn_chunked_enabled()
7948            && ts.iter().all(|&t| t >= 16)
7949            && e.gdn_mma_enabled(c)
7950            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
7951        if !use_vl {
7952            return (0..b)
7953                .map(|s| {
7954                    let (o, t) = (offs[s], ts[s]);
7955                    self.linear_attn_prime_core_pad_view(
7956                        e,
7957                        la,
7958                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
7959                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
7960                        &g4[2].slice(o * num_v..(o + t) * num_v),
7961                        &g4[3].slice(o * num_v..(o + t) * num_v),
7962                        t,
7963                        caches[s],
7964                        il,
7965                        None,
7966                    )
7967                })
7968                .collect();
7969        }
7970        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
7971        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
7972        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
7973        struct SeqBufs {
7974            conv_out: CudaSlice<f32>,
7975            q_g: CudaSlice<f32>,
7976            k_g: CudaSlice<f32>,
7977            v_g: CudaSlice<f32>,
7978            q_l2: CudaSlice<f32>,
7979            k_l2: CudaSlice<f32>,
7980            beta: CudaSlice<f32>,
7981            g_log: CudaSlice<f32>,
7982            gn: CudaSlice<f32>,
7983            gn16: CudaSlice<u8>,
7984        }
7985        let f16o = Self::f16out_on(e, 16);
7986        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
7987        let mut sb = Vec::with_capacity(b);
7988        let mut pres = Vec::with_capacity(b);
7989        for &t in ts.iter().take(b) {
7990            sb.push(SeqBufs {
7991                conv_out: e.uninit(conv_dim * t)?,
7992                q_g: e.uninit(d_state * hk * t)?,
7993                k_g: e.uninit(d_state * hk * t)?,
7994                v_g: e.uninit(d_state * num_v * t)?,
7995                q_l2: e.uninit(d_state * hk * t)?,
7996                k_l2: e.uninit(d_state * hk * t)?,
7997                beta: e.uninit(t * num_v)?,
7998                g_log: e.uninit(t * num_v)?,
7999                gn: e.uninit(d_state * num_v * t)?,
8000                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
8001            });
8002            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
8003        }
8004        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
8005            .map(|s| {
8006                let (o, t) = (offs[s], ts[s]);
8007                let rl = caches[s].recur[il].as_ref().unwrap();
8008                crate::GdnPrepVl {
8009                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
8010                    conv_state: e.addr_f32(&rl.conv_state),
8011                    conv_out: e.addr_f32(&sb[s].conv_out),
8012                    q_g: e.addr_f32(&sb[s].q_g),
8013                    k_g: e.addr_f32(&sb[s].k_g),
8014                    v_g: e.addr_f32(&sb[s].v_g),
8015                    q_l2: e.addr_f32(&sb[s].q_l2),
8016                    k_l2: e.addr_f32(&sb[s].k_l2),
8017                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
8018                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
8019                    beta: e.addr_f32(&sb[s].beta),
8020                    g_log: e.addr_f32(&sb[s].g_log),
8021                    o: e.addr_f32(&pres[s].o),
8022                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
8023                    gn: e.addr_f32(&sb[s].gn),
8024                    gn16: e.addr_u8(&sb[s].gn16),
8025                    kb16: if Engine::l2_v2_on(d_state) {
8026                        e.addr_u8(&pres[s].kb16)
8027                    } else {
8028                        0
8029                    },
8030                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
8031                        e.addr_u8(&pres[s].qb16)
8032                    } else {
8033                        0
8034                    },
8035                    t: t as i32,
8036                    pad: 0,
8037                }
8038            })
8039            .collect();
8040        let args: Vec<crate::GdnSeqVl> = (0..b)
8041            .map(|s| {
8042                let rl = caches[s].recur[il].as_ref().unwrap();
8043                crate::GdnSeqVl {
8044                    kb16: e.addr_u8(&pres[s].kb16),
8045                    gcum: e.addr_f32(&pres[s].gcum),
8046                    beta: e.addr_f32(&sb[s].beta),
8047                    u: e.addr_f32(&pres[s].u),
8048                    wb16: e.addr_u8(&pres[s].wb16),
8049                    y: e.addr_u8(&pres[s].y16),
8050                    ssnap: e.addr_u8(&pres[s].ssnap16),
8051                    state_in: e.addr_f32(&rl.ssm_state),
8052                    state_out: e.addr_f32(&rl.ssm_state_alt),
8053                    q: e.addr_f32(&sb[s].q_l2),
8054                    p: e.addr_f32(&pres[s].p),
8055                    o: e.addr_f32(&pres[s].o),
8056                    k: e.addr_f32(&sb[s].k_l2),
8057                    v: e.addr_f32(&sb[s].v_g),
8058                    g: e.addr_f32(&sb[s].g_log),
8059                    a: e.addr_f32(&pres[s].a),
8060                    w: e.addr_f32(&pres[s].w),
8061                    t: ts[s] as i32,
8062                    nc: pres[s].nc as i32,
8063                }
8064            })
8065            .collect();
8066        e.gdn_prep_vl8(
8067            &prep_args,
8068            la.ssm_conv1d.float_data(),
8069            la.ssm_dt.float_data(),
8070            la.ssm_a.float_data(),
8071            conv_dim,
8072            d_conv,
8073            d_state,
8074            num_v,
8075            num_k,
8076            key_dim,
8077            hk,
8078            eps,
8079        )?;
8080        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
8081        // both standalone mirror launches vanish on the default config.
8082        if !Engine::l2_v2_on(d_state) {
8083            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
8084        }
8085        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
8086        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
8087            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
8088            if !Engine::l2_v2_on(d_state) {
8089                for s in 0..b {
8090                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
8091                }
8092            }
8093            let mut wa = [crate::GdnWVl::default(); 8];
8094            for s in 0..b {
8095                wa[s] = crate::GdnWVl {
8096                    qb16: e.addr_u8(&pres[s].qb16),
8097                    pb16: e.addr_u8(&pres[s].pb16),
8098                };
8099            }
8100            Some(crate::GdnWVl8(wa))
8101        } else {
8102            None
8103        };
8104        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
8105        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
8106        if f16o {
8107            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
8108        }
8109        // per-seq state swap (+ non-f16out tail fallback)
8110        let mut out = Vec::with_capacity(b);
8111        for (s, bufs) in sb.into_iter().enumerate() {
8112            let rl = caches[s].recur[il].as_mut().unwrap();
8113            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8114            let (o, t) = (offs[s], ts[s]);
8115            let SeqBufs { mut gn, gn16, .. } = bufs;
8116            if f16o {
8117                out.push((gn, Some(gn16)));
8118            } else {
8119                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
8120                e.gated_rmsnorm_zv(
8121                    &pres[s].o,
8122                    la.ssm_norm.float_data(),
8123                    &z_v,
8124                    &mut gn,
8125                    d_state,
8126                    num_v * t,
8127                    eps,
8128                )?;
8129                out.push((gn, None));
8130            }
8131        }
8132        Ok(out)
8133    }
8134
8135    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
8136    /// views of the CONCAT projection outputs directly (no per-seq split copies).
8137    /// Same kernels, same values, byte-identical to the Vec shim above.
8138    #[allow(clippy::too_many_arguments)]
8139    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
8140    fn linear_attn_prime_core_pad_view(
8141        &self,
8142        e: &Engine,
8143        la: &LinearAttnLayer,
8144        qkv_mixed: &cudarc::driver::CudaView<f32>,
8145        z: &cudarc::driver::CudaView<f32>,
8146        beta_raw: &cudarc::driver::CudaView<f32>,
8147        alpha: &cudarc::driver::CudaView<f32>,
8148        t: usize,
8149        cache: &mut Cache,
8150        il: usize,
8151        pad_len: Option<&CudaSlice<i32>>,
8152    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
8153        let cfg = &self.cfg;
8154        let geometry = la.geometry;
8155        let d_state = geometry.key_head_dim as usize;
8156        let num_v = geometry.value_heads as usize;
8157        let eps = cfg.rms_eps;
8158        let scale = 1.0 / (d_state as f32).sqrt();
8159
8160        let prep =
8161            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
8162
8163        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
8164        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
8165        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
8166        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
8167        // verify keep the sequential kernel).
8168        let mut o = e.uninit(d_state * num_v * t)?;
8169        let rl = cache.recur[il].as_mut().unwrap();
8170        {
8171            let crate::cache::RecurLayer {
8172                ssm_state,
8173                ssm_state_alt,
8174                ..
8175            } = rl;
8176            e.gdn_scan_prefill(
8177                &prep.q_l2,
8178                &prep.k_l2,
8179                &prep.v_g,
8180                &prep.g_log,
8181                &prep.beta,
8182                prep.kb16.as_ref(),
8183                prep.qb16.as_ref(),
8184                ssm_state,
8185                ssm_state_alt,
8186                &mut o,
8187                num_v,
8188                t,
8189                scale,
8190                prep.hk,
8191            )?;
8192        }
8193        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8194
8195        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
8196        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
8197        let mut gn = e.uninit(d_state * num_v * t)?;
8198        let gn16 = if Self::f16out_on(e, t) {
8199            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
8200            e.gated_rmsnorm_f16out_zv(
8201                &o,
8202                la.ssm_norm.float_data(),
8203                z,
8204                &mut gn,
8205                &mut g16,
8206                d_state,
8207                num_v * t,
8208                eps,
8209            )?;
8210            Some(g16)
8211        } else {
8212            e.gated_rmsnorm_zv(
8213                &o,
8214                la.ssm_norm.float_data(),
8215                z,
8216                &mut gn,
8217                d_state,
8218                num_v * t,
8219                eps,
8220            )?;
8221            None
8222        };
8223        Ok((gn, gn16))
8224    }
8225
8226    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
8227    #[allow(clippy::too_many_arguments)]
8228    fn linear_attn_prime_core_pad(
8229        &self,
8230        e: &Engine,
8231        la: &LinearAttnLayer,
8232        g4: Vec<CudaSlice<f32>>,
8233        t: usize,
8234        cache: &mut Cache,
8235        il: usize,
8236        pad_len: Option<&CudaSlice<i32>>,
8237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8238        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
8239        if let Some(xh) = &gn16
8240            && let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)?
8241        {
8242            return Ok(y);
8243        }
8244        e.matmul(&la.ssm_out, &gn, t)
8245    }
8246
8247    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
8248    ///
8249    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
8250    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
8251    pub fn full_attn(
8252        &self,
8253        e: &Engine,
8254        fa: &FullAttnLayer,
8255        h: &CudaSlice<f32>,
8256        pos_d: &CudaSlice<i32>,
8257        t: usize,
8258        il: usize,
8259    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8260        if self.uses_sliding_gated_moe_program() {
8261            return self.step35_attn(e, fa, h, pos_d, t, il);
8262        }
8263        let cfg = &self.cfg;
8264        let _n_embd = cfg.n_embd as usize;
8265        let geometry = cfg.full_attention_geometry_at(il as u32);
8266        let n_head = geometry.n_head as usize;
8267        let n_head_kv = geometry.n_head_kv as usize;
8268        let head_dim = geometry.head_dim_k as usize;
8269        let eps = cfg.rms_eps;
8270        let scale = geometry.attention_scale();
8271
8272        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
8273        // gate — wq out = n_head*head_dim, no split (see prime-path note).
8274        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8275        // A load-time full-attention TP plan owns the same Q/K/V projections for every
8276        // architecture. Fall back to the original grouped owner-device projection when this
8277        // layer has no TP sidecar.
8278        let mut g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
8279            Some(g3) => g3,
8280            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
8281        };
8282        let v = g3.pop().unwrap();
8283        let mut k = g3.pop().unwrap();
8284        let qf = g3.pop().unwrap();
8285        let (mut q, gate) = if gated {
8286            let mut q = e.uninit(t * n_head * head_dim)?;
8287            let mut gate = e.uninit(t * n_head * head_dim)?;
8288            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
8289            (q, Some(gate))
8290        } else {
8291            (qf, None)
8292        };
8293
8294        // QK-norm (per head_dim row), then partial RoPE.
8295        let mut qn = e.uninit(t * n_head * head_dim)?;
8296        e.rms_norm(
8297            &q,
8298            fa.q_norm.float_data(),
8299            &mut qn,
8300            head_dim,
8301            n_head * t,
8302            eps,
8303        )?;
8304        q = qn;
8305        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
8306        e.rms_norm(
8307            &k,
8308            fa.k_norm.float_data(),
8309            &mut kn,
8310            head_dim,
8311            n_head_kv * t,
8312            eps,
8313        )?;
8314        k = kn;
8315        let rope_dims = geometry.n_rot as usize;
8316        e.rope_neox(
8317            &mut q,
8318            pos_d,
8319            head_dim,
8320            rope_dims,
8321            n_head,
8322            t,
8323            geometry.rope_base,
8324            1.0,
8325        )?;
8326        e.rope_neox(
8327            &mut k,
8328            pos_d,
8329            head_dim,
8330            rope_dims,
8331            n_head_kv,
8332            t,
8333            geometry.rope_base,
8334            1.0,
8335        )?;
8336
8337        // SDPA
8338        let mut attn = e.uninit(t * n_head * head_dim)?;
8339        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
8340        // falls back to naive sdpa.
8341        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
8342            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
8343            e.sdpa_naive(
8344                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
8345            )?;
8346        } else {
8347            e.fa_prefill(
8348                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
8349            )?;
8350        }
8351
8352        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
8353        let attn_g = match &gate {
8354            Some(gate) => {
8355                let mut gsig = e.uninit(t * n_head * head_dim)?;
8356                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8357                let mut ag = e.uninit(t * n_head * head_dim)?;
8358                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8359                ag
8360            }
8361            None => attn,
8362        };
8363
8364        // O follows the same generic load-time TP plan as Q/K/V.
8365        self.full_attn_o(e, fa, &attn_g, t)
8366    }
8367
8368    /// The absorb/decompress operands (`attn_k_b` / `attn_v_b`) are 3D and are ALWAYS the Float
8369    /// arm — on every checkpoint dtype, not just an f32 fixture. Two guards upstream make that a
8370    /// property rather than a hope: `GpuTensor::load_from_source` refuses any quantized non-2D
8371    /// tensor by name (`row_bytes` is derived from `ne[1]`, the MIDDLE axis of a 3D tensor), and
8372    /// `MlaAttnLayer::load` audits residency at load. A quantized `kv_b_proj` is dequantized at
8373    /// the source by `TransformKind::MlaKeyUpSplit`/`MlaValueUpSplit`. This is the last backstop:
8374    /// fail NAMING the constraint rather than through `float_data()`'s norm-flavoured panic.
8375    fn mla_split_operand<'w>(
8376        w: &'w crate::model::GpuTensor,
8377        name: &str,
8378        il: usize,
8379    ) -> &'w CudaSlice<f32> {
8380        match w {
8381            crate::model::GpuTensor::Float { data, .. } => data,
8382            _ => panic!(
8383                "layer {il}: MLA conversion-split operand {name} is not f32-resident. The 3D \
8384                 (d_nope|kv_rank, kv_rank|d_v, n_head) splits have no quantized resident layout: \
8385                 a quantized 3D tensor mis-derives row_bytes in the generic 2D Quant arm, so the \
8386                 source must dequantize the fused kv_b_proj (TensorTransform::SplitMlaKv). \
8387                 Reaching this means both the loader rank guard and MlaAttnLayer::load's \
8388                 residency audit were bypassed"
8389            ),
8390        }
8391    }
8392
8393    /// MLA (multi-head latent attention) mixer core, ABSORBED form — the one arm that serves
8394    /// prefill, chunked prefill and decode (see `cu/mla_attn.cu` FORM CHOICE).
8395    ///
8396    /// `latent` is the layer's latent KV plane; this call APPENDS its own `t` rows at row
8397    /// `slot` and then attends rows `0..slot + t`, which is exactly the oracle's convention
8398    /// that the queries are the LAST `t_q` rows of the cache (`crate::mla::MlaInputs`).
8399    /// Returns the post-`wo` block output [t, n_embd].
8400    #[allow(clippy::too_many_arguments)]
8401    // allow: the parameter list mirrors the kernel/FFI/call contract (rows_exact is the
8402    // verify-batch matmul-class selector, lane/glm5-verify-batch); bundling into a struct
8403    // is a refactor, not a lint fix
8404    fn mla_attn_core(
8405        &self,
8406        e: &Engine,
8407        mla: &crate::hybrid::MlaAttnLayer,
8408        h: &CudaSlice<f32>,
8409        pos_d: &CudaSlice<i32>,
8410        t: usize,
8411        il: usize,
8412        latent: &mut CudaSlice<f32>,
8413        index_plane: Option<IndexerPlanes<'_>>,
8414        slot: usize,
8415        rows_exact: bool,
8416    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8417        let attn = self.mla_attn_core_pre_wo(
8418            e,
8419            mla,
8420            h,
8421            pos_d,
8422            t,
8423            il,
8424            latent,
8425            index_plane,
8426            slot,
8427            rows_exact,
8428        )?;
8429        // Verify-batch wo seam (lane/glm5-verify-batch): the rows arm routes the output
8430        // projection decode-exact, same as every projection inside the core — the wo
8431        // dispatch moved here with the TP split, its routing did not change.
8432        if rows_exact {
8433            e.matmul_rows_exact(&mla.wo, &attn, t)
8434        } else {
8435            e.matmul(&mla.wo, &attn, t)
8436        }
8437    }
8438
8439    /// [`mla_attn_core`] up to (and excluding) the output projection: returns the
8440    /// per-head attention output `[t, n_head * d_v]`. Split out for the glm5 TP-2 seam,
8441    /// whose column-parallel `wo` runs over the cross-rank GATHERED heads — the plain path
8442    /// is the wrapper above, byte-for-byte the pre-split body (the wo matmul moved,
8443    /// nothing else).
8444    #[allow(clippy::too_many_arguments)]
8445    /// The MLA core split into its three segments (lane/glm5-mla-segments-20260904), same kernels
8446    /// in the same order as the single function they replace, so that the decode graph can
8447    /// capture the projections (pre) and the attention (post) with the adjacent KDA runs while
8448    /// the position-dependent middle (append, pool keys, score, select) stays eager. This is the
8449    /// pure refactor step: [`Self::mla_attn_core_pre_wo`] is pre -> mid -> post and every caller
8450    /// is unchanged.
8451    #[allow(clippy::too_many_arguments)]
8452    fn mla_attn_core_pre_wo(
8453        &self,
8454        e: &Engine,
8455        mla: &crate::hybrid::MlaAttnLayer,
8456        h: &CudaSlice<f32>,
8457        pos_d: &CudaSlice<i32>,
8458        t: usize,
8459        il: usize,
8460        latent: &mut CudaSlice<f32>,
8461        index_plane: Option<IndexerPlanes<'_>>,
8462        slot: usize,
8463        rows_exact: bool,
8464    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8465        // MEMRA_MLA_SEG_WS (lane/glm5-mla-capture-20260904, default OFF): run the PRE segment
8466        // into the session's stable buffers. Byte-identical (same kernels, same order, a
8467        // different destination address), and the seam the capture arc needs.
8468        if t == 1 && Engine::mla_seg_ws_on() {
8469            let g = mla.geom;
8470            let q_lora = mla.wq_b.in_features();
8471            let mut ws = e.mla_seg_ws_take(g.n_head, g.d_nope, g.d_rope, g.kv_rank, q_lora)?;
8472            let out = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8473                if MLA_SEG_WS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
8474                    eprintln!(
8475                        "[mla-seg-ws] engaged: the T=1 MLA PRE segment writes the session's \
8476                         stable handoff buffers (MEMRA_MLA_SEG_WS=1)"
8477                    );
8478                }
8479                self.mla_seg_pre(e, mla, h, pos_d, t, il, rows_exact, Some(&mut ws))?;
8480                let gathered = self.mla_seg_mid(
8481                    e,
8482                    mla,
8483                    h,
8484                    MlaMidIn {
8485                        q_an: &ws.q_an,
8486                        c_kv_n: &ws.c_kv_n,
8487                        k_pe: &ws.k_pe,
8488                    },
8489                    latent,
8490                    index_plane,
8491                    t,
8492                    il,
8493                    slot,
8494                    rows_exact,
8495                )?;
8496                self.mla_seg_post(
8497                    e, mla, &ws.q_nope, &ws.q_pe, gathered, latent, t, il, slot, rows_exact,
8498                )
8499            })();
8500            e.mla_seg_ws_put(ws);
8501            return out;
8502        }
8503        let pre = self
8504            .mla_seg_pre(e, mla, h, pos_d, t, il, rows_exact, None)?
8505            .expect("the ws-free PRE segment always returns its own buffers");
8506        let gathered = self.mla_seg_mid(
8507            e,
8508            mla,
8509            h,
8510            MlaMidIn {
8511                q_an: &pre.q_an,
8512                c_kv_n: &pre.c_kv_n,
8513                k_pe: &pre.k_pe,
8514            },
8515            latent,
8516            index_plane,
8517            t,
8518            il,
8519            slot,
8520            rows_exact,
8521        )?;
8522        self.mla_seg_post(
8523            e,
8524            mla,
8525            &pre.q_nope,
8526            &pre.q_pe,
8527            gathered,
8528            latent,
8529            t,
8530            il,
8531            slot,
8532            rows_exact,
8533        )
8534    }
8535
8536    /// Segment PRE of the MLA core: the q and kv projections, their norms, the latent splits and
8537    /// the rope of both position planes. Position enters only through `pos_d` (a device
8538    /// pointer), so this segment is capturable.
8539    #[allow(clippy::too_many_arguments)]
8540    fn mla_seg_pre(
8541        &self,
8542        e: &Engine,
8543        mla: &crate::hybrid::MlaAttnLayer,
8544        h: &CudaSlice<f32>,
8545        pos_d: &CudaSlice<i32>,
8546        t: usize,
8547        il: usize,
8548        rows_exact: bool,
8549        ws: Option<&mut MlaSegWs>,
8550    ) -> Result<Option<MlaPreOut>, Box<dyn std::error::Error>> {
8551        let g = mla.geom;
8552        let cfg = &self.cfg;
8553        let eps = cfg.rms_eps;
8554        let base = cfg.rope_freq_base;
8555        let (nh, dn, dr, r) = (g.n_head, g.d_nope, g.d_rope, g.kv_rank);
8556        assert_eq!(
8557            g.latent_dim,
8558            r + dr,
8559            "layer {il}: MlaGeom latent_dim disagrees with kv_rank + d_rope"
8560        );
8561        // Verify-batch matmul seam (lane/glm5-verify-batch): rows_exact routes every
8562        // projection through the decode-exact classes so each of the t rows is
8563        // bit-identical to the t=1 decode program (matmul_rows_exact contract); false =
8564        // the unchanged dispatch for every other caller (prime keeps its classes).
8565        let mm = |w: &crate::model::GpuTensor,
8566                  x: &CudaSlice<f32>|
8567         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8568            if rows_exact {
8569                e.matmul_rows_exact(w, x, t)
8570            } else {
8571                e.matmul(w, x, t)
8572            }
8573        };
8574
8575        // --- q path: wq_a -> q_a_norm -> wq_b -> per-head [nope | rope] ---
8576        let q_lora = mla.wq_b.in_features();
8577        // POOLED destinations (door `MEMRA_MLA_SEG_WS`, t == 1 only): the same kernels write the
8578        // same values into the session's stable buffers instead of fresh ones, so a captured PRE
8579        // graph and a captured POST graph can hand each other addresses. A layer whose geometry
8580        // differs from the set's refuses by name rather than writing past an end.
8581        if let Some(ws) = ws {
8582            if t != 1 || ws.sig != (nh, dn, dr, r, q_lora) {
8583                return Err(format!(
8584                    "layer {il}: the MLA segment workspace is sized {:?} for t=1 and this call is \
8585                     {:?} at t={t}",
8586                    ws.sig,
8587                    (nh, dn, dr, r, q_lora)
8588                )
8589                .into());
8590            }
8591            let q_a = mm(&mla.wq_a, h)?;
8592            e.rms_norm(
8593                &q_a,
8594                mla.q_a_norm.float_data(),
8595                &mut ws.q_an,
8596                q_lora,
8597                t,
8598                eps,
8599            )?;
8600            let q = mm(&mla.wq_b, &ws.q_an)?;
8601            e.mla_split_latent(&q, &mut ws.q_nope, &mut ws.q_pe, t * nh, dn, dr)?;
8602            e.mla_rope_interleaved(&mut ws.q_pe, pos_d, t, nh, dr, base)?;
8603            let kv = mm(&mla.wkv_a, h)?;
8604            let mut c_kv = e.uninit(t * r)?;
8605            e.mla_split_latent(&kv, &mut c_kv, &mut ws.k_pe, t, r, dr)?;
8606            e.rms_norm(&c_kv, mla.kv_a_norm.float_data(), &mut ws.c_kv_n, r, t, eps)?;
8607            e.mla_rope_interleaved(&mut ws.k_pe, pos_d, t, 1, dr, base)?;
8608            return Ok(None);
8609        }
8610        let q_a = mm(&mla.wq_a, h)?;
8611        let mut q_an = e.uninit(t * q_lora)?;
8612        e.rms_norm(&q_a, mla.q_a_norm.float_data(), &mut q_an, q_lora, t, eps)?;
8613        let q = mm(&mla.wq_b, &q_an)?;
8614        // Per head the row is [nope | rope] contiguous, so t*nh rows of width dn+dr split with
8615        // the same kernel the latent row uses — the two layouts are the same shape.
8616        let mut q_nope = e.uninit(t * nh * dn)?;
8617        let mut q_pe = e.uninit((t * nh * dr).max(1))?;
8618        e.mla_split_latent(&q, &mut q_nope, &mut q_pe, t * nh, dn, dr)?;
8619        // NoPE (glm5_next, rope_head_dim 0): no rope plane exists. The launcher is a no-op, but
8620        // the allocation above is still non-empty so nothing downstream holds a null slice.
8621        e.mla_rope_interleaved(&mut q_pe, pos_d, t, nh, dr, base)?;
8622
8623        // --- kv path: wkv_a -> [c_kv (rms-normed) | k_pe (roped, NOT normed)] ---
8624        let kv = mm(&mla.wkv_a, h)?;
8625        let mut c_kv = e.uninit(t * r)?;
8626        let mut k_pe = e.uninit((t * dr).max(1))?;
8627        e.mla_split_latent(&kv, &mut c_kv, &mut k_pe, t, r, dr)?;
8628        let mut c_kv_n = e.uninit(t * r)?;
8629        e.rms_norm(&c_kv, mla.kv_a_norm.float_data(), &mut c_kv_n, r, t, eps)?;
8630        e.mla_rope_interleaved(&mut k_pe, pos_d, t, 1, dr, base)?;
8631
8632        // --- DSA k-pool selection, BEFORE attending: the indexer's own state row for each of
8633        // this call's tokens is appended first, so a query sees itself exactly as the latent
8634        // plane already lets it (the reference concatenates into the indexer cache, then scores).
8635        Ok(Some(MlaPreOut {
8636            q_nope,
8637            q_pe,
8638            q_an,
8639            c_kv_n,
8640            k_pe,
8641        }))
8642    }
8643
8644    /// Segment MID of the MLA core: the latent append at `slot` and the DSA k-pool selection.
8645    /// Both derive launch geometry from the host-side `slot`, which is why this segment is the
8646    /// one the decode graph leaves eager.
8647    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
8648    fn mla_seg_mid(
8649        &self,
8650        e: &Engine,
8651        mla: &crate::hybrid::MlaAttnLayer,
8652        h: &CudaSlice<f32>,
8653        planes: MlaMidIn<'_>,
8654        latent: &mut CudaSlice<f32>,
8655        index_plane: Option<IndexerPlanes<'_>>,
8656        t: usize,
8657        il: usize,
8658        slot: usize,
8659        rows_exact: bool,
8660    ) -> Result<Option<(CudaSlice<i32>, usize)>, Box<dyn std::error::Error>> {
8661        let g = mla.geom;
8662        let (dr, r) = (g.d_rope, g.kv_rank);
8663        e.mla_append_latent(latent, planes.c_kv_n, planes.k_pe, slot, t, r, dr)?;
8664        let q_an = planes.q_an;
8665        let gathered = match (&mla.index, index_plane) {
8666            (Some(indexer), Some(plane)) => {
8667                Some(self.mla_kpool_select(e, indexer, h, q_an, plane, t, slot, il, rows_exact)?)
8668            }
8669            (Some(_), None) => {
8670                return Err(format!(
8671                    "layer {il} declares a DSA k-pool indexer but no indexer state plane was \
8672                     supplied — the ModelPlan must declare StatePlan::LatentKvCache with a \
8673                     non-zero index_width for it"
8674                )
8675                .into());
8676            }
8677            (None, _) => None,
8678        };
8679
8680        // --- absorbed MLA core over the latent plane ---
8681        Ok(gathered)
8682    }
8683
8684    /// Segment POST of the MLA core: the absorbed query, the attention over the selected (or
8685    /// full) latent rows, the value decompression. Its geometry is the selection width, which
8686    /// is constant past `topk * pool` tokens, so this segment is capturable there.
8687    #[allow(clippy::too_many_arguments)]
8688    fn mla_seg_post(
8689        &self,
8690        e: &Engine,
8691        mla: &crate::hybrid::MlaAttnLayer,
8692        q_nope: &CudaSlice<f32>,
8693        q_pe: &CudaSlice<f32>,
8694        gathered: Option<(CudaSlice<i32>, usize)>,
8695        latent: &CudaSlice<f32>,
8696        t: usize,
8697        il: usize,
8698        slot: usize,
8699        rows_exact: bool,
8700    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8701        let g = mla.geom;
8702        let (nh, dn, dr, dv, r) = (g.n_head, g.d_nope, g.d_rope, g.d_v, g.kv_rank);
8703        let t_kv = slot + t;
8704        let wk_b = Self::mla_split_operand(&mla.wk_b, "attn_k_b", il);
8705        let wv_b = Self::mla_split_operand(&mla.wv_b, "attn_v_b", il);
8706
8707        // MEMRA_MLA_TC_PREFILL door (default OFF; flag read per call — the rollback seam).
8708        // Engagement conditions, every one load-bearing:
8709        //   * a gathered selection exists — the door serves the DSA arm only; the dense
8710        //     absorbed arm (GLM-5.2, no indexer) keeps the f32 kernel it was gated on;
8711        //   * d_rope == 0 (NoPE) — the TC kernel treats the latent row as both K and V,
8712        //     which is only the whole truth when there is no rope plane;
8713        //   * kv_rank == 512 — the kernel's stamped head dim (glm5_next / GLM-5.2 class);
8714        //   * t >= 16 — prefill widths only. Decode (t == 1) and short resumes NEVER enter,
8715        //     which is what the decode byte-identity gate proves rather than assumes.
8716        // Anything else falls through to the unchanged f32 kernels below — behavior identical
8717        // to the flag being off.
8718        // A chain returning Ok(None) is a cuBLASLt shape DECLINE (announced once per shape);
8719        // the let-chain then simply does not match and the f32 kernels below serve the call.
8720        // glm5 TP composition guard (lane/glm5-composition): the TC prefill chain's gate
8721        // ran on the FULL-head geometry only; a head shard (any rank) declines it by name
8722        // and falls through to the f32 kernels below — behavior identical to the flag
8723        // being off for that layer, announced once. The composed door re-gates on the box
8724        // (real-artifact kv_rank 512 shapes; the rig fixtures are kv_rank 16 and never
8725        // reach this chain).
8726        // The announce shares the chain's OWN conjuncts (gathered + !portable_mma_gated),
8727        // so it can never blame TP for a decline the missing DSA gather or the MMA gate
8728        // caused (#82 review).
8729        if mla.tp_shard
8730            && gathered.is_some()
8731            && dr == 0
8732            && r == 512
8733            && t >= 16
8734            && !crate::portable_mma_gated()
8735            && mla_tc_prefill_enabled()
8736        {
8737            static TP_TC_DECLINE: std::sync::Once = std::sync::Once::new();
8738            TP_TC_DECLINE.call_once(|| {
8739                eprintln!(
8740                    "[mla-tc-prefill] DECLINED on glm5-TP head shards: the door's gate ran \
8741                     on full-head geometry; shards ride the f32 prefill kernels until the \
8742                     TP composition gate lands (pin MEMRA_MLA_TC_PREFILL=0 to silence)"
8743                );
8744            });
8745        }
8746        if let Some((idx, slots)) = &gathered
8747            && dr == 0
8748            && r == 512
8749            && t >= 16
8750            && !rows_exact // verify-batch stays on the decode-exact classes (t <= 15 anyway)
8751            && !crate::portable_mma_gated()
8752            && !mla.tp_shard
8753            && mla_tc_prefill_enabled()
8754            && let Some(attn) = self.mla_tc_prefill_chain(
8755                e, wk_b, wv_b, q_nope, latent, idx, *slots, t, t_kv, nh, dn, dv, r, g.scale,
8756            )?
8757        {
8758            return Ok(attn);
8759        }
8760
8761        let mut q_lat = e.uninit(t * nh * r)?;
8762        e.mla_absorb_q(q_nope, wk_b, &mut q_lat, t, nh, dn, r)?;
8763        let mut o_lat = e.uninit(t * nh * r)?;
8764        match &gathered {
8765            Some((idx, slots)) => e.mla_attn_gathered(
8766                &q_lat, q_pe, latent, idx, &mut o_lat, nh, r, dr, t, *slots, g.scale,
8767            )?,
8768            None => e.mla_attn_absorbed(
8769                &q_lat, q_pe, latent, &mut o_lat, nh, r, dr, t, t_kv, g.scale,
8770            )?,
8771        }
8772        let mut attn = e.uninit(t * nh * dv)?;
8773        e.mla_decompress_v(&o_lat, wv_b, &mut attn, t, nh, dv, r)?;
8774
8775        Ok(attn)
8776    }
8777
8778    /// The MEMRA_MLA_TC_PREFILL chain: absorb and decompress as strided-batched bf16
8779    /// tensor-core GEMMs, attention as the gathered bf16 MMA kernel. Returns `Ok(None)` when
8780    /// cuBLASLt declines a GEMM shape (announced once per shape) so the caller falls back to
8781    /// the f32 kernels; every other failure is a hard error.
8782    ///
8783    /// FORM CHOICE, stated for the record (the dual-form MLA law): every fast engine runs
8784    /// MATERIALIZED (per-head MHA) attention at DENSE prefill and absorbed MQA at decode.
8785    /// glm5_next prefill is NOT dense: the DSA indexer caps every query at topk+tail rows and
8786    /// selects ONE list per query SHARED ACROSS ALL 64 HEADS. That shared list is what makes
8787    /// the ABSORBED form the GEMM-shaped one here — the head axis is the MMA m, the shared
8788    /// latent rows are one B operand per tile — while materializing K/V would give every head
8789    /// its own K plane and destroy exactly that sharing (back to per-(query,head) matvecs on
8790    /// the gathered walk). It is also FlashMLA's own sparse-prefill geometry (q 576/512 over
8791    /// gathered latent rows). Queries whose selection is trivial (visible <= topk: the lists
8792    /// ARE the full causal prefix, emitted by the selector itself) ride the SAME kernel with
8793    /// the identity gather — there is no separate dense program to gate.
8794    ///
8795    /// Transient cost per (layer, chunk) at the census shape (t=2313, t_kv=4626): bf16 q_lat
8796    /// 152 MB + bf16 latent window 4.7 MB + bf16 q_nope/o_lat copies ~230 MB — all freed with
8797    /// the call. (The materialized-K/V alternative would have been 4096 x 64 x (256+256) x 2B
8798    /// = 256 MB/layer-chunk of K/V ALONE, plus the per-head-K program cost above.) Weight
8799    /// bf16 converts (wk_b/wv_b, 8.4M elems each) run per call, ~50 us class; a resident
8800    /// mirror is a later diet, not correctness.
8801    #[allow(clippy::too_many_arguments)]
8802    fn mla_tc_prefill_chain(
8803        &self,
8804        e: &Engine,
8805        wk_b: &CudaSlice<f32>,
8806        wv_b: &CudaSlice<f32>,
8807        q_nope: &CudaSlice<f32>,
8808        latent: &CudaSlice<f32>,
8809        idx: &CudaSlice<i32>,
8810        width: usize,
8811        t: usize,
8812        t_kv: usize,
8813        nh: usize,
8814        dn: usize,
8815        dv: usize,
8816        r: usize,
8817        scale: f32,
8818    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
8819        // Once-per-shape decline announce (the bf16_tc_gemm pattern): a door that quietly
8820        // stops engaging reads exactly like a door that never helped.
8821        fn declined(stage: &str, m: usize, n: usize, k: usize, batch: usize) {
8822            type ShapeSet = std::collections::HashSet<(usize, usize, usize, usize)>;
8823            static SAID: std::sync::Mutex<Option<ShapeSet>> = std::sync::Mutex::new(None);
8824            let mut g = SAID.lock().unwrap();
8825            if g.get_or_insert_with(std::collections::HashSet::new)
8826                .insert((m, n, k, batch))
8827            {
8828                eprintln!(
8829                    "[mla-tc-prefill] DECLINED at {stage} m={m} n={n} k={k} batch={batch} \
8830                     (no cuBLASLt heuristic) — this call falls back to the f32 MLA kernels"
8831                );
8832            }
8833        }
8834        // Weights and activations to bf16. The converts require n % 4 == 0; every operand here
8835        // is a multiple of the head dims (dn/dv/r all >= 16 and % 4 == 0 on the shapes the door
8836        // admits), asserted rather than assumed.
8837        for (name, n) in [
8838            ("wk_b", nh * r * dn),
8839            ("wv_b", nh * dv * r),
8840            ("q_nope", t * nh * dn),
8841            ("latent", t_kv * r),
8842        ] {
8843            debug_assert!(
8844                n.is_multiple_of(4),
8845                "mla-tc-prefill: {name} elems {n} % 4 != 0"
8846            );
8847            let _ = (name, n);
8848        }
8849        let wk_bf = e.f32_to_bf16(wk_b, nh * r * dn)?;
8850        let wv_bf = e.f32_to_bf16(wv_b, nh * dv * r)?;
8851        let qn_bf = e.f32_to_bf16(q_nope, t * nh * dn)?;
8852        // absorb: per head h, q_lat[:,h,:] [t, r] = q_nope[:,h,:] [t, dn] @ W_uk[h] [r, dn]^T.
8853        // wk_b is the conversion-split (h, l, p) plane, contiguous in p == the reduction axis:
8854        // per head it IS the [n=r, k=dn] row-major operand. bf16 out feeds the attention kernel.
8855        let mut q_lat_bf = e.alloc_u8_uninit(t * nh * r * 2)?;
8856        if !e.mla_bf16_gemm_sb_bf16out(
8857            &wk_bf,
8858            &qn_bf,
8859            &mut q_lat_bf,
8860            t,
8861            r,
8862            dn,
8863            nh * dn,
8864            dn,
8865            nh * r,
8866            r,
8867            nh,
8868        )? {
8869            declined("absorb", t, r, dn, nh);
8870            return Ok(None);
8871        }
8872        // The latent window rows 0..t_kv (this call's rows were appended above), bf16.
8873        let cache_bf = e.f32_to_bf16(latent, t_kv * r)?;
8874        let mut o_lat = e.uninit(t * nh * r)?;
8875        e.mla_attn_gathered_tc(
8876            &q_lat_bf, &cache_bf, idx, &mut o_lat, nh, r, t, width, scale,
8877        )?;
8878        // decompress: per head h, attn[:,h,:] [t, dv] = o_lat[:,h,:] [t, r] @ W_uv[h] [dv, r]^T.
8879        // wv_b is (h, j, l), contiguous in l == the reduction axis: per head [n=dv, k=r].
8880        let o_bf = e.f32_to_bf16(&o_lat, t * nh * r)?;
8881        let mut attn = e.uninit(t * nh * dv)?;
8882        if !e.mla_bf16_gemm_sb_f32out(
8883            &wv_bf,
8884            &o_bf,
8885            &mut attn,
8886            t,
8887            dv,
8888            r,
8889            nh * r,
8890            r,
8891            nh * dv,
8892            dv,
8893            nh,
8894        )? {
8895            declined("decompress", t, dv, r, nh);
8896            return Ok(None);
8897        }
8898        crate::MLA_TC_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8899        {
8900            static ANNOUNCED: std::sync::Once = std::sync::Once::new();
8901            ANNOUNCED.call_once(|| {
8902                eprintln!(
8903                    "[mla-tc-prefill] engaged: absorb/decompress = strided-batched bf16 TC \
8904                     GEMMs, attention = fa_mla_gathered_bf16 (t={t}, t_kv={t_kv}, nh={nh}, \
8905                     width={width}); dispatches counted in MLA_TC_PREFILL_DISPATCHES"
8906                );
8907            });
8908        }
8909        Ok(Some(attn))
8910    }
8911
8912    /// Layer-scoped wrapper: names the layer in any selection failure.
8913    #[allow(clippy::too_many_arguments)]
8914    fn mla_kpool_select(
8915        &self,
8916        e: &Engine,
8917        indexer: &crate::hybrid::MlaIndexer,
8918        h: &CudaSlice<f32>,
8919        q_resid: &CudaSlice<f32>,
8920        plane: IndexerPlanes<'_>,
8921        t: usize,
8922        slot: usize,
8923        il: usize,
8924        rows_exact: bool,
8925    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
8926        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, rows_exact).map_err(
8927            |source| -> Box<dyn std::error::Error> {
8928                format!("layer {il}: DSA k-pool selection failed: {source}").into()
8929            },
8930        )
8931    }
8932
8933    /// DSA k-pool indexer: append this call's packed indexer state, then select the cache rows
8934    /// each query may attend. Returns the per-query position list and its width (`-1` padded).
8935    ///
8936    /// The program is `Glm5NextTextIndexer.forward`
8937    /// (research/glm53-flash-bringup-20260827/modular_glm5_next-ref.py:771), transcribed in
8938    /// `memra_reference::kpool_allowed_tokens`, which is this path's oracle:
8939    ///   1. `k = LayerNorm_affine(wk(x))` — LayerNorm WITH BIAS at eps 1e-5, NOT the model's
8940    ///      RMSNorm at `rms_norm_eps`; `gate = index_kpool_compress_gate(x)`. Both are cached.
8941    ///   2. Every COMPLETE pool of `pool` consecutive cached tokens collapses to one key by a
8942    ///      per-channel softmax over (gate + positional embedding).
8943    ///   3. `score[i][p] = sum_h relu(q[i][h] . pool_key[p] * d^-1/2) * weights_proj(x)[i][h] *
8944    ///      heads^-1/2`, with pools whose last token is invisible to the query masked out.
8945    ///   4. Top `top_k / pool` pools expand back to raw rows; the incomplete tail is appended raw.
8946    ///
8947    /// `q_resid` is `q_a_layernorm(q_a_proj(x))` — the SAME tensor the MLA query up-projection
8948    /// consumes, which is why the indexer is scored here rather than before the core.
8949    #[allow(clippy::too_many_arguments)]
8950    pub fn mla_kpool_indices(
8951        e: &Engine,
8952        indexer: &crate::hybrid::MlaIndexer,
8953        h: &CudaSlice<f32>,
8954        q_resid: &CudaSlice<f32>,
8955        plane: IndexerPlanes<'_>,
8956        t: usize,
8957        slot: usize,
8958    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
8959        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, false)
8960    }
8961
8962    /// [`Self::mla_kpool_indices`] with the verify-batch matmul-class selector
8963    /// (lane/glm5-verify-batch): `rows_exact` routes the indexer's four projections
8964    /// through the decode-exact classes so each of the t rows is bit-identical to the
8965    /// t=1 decode program; `false` is the unchanged dispatch.
8966    #[allow(clippy::too_many_arguments)]
8967    pub fn mla_kpool_indices_ex(
8968        e: &Engine,
8969        indexer: &crate::hybrid::MlaIndexer,
8970        h: &CudaSlice<f32>,
8971        q_resid: &CudaSlice<f32>,
8972        plane: IndexerPlanes<'_>,
8973        t: usize,
8974        slot: usize,
8975        rows_exact: bool,
8976    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
8977        let mm = |w: &crate::model::GpuTensor,
8978                  x: &CudaSlice<f32>|
8979         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8980            if rows_exact {
8981                e.matmul_rows_exact(w, x, t)
8982            } else {
8983                e.matmul(w, x, t)
8984            }
8985        };
8986        /// `nn.LayerNorm` default epsilon. The indexer's k_norm is a LayerNorm, so it does NOT
8987        /// take the model's `rms_norm_eps` (census: "eps 1e-5, NOT rms_norm_eps"); they coincide
8988        /// numerically on GLM-5.3-Flash and the constant keeps them from being coupled.
8989        const INDEX_NORM_EPS: f32 = 1e-5;
8990
8991        let ig = indexer.geom;
8992        let d = ig.head_dim;
8993        let t_kv = slot + t;
8994        let IndexerPlanes {
8995            state: plane,
8996            pool_keys: pool_key_plane,
8997            ready: pools_ready,
8998            state_ring_rows,
8999            capacity_tokens,
9000        } = plane;
9001
9002        // TAIL RING. The plane's rows are read EXACTLY ONCE, by the pool-key build of the pool
9003        // each row belongs to, so a ring of `ring` rows holds everything still live. The state
9004        // plan does not carry `pool`, so the allocator books physical rows and the EFFECTIVE ring
9005        // is rounded down here: a ring that is not a whole number of pools would split a pool
9006        // across the wrap. ONE POOL is the whole correctness floor (lane/glm53-ring-sizing): the
9007        // drain below serves any `t` from any ring at or above it, so `ring` never bounds a
9008        // prompt.
9009        let ring = if state_ring_rows == 0 {
9010            0
9011        } else {
9012            state_ring_rows / ig.pool * ig.pool
9013        };
9014        if state_ring_rows > 0 && ring == 0 {
9015            return Err(format!(
9016                "indexer tail ring of {state_ring_rows} rows cannot hold one pool of {}; \
9017                 raise MEMRA_DSA_INDEX_RING or set it to 0 for the flat plane",
9018                ig.pool
9019            )
9020            .into());
9021        }
9022        // RESIDENCY TRIPWIRE. `*pools_ready` counts pools whose keys were built over rows that are
9023        // now history. If the cache ever rewound past `slot` without clamping it (see
9024        // `LatentKvLayer::truncate_index_pool_keys`), those keys were built over rows this call is
9025        // about to overwrite — a silent wrong selection. Fail here instead.
9026        if *pools_ready > slot / ig.pool {
9027            return Err(format!(
9028                "resident k-pool key plane claims {} finished pools but the cache holds only {} \
9029                 complete pools before this call ({slot} rows / pool {}) — a rewind reduced the \
9030                 latent length without clamping index_pools_ready",
9031                *pools_ready,
9032                slot / ig.pool,
9033                ig.pool
9034            )
9035            .into());
9036        }
9037
9038        // 1. packed state rows [k | gate], appended at `slot` — the same [a|b] row shape the
9039        //    latent plane uses, so `mla_append_latent` packs it with no new kernel.
9040        let k_raw = mm(&indexer.wk, h)?;
9041        let mut k_norm = e.uninit(t * d)?;
9042        e.layer_norm_bias(
9043            &k_raw,
9044            indexer.k_norm_w.float_data(),
9045            indexer.k_norm_b.float_data(),
9046            &mut k_norm,
9047            d,
9048            t,
9049            INDEX_NORM_EPS,
9050        )?;
9051        let gate = mm(&indexer.kpool_gate, h)?;
9052
9053        // 2. pool keys over every COMPLETE pool in the cache — INCREMENTALLY. A pool's key is a
9054        //    function of its own `pool` state rows (append-only, never rewritten) and the constant
9055        //    `kpool_ape`, so it is final the instant the pool's last row lands. Only pools
9056        //    `[*pools_ready, n_pools)` are built; the rest are already resident and bit-identical
9057        //    to a rebuild. This turns the old O(t_kv * d) per-call pass into O(t * d).
9058        let n_pools = t_kv / ig.pool;
9059        let select_k = ig.select_k(n_pools);
9060        let width = ig.index_width(n_pools);
9061        // Sized to the SESSION's capacity, so a session that primes and then decodes never
9062        // reallocates (a fresh buffer would drop every resident key, and under the ring the rows
9063        // to rebuild them from are gone). `capacity_tokens` is that capacity; a pool covers
9064        // `ig.pool` tokens, so the key plane is `pool` times SHORTER than a flat state plane —
9065        // 32 f32 per token against 256. It is NOT read off `plane.len()` any more: once the state
9066        // plane is a ring, its length is one call's tail, not the context.
9067        // `.max(1)` keeps the slice non-null at t_kv < pool, where no complete pool exists yet.
9068        let capacity_pools = capacity_tokens / ig.pool;
9069        let need = (capacity_pools * d).max(n_pools * d).max(1);
9070        if pool_key_plane.as_ref().is_none_or(|k| k.len() < need) {
9071            *pool_key_plane = Some(e.uninit(need)?);
9072            *pools_ready = 0;
9073        }
9074        let pool_keys = pool_key_plane
9075            .as_mut()
9076            .expect("resident pool-key plane just allocated");
9077
9078        // THE DRAIN. The state plane is written by exactly one kernel and read by exactly one,
9079        // and a row's single read is the pool-key build of the pool that row belongs to. So the
9080        // rows that must be live at any instant are `[*pools_ready * pool, cur)`: everything
9081        // below has been read, everything above is not written yet, and the two kernels can be
9082        // interleaved in sub-ranges of the call instead of run once each over the whole call.
9083        //
9084        // That is what makes the ring size a WORKING-SET choice rather than a bound on `t`:
9085        // `index_ring_take` hands back how many rows fit before the ring must be drained, the
9086        // build drains it, and the loop continues. `k_norm`/`gate` are computed ONCE for the
9087        // whole call above and walked by source-row offset, so the values, their order, and the
9088        // ring addresses they land on are exactly what a single whole-call append produced.
9089        // A flat plane (`ring == 0`) takes the whole call in one iteration, byte for byte.
9090        let ape = indexer.kpool_ape.float_data();
9091        let mut cur = slot;
9092        let mut appended = 0usize;
9093        while appended < t {
9094            let take =
9095                crate::cache::index_ring_take(ring, ig.pool, *pools_ready, cur, t - appended)
9096                    .ok_or_else(|| -> Box<dyn std::error::Error> {
9097                        format!(
9098                            "indexer tail ring lapped: {ring} rows cannot hold the {} rows still \
9099                         owed to unbuilt pools at row {cur} (pools_ready {}, pool {}, slot \
9100                         {slot}, t {t}). The pool-key plane was reset or the cache rewound \
9101                         without clamping index_pools_ready, so rows this call must read were \
9102                         already overwritten. Raise MEMRA_DSA_INDEX_RING, or set \
9103                         MEMRA_DSA_INDEX_RING=0 for the flat plane",
9104                            cur.saturating_sub((*pools_ready).saturating_mul(ig.pool)),
9105                            *pools_ready,
9106                            ig.pool
9107                        )
9108                        .into()
9109                    })?;
9110            debug_assert!(take > 0 && appended + take <= t);
9111            e.mla_index_append(plane, &k_norm, &gate, appended, cur, take, d, d, ring)?;
9112            cur += take;
9113            appended += take;
9114            let ready_now = cur / ig.pool;
9115            e.mla_kpool_pool_keys(
9116                plane,
9117                ape,
9118                pool_keys,
9119                (*pools_ready).min(ready_now),
9120                ready_now,
9121                ig.pool,
9122                d,
9123                ring,
9124            )?;
9125            *pools_ready = ready_now;
9126        }
9127        debug_assert!(t == 0 || *pools_ready == n_pools);
9128        let pool_keys = &*pool_keys;
9129
9130        // 3. score + head mix, 4. top-k -> raw rows + tail.
9131        let q_index = mm(&indexer.wq_b, q_resid)?;
9132        let head_weights = mm(&indexer.weights_proj, h)?;
9133        let mut score = e.uninit((t * n_pools).max(1))?;
9134        e.mla_kpool_score(
9135            &q_index,
9136            pool_keys,
9137            &head_weights,
9138            &mut score,
9139            t,
9140            ig.heads,
9141            d,
9142            n_pools,
9143            ig.pool,
9144            slot,
9145            (d as f32).powf(-0.5),
9146            (ig.heads as f32).powf(-0.5),
9147        )?;
9148        let mut idx = e.uninit_i32(t * width)?;
9149        e.mla_kpool_select(
9150            &score,
9151            &mut idx,
9152            t,
9153            n_pools,
9154            ig.pool,
9155            select_k,
9156            width,
9157            slot,
9158            ig.always_select_tail,
9159        )?;
9160        Ok((idx, width))
9161    }
9162
9163    /// STATELESS MLA arm (`HybridModel::forward`): the latent plane lives for this call only,
9164    /// sized to the request. Same math as the cached arm — it is the same core with slot 0.
9165    pub fn mla_attn(
9166        &self,
9167        e: &Engine,
9168        mla: &crate::hybrid::MlaAttnLayer,
9169        h: &CudaSlice<f32>,
9170        pos_d: &CudaSlice<i32>,
9171        t: usize,
9172        il: usize,
9173    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9174        if mla.tp.is_some() {
9175            return Err(format!(
9176                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the stateless \
9177                 mixer path is unwired for a head shard"
9178            )
9179            .into());
9180        }
9181        let mut latent = e.uninit(t * mla.geom.latent_dim)?;
9182        let mut index_plane = match mla.index.as_ref() {
9183            Some(indexer) => Some(e.uninit(t * indexer.geom.state_width())?),
9184            None => None,
9185        };
9186        // No residency across calls here — the planes die with the call, so `ready` starts at 0 and
9187        // every pool is built exactly once, which is what the cached arm also does on its prime.
9188        let mut pool_keys = None;
9189        let mut pools_ready = 0usize;
9190        let planes = index_plane.as_mut().map(|state| IndexerPlanes {
9191            state,
9192            pool_keys: &mut pool_keys,
9193            ready: &mut pools_ready,
9194            // Per-call plane, sized to the request: no ring, capacity is the request itself.
9195            state_ring_rows: 0,
9196            capacity_tokens: t,
9197        });
9198        self.mla_attn_core(e, mla, h, pos_d, t, il, &mut latent, planes, 0, false)
9199    }
9200
9201    /// STATEFUL MLA arm (prime and T=1 decode): appends into the session's latent plane and
9202    /// attends the whole history. `cache.latent[il]` is allocated by the `LatentKvCache` arm of
9203    /// the cache allocator; a `None` here means the ModelPlan and the loaded mixer disagree.
9204    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
9205    pub fn mla_attn_cached(
9206        &self,
9207        e: &Engine,
9208        mla: &crate::hybrid::MlaAttnLayer,
9209        h: &CudaSlice<f32>,
9210        pos_d: &CudaSlice<i32>,
9211        t: usize,
9212        il: usize,
9213        cache: &mut Cache,
9214    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9215        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, false)
9216    }
9217
9218    /// [`Self::mla_attn_cached`] on the VERIFY-BATCH matmul classes (lane/glm5-verify-batch):
9219    /// the SAME core at t=K+1 rows with every internal projection routed decode-exact
9220    /// (`matmul_rows_exact`), so row r of the batched call is bit-identical to the t=1
9221    /// `mla_attn_cached` call the per-row verify walk makes at position pos0+r. Causality
9222    /// within the batch is per-query by construction: the kpool selection masks pools
9223    /// invisible to each query and appends each query's OWN raw tail
9224    /// (`first_pos + t + 1`), and the gathered attention walks each query's own idx list
9225    /// (-1 padding arithmetic-invariant). Held by `glm5_tparallel_verify_gpu` gates 1+2
9226    /// running the batched arm. ONLY the glm5 verify-batch walk calls this.
9227    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
9228    pub fn mla_attn_cached_rows_exact(
9229        &self,
9230        e: &Engine,
9231        mla: &crate::hybrid::MlaAttnLayer,
9232        h: &CudaSlice<f32>,
9233        pos_d: &CudaSlice<i32>,
9234        t: usize,
9235        il: usize,
9236        cache: &mut Cache,
9237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9238        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, true)
9239    }
9240
9241    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
9242    fn mla_attn_cached_inner(
9243        &self,
9244        e: &Engine,
9245        mla: &crate::hybrid::MlaAttnLayer,
9246        h: &CudaSlice<f32>,
9247        pos_d: &CudaSlice<i32>,
9248        t: usize,
9249        il: usize,
9250        cache: &mut Cache,
9251        rows_exact: bool,
9252    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9253        // glm5 TP fail-closed choke point, covering BOTH plain entries (decode/prime AND
9254        // the verify-batch rows arm): a TP-sharded layer holds heads/2 and a per-rank
9255        // latent replica — running it on the plain path would compute a silently-halved
9256        // mixer against the wrong plane, so it refuses by name instead.
9257        if mla.tp.is_some() {
9258            return Err(format!(
9259                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the plain mixer \
9260                 path is unwired for a head shard — only the TP decode/prime walk may \
9261                 execute it (rows_exact={rows_exact})"
9262            )
9263            .into());
9264        }
9265        // Read before the layer borrow: it sizes the resident pool-key plane, which the ring'd
9266        // state plane's own length can no longer stand in for.
9267        let max_ctx = cache.max_ctx;
9268        let layer = cache.latent[il].as_mut().ok_or_else(|| {
9269            format!(
9270                "layer {il} is Mixer::Mla but the cache has no latent plane — the ModelPlan \
9271                 must declare StatePlan::LatentKvCache for it"
9272            )
9273        })?;
9274        let attn =
9275            self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?;
9276        // Verify-batch wo seam: the rows arm keeps its decode-exact output projection —
9277        // the wo dispatch moved here with the TP split, its routing did not change.
9278        if rows_exact {
9279            e.matmul_rows_exact(&mla.wo, &attn, t)
9280        } else {
9281            e.matmul(&mla.wo, &attn, t)
9282        }
9283    }
9284
9285    /// The stateful MLA call against ONE latent plane, up to (and excluding) the output
9286    /// projection. The plain path wraps it above (canonical plane + `wo`); the glm5 TP-2
9287    /// walk calls it once per rank (root shard on the canonical plane, peer shard on the
9288    /// replicated peer plane) and joins the halves through the column-parallel `wo`.
9289    #[allow(clippy::too_many_arguments)]
9290    pub(crate) fn mla_attn_cached_pre_wo(
9291        &self,
9292        e: &Engine,
9293        mla: &crate::hybrid::MlaAttnLayer,
9294        h: &CudaSlice<f32>,
9295        pos_d: &CudaSlice<i32>,
9296        t: usize,
9297        il: usize,
9298        layer: &mut memra_kv::LatentKvLayer,
9299        max_ctx: usize,
9300        rows_exact: bool,
9301    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9302        let slot = layer.len;
9303        let width = layer.width;
9304        assert_eq!(
9305            width, mla.geom.latent_dim,
9306            "layer {il}: cache latent width {width} != MlaGeom latent_dim {}",
9307            mla.geom.latent_dim
9308        );
9309        let capacity = layer.rows.len() / width;
9310        if slot + t > capacity {
9311            return Err(format!(
9312                "layer {il}: latent cache overflow — {slot} + {t} rows exceeds capacity {capacity}"
9313            )
9314            .into());
9315        }
9316        if mla.index.is_some() && layer.index_rows.is_none() {
9317            return Err(format!(
9318                "layer {il} loaded a DSA k-pool indexer but its latent cache carries no indexer \
9319                 state plane — StatePlan::LatentKvCache declared index_width 0 for a layer whose \
9320                 SparseIndexPlan is Own {{ kpool: Some(..) }}"
9321            )
9322            .into());
9323        }
9324        // Both planes are borrowed for the whole core call, so len bookkeeping happens after.
9325        // ONE `len` covers both: they are appended in the same call and must never drift.
9326        // The resident pool-key plane rides along: it is state that must SURVIVE the call, so the
9327        // core writes `ready` back through the borrow and it is restored with the buffers.
9328        let mut rows = std::mem::replace(&mut layer.rows, e.uninit(0)?);
9329        let mut index_rows = layer.index_rows.take();
9330        let mut pool_keys = layer.index_pool_keys.take();
9331        let mut pools_ready = layer.index_pools_ready;
9332        let index_ring_rows = layer.index_ring_rows.unwrap_or(0);
9333        let planes = index_rows.as_mut().map(|state| IndexerPlanes {
9334            state,
9335            pool_keys: &mut pool_keys,
9336            ready: &mut pools_ready,
9337            state_ring_rows: index_ring_rows,
9338            capacity_tokens: max_ctx,
9339        });
9340        let out =
9341            self.mla_attn_core_pre_wo(e, mla, h, pos_d, t, il, &mut rows, planes, slot, rows_exact);
9342        layer.rows = rows;
9343        layer.index_rows = index_rows;
9344        layer.index_pool_keys = pool_keys;
9345        // A FAILED core leaves `len` where it was, so the resident plane must go back too: it may
9346        // have advanced over pools built from rows a retry is about to rewrite with different
9347        // inputs. Clamping here (rather than letting the next call's tripwire fire) makes
9348        // retry-after-error correct instead of merely loud.
9349        layer.index_pools_ready = if out.is_ok() {
9350            pools_ready
9351        } else if let Some(indexer) = mla.index.as_ref() {
9352            pools_ready.min(layer.len / indexer.geom.pool)
9353        } else {
9354            pools_ready
9355        };
9356        let out = out?;
9357        // Resolve the layer's RESIDENT pool copy (the state plan does not carry `pool`; the
9358        // latent-plane snapshot path reads this field to address the tail ring). A nonzero
9359        // resident value that disagrees with the loaded geometry is corruption, not a race:
9360        // there is exactly one geometry per loaded layer.
9361        if let Some(indexer) = mla.index.as_ref() {
9362            let pool = indexer.geom.pool;
9363            if layer.index_pool != 0 && layer.index_pool != pool {
9364                return Err(format!(
9365                    "layer {il}: resident indexer pool {} != loaded geometry pool {pool}",
9366                    layer.index_pool,
9367                )
9368                .into());
9369            }
9370            layer.index_pool = pool;
9371        }
9372        layer.len = slot + t;
9373        let len_i32 = i32::try_from(layer.len).map_err(|_| "latent length exceeds i32 mirror")?;
9374        // Door H (`MEMRA_GLM5_HTOD_DIET`): the async `i32_set_k` launch instead of this
9375        // SYNCHRONIZING pageable 4-byte copy — 11 of these per round, one per MLA trunk layer.
9376        e.i32_mirror_store(&mut layer.len_d, len_i32)?;
9377        Ok(out)
9378    }
9379
9380    /// The glm5 TP MLA walk for one prime/decode call (`mla` is the ROOT head shard; its
9381    /// sidecar carries the peer shards + runtime). Replicated per-token work runs on EVERY
9382    /// rank from identical inputs (wq_a/wkv_a/indexer/k-pool selection — identical bytes by
9383    /// determinism on uniform hardware, gate-held); each rank attends with its heads over
9384    /// its OWN latent replica; the attention parts are gathered through the armed transport
9385    /// and each rank's COLUMN-parallel `wo` slice computes its slice of the output with the
9386    /// plain matvec kernel — no cross-rank arithmetic anywhere.
9387    #[allow(clippy::too_many_arguments)]
9388    pub(crate) fn mla_tp_attn_cached(
9389        &self,
9390        e: &Engine,
9391        mla: &crate::hybrid::MlaAttnLayer,
9392        h: &CudaSlice<f32>,
9393        pos_d: &CudaSlice<i32>,
9394        t: usize,
9395        il: usize,
9396        cache: &mut Cache,
9397        rows_exact: bool,
9398    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9399        let tp = mla
9400            .tp
9401            .as_ref()
9402            .ok_or("mla_tp_attn_cached called on an unsharded layer")?;
9403        let rt = &tp.rt;
9404        let ranks = tp.ranks();
9405        let g = mla.geom; // SHARD geometry: n_head = full/ranks
9406        let hl = g.n_head;
9407        let dv = g.d_v;
9408        let full_heads = tp.full_heads;
9409        let n_embd = tp.n_embd;
9410        let hh = n_embd / ranks;
9411        let max_ctx = cache.max_ctx;
9412
9413        // HOP 1 — fan-out of the mixer input and positions to every peer rank. Both move the
9414        // WHOLE buffer, exactly as the v1 arm did, so the transport arms move identical
9415        // byte ranges (lane/glm5-tp-transport).
9416        let hop = rt.hop(e);
9417        let h_peers = crate::tp_transport::fanout_f32(&hop, h, h.len())?;
9418        let pos_peers = crate::tp_transport::fanout_i32(&hop, pos_d, pos_d.len())?;
9419
9420        // Peer replica planes (lazily geometry-cloned from the canonical plane).
9421        {
9422            let canonical = cache.latent[il].as_ref().ok_or_else(|| {
9423                format!("layer {il}: glm5 TP MLA walk found no canonical latent plane")
9424            })?;
9425            crate::glm5_tp::ensure_mla_peer_latent(
9426                rt,
9427                canonical,
9428                &mut cache.glm5_tp_latent_peer[il],
9429            )?;
9430        }
9431
9432        // Peer passes first (each rank's heads over its replica), then root (canonical
9433        // plane unchanged) — v1's issue order at two ranks. `rows_exact` threads the
9434        // caller's matmul class through every rank: false = the prime/decode walk
9435        // (byte-for-byte the pre-composition arm), true = the spec x TP verify walk
9436        // (lane/glm5-composition) riding the same rows-exact classes as the unsharded
9437        // verify walk.
9438        let mut attn: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
9439        for r in 1..ranks {
9440            let layer = &mut cache.glm5_tp_latent_peer[il].as_mut().unwrap()[r - 1];
9441            attn[r] = Some(self.mla_attn_cached_pre_wo(
9442                &rt.peers[r - 1],
9443                &tp.peers[r - 1],
9444                &h_peers[r - 1],
9445                &pos_peers[r - 1],
9446                t,
9447                il,
9448                layer,
9449                max_ctx,
9450                rows_exact,
9451            )?);
9452        }
9453        attn[0] = {
9454            let layer = cache.latent[il].as_mut().unwrap();
9455            Some(self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?)
9456        };
9457
9458        // HOP 2 — gather the per-head parts into the FULL [t, full_heads*dv] layout on
9459        // every rank. `full_heads * dv == ranks * (hl * dv)` by the shard map.
9460        let part = hl * dv;
9461        debug_assert_eq!(full_heads * dv, ranks * part);
9462        let attn_refs: Vec<&CudaSlice<f32>> = attn
9463            .iter()
9464            .map(|a| a.as_ref().expect("filled above"))
9465            .collect();
9466        let fulls = crate::tp_transport::gather_parts(&hop, &attn_refs, t, part)?;
9467
9468        // Column-parallel wo slices + output concat (pure movement). The verify walk's
9469        // wo rides the rows-exact class, exactly like the unsharded verify walk's wo.
9470        let mut ys = Vec::with_capacity(ranks);
9471        if rows_exact {
9472            ys.push(e.matmul_rows_exact(&mla.wo, &fulls[0], t)?);
9473            for r in 1..ranks {
9474                ys.push(rt.peers[r - 1].matmul_rows_exact(&tp.peers[r - 1].wo, &fulls[r], t)?);
9475            }
9476        } else {
9477            ys.push(e.matmul(&mla.wo, &fulls[0], t)?);
9478            for r in 1..ranks {
9479                ys.push(rt.peers[r - 1].matmul(&tp.peers[r - 1].wo, &fulls[r], t)?);
9480            }
9481        }
9482        // HOP 3 — concat the column parts into the mixer output on ROOT.
9483        debug_assert_eq!(n_embd, ranks * hh);
9484        let y_refs: Vec<&CudaSlice<f32>> = ys.iter().collect();
9485        crate::tp_transport::concat_parts_on_root(&hop, &y_refs, t, hh)
9486    }
9487
9488    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
9489    pub fn linear_attn(
9490        &self,
9491        e: &Engine,
9492        la: &LinearAttnLayer,
9493        h: &CudaSlice<f32>,
9494        t: usize,
9495    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9496        let cfg = &self.cfg;
9497        let _n_embd = cfg.n_embd as usize;
9498        let geometry = la.geometry;
9499        let d_state = geometry.key_head_dim as usize;
9500        let num_k = geometry.key_heads as usize;
9501        let num_v = geometry.value_heads as usize;
9502        let d_conv = geometry.conv_kernel as usize;
9503        let head_k = d_state;
9504        let head_v = geometry.value_head_dim as usize;
9505        let key_dim = head_k * num_k; // 2048
9506        let value_dim = head_v * num_v; // 4096
9507        let conv_dim = key_dim * 2 + value_dim; // 8192
9508        let eps = cfg.rms_eps;
9509        let scale = 1.0 / (d_state as f32).sqrt();
9510
9511        // projections
9512        // grouped: one f16 activation convert feeds all four projections (matmul_group)
9513        let mut g4 = e.matmul_group(
9514            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
9515            h,
9516            t,
9517        )?;
9518        let alpha = g4.pop().unwrap(); // [T, num_v]
9519        let beta_raw = g4.pop().unwrap(); // [T, num_v]
9520        let z = g4.pop().unwrap(); // [T, value_dim]
9521        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
9522
9523        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
9524        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
9525        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
9526        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
9527        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
9528        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
9529        let _ = (head_k, head_v);
9530        let mut q_g = e.uninit(d_state * num_v * t)?;
9531        let mut k_g = e.uninit(d_state * num_v * t)?;
9532        let mut v_g = e.uninit(d_state * num_v * t)?;
9533        e.ssm_conv1d_gdn(
9534            &qkv_mixed,
9535            la.ssm_conv1d.float_data(),
9536            &mut q_g,
9537            &mut k_g,
9538            &mut v_g,
9539            conv_dim,
9540            t,
9541            d_conv,
9542            d_state,
9543            num_v,
9544            num_k,
9545            key_dim,
9546        )?;
9547        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
9548        let mut q_l2 = e.uninit(d_state * num_v * t)?;
9549        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
9550        let mut k_l2 = e.uninit(d_state * num_v * t)?;
9551        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
9552        let v_gd = v_g;
9553
9554        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
9555        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
9556        let mut beta = e.uninit(t * num_v)?;
9557        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
9558        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
9559        let mut g_log = e.uninit(t * num_v)?;
9560        e.gdn_glog(
9561            &alpha,
9562            la.ssm_dt.float_data(),
9563            la.ssm_a.float_data(),
9564            &mut g_log,
9565            num_v,
9566            t,
9567        )?;
9568
9569        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
9570        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
9571        let mut state_out = e.zeros(d_state * d_state * num_v)?;
9572        let mut o = e.uninit(d_state * num_v * t)?;
9573        e.gdn_scan_prefill(
9574            &q_l2,
9575            &k_l2,
9576            &v_gd,
9577            &g_log,
9578            &beta,
9579            None,
9580            None,
9581            &state_in,
9582            &mut state_out,
9583            &mut o,
9584            num_v,
9585            t,
9586            scale,
9587            num_v,
9588        )?;
9589
9590        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
9591        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
9592        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
9593        // o rows are (t*num_v+vh) too. Good.
9594        let mut gn = e.uninit(d_state * num_v * t)?;
9595        e.gated_rmsnorm(
9596            &o,
9597            la.ssm_norm.float_data(),
9598            &z,
9599            &mut gn,
9600            d_state,
9601            num_v * t,
9602            eps,
9603        )?;
9604
9605        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
9606        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
9607        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
9608        let out = e.matmul(&la.ssm_out, &gn, t)?;
9609        Ok(out)
9610    }
9611}
9612
9613impl HybridModel {
9614    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
9615    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
9616    ///
9617    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
9618    /// different 860160-byte block than the same expert of layer 7).
9619    ///
9620    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
9621    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
9622    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
9623    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
9624    pub fn moe_ffn_il(
9625        &self,
9626        e: &Engine,
9627        m: &MoeWeights,
9628        z: &CudaSlice<f32>,
9629        t: usize,
9630        il: u16,
9631    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9632        Self::moe_ffn_inner(
9633            e,
9634            m,
9635            z,
9636            None,
9637            t,
9638            &self.cfg,
9639            il,
9640            self.max_moe_block(),
9641            false,
9642            None,
9643            self.uses_sliding_gated_moe_program(),
9644            false,
9645        )
9646    }
9647
9648    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
9649    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
9650    pub fn moe_ffn_il_prefill(
9651        &self,
9652        e: &Engine,
9653        m: &MoeWeights,
9654        z: &CudaSlice<f32>,
9655        t: usize,
9656        il: u16,
9657    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9658        Self::moe_ffn_inner(
9659            e,
9660            m,
9661            z,
9662            None,
9663            t,
9664            &self.cfg,
9665            il,
9666            self.max_moe_block(),
9667            true,
9668            Some(&self.step_grouped_prefill),
9669            self.uses_sliding_gated_moe_program(),
9670            false,
9671        )
9672    }
9673
9674    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
9675    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
9676    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
9677    pub fn moe_ffn_il_zq8(
9678        &self,
9679        e: &Engine,
9680        m: &MoeWeights,
9681        z: &CudaSlice<f32>,
9682        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
9683        t: usize,
9684        il: u16,
9685    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9686        Self::moe_ffn_inner(
9687            e,
9688            m,
9689            z,
9690            zq8,
9691            t,
9692            &self.cfg,
9693            il,
9694            self.max_moe_block(),
9695            false,
9696            None,
9697            self.uses_sliding_gated_moe_program(),
9698            false,
9699        )
9700    }
9701
9702    /// Verify-rows twin of [`Self::moe_ffn_il_zq8`] (lane/glm5-vrest): the SAME routing and
9703    /// dispatch decisions with the pairs-shaped batched routed-expert arm armed. Only the
9704    /// verify walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`, t>=2) calls this; every
9705    /// unqualified shape inside falls closed to the byte-identical sequential loop.
9706    pub(crate) fn moe_ffn_il_zq8_vrows(
9707        &self,
9708        e: &Engine,
9709        m: &MoeWeights,
9710        z: &CudaSlice<f32>,
9711        t: usize,
9712        il: u16,
9713    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9714        Self::moe_ffn_inner(
9715            e,
9716            m,
9717            z,
9718            None,
9719            t,
9720            &self.cfg,
9721            il,
9722            self.max_moe_block(),
9723            false,
9724            None,
9725            self.uses_sliding_gated_moe_program(),
9726            true,
9727        )
9728    }
9729
9730    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
9731    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
9732    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
9733    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
9734    ///
9735    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
9736    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
9737    pub(crate) fn moe_ffn(
9738        e: &Engine,
9739        m: &MoeWeights,
9740        z: &CudaSlice<f32>,
9741        t: usize,
9742        cfg: &ModelConfig,
9743        il: u16,
9744        max_block: usize,
9745    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9746        Self::moe_ffn_inner(
9747            e, m, z, None, t, cfg, il, max_block, false, None, false, false,
9748        )
9749    }
9750
9751    #[allow(clippy::too_many_arguments)]
9752    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
9753    pub(crate) fn moe_ffn_inner(
9754        e: &Engine,
9755        m: &MoeWeights,
9756        z: &CudaSlice<f32>,
9757        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
9758        t: usize,
9759        cfg: &ModelConfig,
9760        il: u16,
9761        max_block: usize,
9762        prefill: bool,
9763        grouped_prefill: Option<&std::sync::Mutex<crate::hybrid::StepEpGroupedPrefill>>,
9764        sliding_gated_moe: bool,
9765        vrows: bool,
9766    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9767        let worker_io = crate::spill_pread::worker_enabled();
9768        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
9769        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
9770            e.with_moe_cache(max_block, |cache, _| {
9771                cache.begin_forward_epoch(il, t);
9772                if worker_io {
9773                    cache.begin_worker_scope();
9774                }
9775                Ok(())
9776            })?;
9777        }
9778        if let Some(ep) = &m.glm5_ep {
9779            // glm5 TP-2 EP walk (MEMRA_GLM5_TP): whole-expert halves, root router, slot-ordered
9780            // canonical combine. Every other arm of this function is unreachable for an
9781            // EP-armed layer by construction. `prefill` keys the EP grouped-prime arm
9782            // (MEMRA_GLM5_EP_GROUPED_PRIME) exactly as it keys the plain grouped arm below.
9783            return Self::moe_ffn_glm5_ep(e, m, ep, z, zq8, t, cfg, il, prefill);
9784        }
9785        if m.step_ep.is_some() || m.step_tp.is_some() {
9786            let moe = cfg
9787                .moe
9788                .as_ref()
9789                .ok_or("Step distributed execution requires MoE model metadata")?;
9790            let n_embd = cfg.n_embd as usize;
9791            let n_expert = moe.expert_count as usize;
9792            let n_used = moe.expert_used_count as usize;
9793            let sigmoid = cfg
9794                .sigmoid_router()
9795                .ok_or("Step distributed execution requires the Step sigmoid router")?;
9796            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9797            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
9798            let grouped_prefill_requested = prefill && step_ep_grouped_prefill_enabled()?;
9799            if grouped_prefill_requested && !step_tp_prefill_enabled()? {
9800                return Err(
9801                    "MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into(),
9802                );
9803            }
9804            if grouped_prefill_requested && !step_grouped_prefill_shape(true, prefill, t) {
9805                return Err(format!(
9806                    "Step grouped prefill tokens {t} are outside the qualified {}..={} range",
9807                    PRIME_MIN_T,
9808                    crate::cache::PRIME_CHUNK_MAX_TOKENS,
9809                )
9810                .into());
9811            }
9812            let grouped_decode_shape = step_grouped_decode_shape(prefill, t);
9813            let grouped_prefill_shape =
9814                step_grouped_prefill_shape(grouped_prefill_requested, prefill, t);
9815            if let Some(ep) = m.step_ep.as_ref().filter(|ep| {
9816                ep.grouped_decode.is_some() && (grouped_decode_shape || grouped_prefill_shape)
9817            }) {
9818                let (selected, route_weights) =
9819                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
9820                crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
9821                Self::trace_moe_routes(il, t, &selected, &route_weights)?;
9822                Self::trace_moe_input(e, il, t, n_embd, z)?;
9823                let selected = selected
9824                    .iter()
9825                    .map(|&expert| expert as usize)
9826                    .collect::<Vec<_>>();
9827
9828                // The narrow route readback above orders the owning-stage producer. The grouped
9829                // runtime then copies the resident root activation into its persistent rank inputs.
9830                e.stream().synchronize()?;
9831                let execute = |state: &mut crate::hybrid::StepEpGroupedDecode| {
9832                    state.projection.set_activation_limit(ep.activation_limit)?;
9833                    ep.runtime
9834                        .refresh_step_grouped_expert_parallel_gate_from_root_device(
9835                            ep.experts.e4m3()?,
9836                            &mut state.projection,
9837                            z,
9838                            t,
9839                            &selected,
9840                        )?;
9841                    ep.runtime.refresh_step_grouped_expert_parallel_combine(
9842                        &state.projection,
9843                        &mut state.combine,
9844                        &route_weights,
9845                    )?;
9846                    ep.runtime.execute_step_grouped_expert_parallel_gate(
9847                        ep.experts.e4m3()?,
9848                        &mut state.projection,
9849                    )?;
9850                    ep.runtime.execute_step_grouped_expert_parallel_combine(
9851                        &state.projection,
9852                        &mut state.combine,
9853                    )?;
9854                    let mut output = ep.runtime.copy_step_grouped_expert_parallel_combine_root(
9855                        &state.projection,
9856                        &state.combine,
9857                        e,
9858                    )?;
9859                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9860                    if prefill {
9861                        // A shared plan may be reused by the next layer on a different runtime
9862                        // stream. Complete the owning-stage copy before its source is overwritten.
9863                        e.stream().synchronize()?;
9864                    }
9865                    eprintln!(
9866                        "[step-tp-ep-grouped] execute layer={il} tokens={t} devices={:?} \
9867                         attention_layout=tensor-parallel expert_layout=expert-parallel \
9868                         expert_transport={} native_p2p=true route_control=host-narrow \
9869                         input=root-device projection_workspaces=persistent \
9870                         combine=root-device output=owning-stage-device \
9871                         prefill={prefill} batched_decode=false capacity={} \
9872                         performance_claim=false",
9873                        ep.devices,
9874                        ep.runtime.transport_label(),
9875                        state.projection.max_tokens(),
9876                    );
9877                    Ok::<_, Box<dyn std::error::Error>>(output)
9878                };
9879
9880                if grouped_prefill_shape {
9881                    let grouped_prefill = grouped_prefill
9882                        .ok_or("Step grouped prefill has no model-scoped executor")?;
9883                    let mut shared = grouped_prefill
9884                        .lock()
9885                        .map_err(|_| "Step grouped prefill state lock is poisoned")?;
9886                    let needs_prepare = shared.state.as_ref().is_none_or(|state| {
9887                        state.devices != ep.devices
9888                            || state.grouped.projection.max_tokens() < t
9889                            || state.grouped.projection.input_width() != n_embd
9890                            || state.grouped.projection.expert_width()
9891                                != moe.expert_ff_length as usize
9892                    });
9893                    if needs_prepare {
9894                        let seed_input = vec![0.0f32; n_embd];
9895                        let seed_selected = &selected[..n_used];
9896                        let seed_weights = &route_weights[..n_used];
9897                        let projection = ep
9898                            .runtime
9899                            .prepare_step_grouped_expert_parallel_gate_with_capacity(
9900                                ep.experts.e4m3()?,
9901                                &seed_input,
9902                                1,
9903                                seed_selected,
9904                                ep.activation_limit,
9905                                t,
9906                            )?;
9907                        let combine = ep.runtime.prepare_step_grouped_expert_parallel_combine(
9908                            &projection,
9909                            seed_weights,
9910                        )?;
9911                        shared.state = Some(crate::hybrid::StepEpGroupedPrefillState {
9912                            devices: ep.devices.clone(),
9913                            grouped: crate::hybrid::StepEpGroupedDecode {
9914                                projection,
9915                                combine,
9916                            },
9917                        });
9918                        eprintln!(
9919                            "[step-tp-ep-grouped-prefill] prepare capacity={t} devices={:?} \
9920                             shared_across_layers=true performance_claim=false",
9921                            ep.devices,
9922                        );
9923                    }
9924                    return execute(
9925                        &mut shared
9926                            .state
9927                            .as_mut()
9928                            .expect("Step grouped prefill state prepared above")
9929                            .grouped,
9930                    );
9931                }
9932
9933                let mut grouped = ep
9934                    .grouped_decode
9935                    .as_ref()
9936                    .expect("grouped decode presence checked above")
9937                    .lock()
9938                    .map_err(|_| "Step grouped decode state lock is poisoned")?;
9939                return execute(&mut grouped);
9940            }
9941            if grouped_prefill_shape {
9942                return Err(
9943                    "Step grouped prefill requires native-P2P expert-owner device arithmetic"
9944                        .into(),
9945                );
9946            }
9947            // MEMRA_STEP_TP_DEV_ROUTER=1 (t=1): route on device and feed the device-routed
9948            // expert program — the per-layer host logits readback (the last per-layer host
9949            // sync) disappears. Selection tie-breaking may differ from the host router:
9950            // numeric-class door, run-gen argmax gate + boot battery.
9951            // STEP TP2 GEMM PRIME (2026-08-27, TTFT lane): a prime chunk's routed MoE goes
9952            // through ONE grouped f16 GEMM per projection over the resident NVFP4 banks —
9953            // the per-token device routes below cost 240 s at m=4092 (measured), the grouped
9954            // lane's sizing rows run 170-270 TFLOP/s. Router selections come from the same
9955            // sigmoid host oracle the EP arm uses; shexp rides the canonical grouped add.
9956            // t>=16 alone keys the branch: the batch prime reaches here through moe_ffn_il,
9957            // whose `prefill` is FALSE (only the _prefill twin sets it), and no other step37
9958            // route runs t>=16 — verify walks t<=8, decode t=1. Requiring `prefill` made the
9959            // first gate arm skip this branch entirely and wake the generic f16g arm instead
9960            // (48 s + kq_gemm_sk rc=1001, 2026-08-27).
9961            if t >= 16
9962                && crate::step_gemm_prime_on()
9963                && let Some(tp) = &m.step_tp
9964                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
9965            {
9966                // MEMRA_PRIME_PROF=1 sub-split of the moe bucket. The phase timer put
9967                // 1788 ms of a 3093 ms chunk here, but forcing the 32-row tile form (4x
9968                // more weight dequant) moved it only 5% — so the grouped GEMM is not
9969                // obviously what dominates. The router below is a HOST oracle: sigmoid +
9970                // top-8 over 288 experts for every one of 4096 tokens, per layer, which
9971                // is a D2H copy and a full pipeline drain 42 times per chunk. Attribute
9972                // it before optimizing the kernel it sits in front of.
9973                let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1") && t >= 16;
9974                let mut mt = std::time::Instant::now();
9975                let (selected, route_weights) =
9976                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
9977                let sel_i32: Vec<i32> = selected.iter().map(|&x| x as i32).collect();
9978                let d_router = if mprof {
9979                    let _ = e.stream().synchronize();
9980                    let v = mt.elapsed().as_secs_f64() * 1e3;
9981                    mt = std::time::Instant::now();
9982                    v
9983                } else {
9984                    0.0
9985                };
9986                // MEMRA_MOE_DETERM=1: run the WHOLE grouped routine twice on identical
9987                // inputs and diff. The standalone harness cleared the grouped GEMM kernels
9988                // (8 invocations, both lanes, maxdiff 0.0 over 20.9M elements) but it does
9989                // not model the cross-device join/scatter or the o_proj-style reduction,
9990                // and the loader refuses both topologies (TP1, same-device) that would
9991                // isolate those by env. This tests the un-excluded region directly, in
9992                // the place it actually runs.
9993                //
9994                // The prime is nondeterministic: same prompt, one forward, temperature=0,
9995                // max_tokens=1, and the first token varies across reps. That blocks
9996                // MEMRA_PP_BF16's correctness receipt and invalidates every byte-identity
9997                // gate taken through the server. This probe also yields the jitter
9998                // MAGNITUDE, which any tolerance band needs.
9999                let mdet =
10000                    std::env::var("MEMRA_MOE_DETERM").as_deref() == Ok("1") && t >= 16 && il < 4;
10001                if mdet {
10002                    let a = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
10003                        bank,
10004                        e,
10005                        z,
10006                        t,
10007                        &sel_i32,
10008                        &route_weights,
10009                        n_used,
10010                        tp.activation_limit,
10011                    )?;
10012                    let b = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
10013                        bank,
10014                        e,
10015                        z,
10016                        t,
10017                        &sel_i32,
10018                        &route_weights,
10019                        n_used,
10020                        tp.activation_limit,
10021                    )?;
10022                    let (ha, hb) = (e.dtoh(&a)?, e.dtoh(&b)?);
10023                    let mut md = 0.0f32;
10024                    let mut ndiff = 0usize;
10025                    for (x, y) in ha.iter().zip(hb.iter()) {
10026                        let d = (x - y).abs();
10027                        if d > 0.0 {
10028                            ndiff += 1;
10029                        }
10030                        if d > md {
10031                            md = d;
10032                        }
10033                    }
10034                    eprintln!(
10035                        "[moe-determ] il={il} t={t} maxdiff={md:.3e} \
10036                                 differing={ndiff}/{} -> {}",
10037                        ha.len(),
10038                        if ndiff == 0 {
10039                            "IDENTICAL"
10040                        } else {
10041                            "NONDETERMINISTIC"
10042                        }
10043                    );
10044                }
10045                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
10046                    bank,
10047                    e,
10048                    z,
10049                    t,
10050                    &sel_i32,
10051                    &route_weights,
10052                    n_used,
10053                    tp.activation_limit,
10054                )?;
10055                let d_gemm = if mprof {
10056                    let _ = e.stream().synchronize();
10057                    let v = mt.elapsed().as_secs_f64() * 1e3;
10058                    mt = std::time::Instant::now();
10059                    v
10060                } else {
10061                    0.0
10062                };
10063                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10064                if mprof {
10065                    let _ = e.stream().synchronize();
10066                    let d_shared = mt.elapsed().as_secs_f64() * 1e3;
10067                    // Per LAYER, not accumulated: the four trunk phases already carry the
10068                    // per-chunk totals, and one line per layer is what shows whether the
10069                    // cost is flat across layers or concentrated in a few.
10070                    eprintln!(
10071                        "[moe-prof] il={il} t={t} router={d_router:.1}ms \
10072                                 gemm={d_gemm:.1}ms shared={d_shared:.1}ms"
10073                    );
10074                }
10075                return Ok(output);
10076            }
10077            if t == 1
10078                && crate::tp::step_nvfp4_dev_routes_enabled()?
10079                && crate::tp::step_tp_dev_router_enabled()?
10080                && let Some(tp) = &m.step_tp
10081                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
10082            {
10083                let (sf, route_norm) = sigmoid;
10084                // MEMRA_ROUTES_PRESTAGE=1: issue the rank input pull + quantize NOW,
10085                // before the router — the rank streams overlap the gemv+topk.
10086                // MEMRA_DEV1_ROUTER=1 rides the prestage hook: rank1 routes LOCALLY
10087                // from its own z copy (replicated deterministic router — identical
10088                // bits in, identical sel/w out) and starts its sweep without
10089                // waiting the root's sel broadcast.
10090                static D1_ROUTER: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10091                let d1_router = *D1_ROUTER
10092                    .get_or_init(|| std::env::var("MEMRA_DEV1_ROUTER").as_deref() == Ok("1"));
10093                if d1_router {
10094                    let (sf_h, rn_h) = sigmoid;
10095                    let n_ex = m.gate_exps.n_expert;
10096                    let act_ct = m.active_count();
10097                    let _ = tp.runtime.nvfp4_routes_prestage_with(
10098                        bank,
10099                        e,
10100                        z,
10101                        |rank1, in1, sel1, w1| {
10102                            let mut guard = DEV1_ROUTER_REPS
10103                                .lock()
10104                                .map_err(|_| "dev1 router replica lock")?;
10105                            let (reps, scratch) =
10106                                guard.get_or_insert_with(|| (Default::default(), None));
10107                            if !reps.contains_key(&il) {
10108                                use cudarc::driver::DevicePtr;
10109                                let (g1, p1, a1) = (
10110                                    rank1.htod(&vec![0.0f32; n_ex * n_embd])?,
10111                                    rank1.htod(&vec![0.0f32; n_ex])?,
10112                                    rank1.alloc_u8_uninit(n_ex)?,
10113                                );
10114                                for (src, dst_len, dst) in [
10115                                    (
10116                                        {
10117                                            let s = e.stream();
10118                                            let (p, _g) = m.gate_inp.float_data().device_ptr(&s);
10119                                            p
10120                                        },
10121                                        n_ex * n_embd * 4,
10122                                        {
10123                                            let s = rank1.stream();
10124                                            let (p, _g) = g1.device_ptr(&s);
10125                                            p
10126                                        },
10127                                    ),
10128                                    (
10129                                        {
10130                                            let s = e.stream();
10131                                            let (p, _g) = m.exp_probs_b_dev.device_ptr(&s);
10132                                            p
10133                                        },
10134                                        n_ex * 4,
10135                                        {
10136                                            let s = rank1.stream();
10137                                            let (p, _g) = p1.device_ptr(&s);
10138                                            p
10139                                        },
10140                                    ),
10141                                    (
10142                                        {
10143                                            let s = e.stream();
10144                                            let (p, _g) = m.active_experts_dev.device_ptr(&s);
10145                                            p
10146                                        },
10147                                        n_ex,
10148                                        {
10149                                            let s = rank1.stream();
10150                                            let (p, _g) = a1.device_ptr(&s);
10151                                            p
10152                                        },
10153                                    ),
10154                                ] {
10155                                    crate::tp::raw_copy_bytes(dst, src, dst_len, rank1)?;
10156                                }
10157                                rank1.stream().synchronize()?;
10158                                reps.insert(il, (g1, p1, a1));
10159                            }
10160                            if scratch.is_none() {
10161                                *scratch = Some(rank1.htod(&vec![0.0f32; n_ex])?);
10162                            }
10163                            let (g1, p1, a1) = reps.get(&il).expect("armed above");
10164                            let logits1 = scratch.as_mut().expect("armed above");
10165                            rank1.router_gemv_into(g1, in1, logits1, n_embd, n_ex, 1)?;
10166                            rank1.moe_router_sigmoid_topk_into(
10167                                logits1, 1, n_ex, n_used, act_ct, p1, a1, sf_h, rn_h, sel1, w1,
10168                            )?;
10169                            Ok(true)
10170                        },
10171                    )?;
10172                } else {
10173                    let _ = tp.runtime.nvfp4_routes_prestage(bank, e, z)?;
10174                }
10175                // Persistent selection buffers: the allocating topk built two fresh
10176                // slices per layer; sel/w land in process-static rows instead
10177                // (host-op diet — same kernel, same bytes).
10178                #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
10179                static SELW: std::sync::Mutex<Option<(usize, CudaSlice<i32>, CudaSlice<f32>)>> =
10180                    std::sync::Mutex::new(None);
10181                let mut selw = SELW.lock().map_err(|_| "selw lock poisoned")?;
10182                if selw.as_ref().is_none_or(|(d, ..)| *d != e.ctx().ordinal()) {
10183                    *selw = Some((
10184                        e.ctx().ordinal(),
10185                        e.htod_i32(&vec![0i32; n_used])?,
10186                        e.htod(&vec![0.0f32; n_used])?,
10187                    ));
10188                }
10189                let (_, sel_d, w_d) = selw.as_mut().expect("armed above");
10190                e.moe_router_sigmoid_topk_into(
10191                    &logits,
10192                    t,
10193                    n_expert,
10194                    n_used,
10195                    m.active_count(),
10196                    &m.exp_probs_b_dev,
10197                    &m.active_experts_dev,
10198                    sf,
10199                    route_norm,
10200                    sel_d,
10201                    w_d,
10202                )?;
10203                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
10204                // FAIL-CLOSED for the route taps: this walk keeps the selection
10205                // device-side, so `trace_moe_routes` (MEMRA_MOE_TRACE /
10206                // MEMRA_MOE_WEIGHT_TRACE) never sees its rows. Every other MoE walk is
10207                // either host-routed (the taps ride the existing readback) or forced to
10208                // the host-visible path by observation mode — this one is neither. A
10209                // trace that silently misses whole layers poisons any placement mint
10210                // built on it (LAW:coactivation-expert-placement measurement leg), so an
10211                // armed tap refuses by name instead of dropping rows.
10212                if std::env::var("MEMRA_MOE_TRACE").is_ok()
10213                    || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
10214                {
10215                    return Err("MEMRA_MOE_TRACE/MEMRA_MOE_WEIGHT_TRACE cannot trace the \
10216                         device-routed step TP walk (selection never returns to host; \
10217                         tracing would add a new sync). Route through the host-router \
10218                         arm — refused rather than silently dropping rows"
10219                        .into());
10220                }
10221                crate::moe_sel_dump::refuse_device_only("the device-routed step TP walk")?;
10222                // MEMRA_SHEXP_OVERLAP=1: issue the shared expert from the routes
10223                // PREJOIN hook so it executes while the peer rank drains its sweep
10224                // (fills dev0's join wait); apply adds the identical values after.
10225                static SHEXP_OV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10226                let shexp_ov = *SHEXP_OV
10227                    .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
10228                // MEMRA_SHEXP_DEV1=1 (supersedes the dev0 overlap): the shared
10229                // expert runs on rank1 — the idle device — same kernels, same
10230                // split program, down row root-resident: bit-identical.
10231                static SHEXP_D1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10232                let shexp_d1 = *SHEXP_D1
10233                    .get_or_init(|| std::env::var("MEMRA_SHEXP_DEV1").as_deref() == Ok("1"))
10234                    && tp.runtime.rank_engine(1).is_some();
10235                // MOE TAIL FUSION M1 (MEMRA_TAIL_ADD3=0 reverts): pre-arm the
10236                // overlap ws + ones row and hand their RAW pointers to the routed
10237                // run — the join add folds the shexp apply into one launch.
10238                static TAIL3: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10239                let tail3 =
10240                    *TAIL3.get_or_init(|| std::env::var("MEMRA_TAIL_ADD3").as_deref() != Ok("0"));
10241                let mut ov_issued = false;
10242                let mut d1_issued = false;
10243                let mut tail_folded = false;
10244                let mut output = if shexp_d1 {
10245                    let rank1 = tp.runtime.rank_engine(1).expect("checked above");
10246                    tp.runtime
10247                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin(
10248                            bank,
10249                            e,
10250                            z,
10251                            sel_d,
10252                            w_d,
10253                            n_used,
10254                            tp.activation_limit,
10255                            || {
10256                                d1_issued =
10257                                    Self::shexp_dev1_issue(e, rank1, m, z, cfg, il, n_embd)?;
10258                                Ok(())
10259                            },
10260                        )?
10261                } else if shexp_ov {
10262                    // Raw sh/ones pointers for the fused tail (persistent statics;
10263                    // pointers stable, no lock held across the routed call). The
10264                    // sh CONTENT is written by the prejoin-issued kernels earlier
10265                    // on e's stream — stream order covers the fused add.
10266                    let post_add = if tail3 {
10267                        Self::shexp_overlap_tail_ptrs(e, m, cfg, n_embd)?
10268                    } else {
10269                        None
10270                    };
10271                    let used_post = post_add.is_some();
10272                    let out = tp
10273                        .runtime
10274                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3(
10275                            bank,
10276                            e,
10277                            z,
10278                            sel_d,
10279                            w_d,
10280                            n_used,
10281                            tp.activation_limit,
10282                            || {
10283                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
10284                                Ok(())
10285                            },
10286                            post_add,
10287                        )?;
10288                    // ov_issued false with post_add armed = an early-return arm
10289                    // (the GRAPH door) skipped the prejoin AND ignored post_add —
10290                    // fall through to the normal shexp add (battery v22 receipt:
10291                    // the strict error here failed every graph-door boot).
10292                    if used_post && ov_issued {
10293                        tail_folded = true; // apply folded into the join add
10294                    }
10295                    out
10296                } else {
10297                    tp.runtime.run_tensor_parallel_routes_nvfp4_device_routed(
10298                        bank,
10299                        e,
10300                        z,
10301                        sel_d,
10302                        w_d,
10303                        n_used,
10304                        tp.activation_limit,
10305                    )?
10306                };
10307                if output.len() != t * n_embd {
10308                    return Err(format!(
10309                        "Step tp routed output has {} values, expected {t}x{n_embd}",
10310                        output.len()
10311                    )
10312                    .into());
10313                }
10314                if tail_folded {
10315                    // shexp already folded into the join add (MOE TAIL FUSION M1)
10316                } else if d1_issued {
10317                    Self::shexp_dev1_apply(e, &mut output, n_embd)?;
10318                } else if ov_issued {
10319                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
10320                } else {
10321                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10322                }
10323                static DR_LOGGED: std::sync::atomic::AtomicU64 =
10324                    std::sync::atomic::AtomicU64::new(0);
10325                let layer_bit = 1u64 << (il as u64 % 64);
10326                if DR_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
10327                    == 0
10328                {
10329                    eprintln!(
10330                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
10331                                 expert_transport={} native_p2p={} router=device \
10332                                 activation=host-canonical accumulation=host-canonical \
10333                                 output=e-device io=device performance_claim=false \
10334                                 (logged once per layer)",
10335                        tp.devices,
10336                        tp.runtime.transport_label(),
10337                        tp.runtime.native_p2p(),
10338                    );
10339                }
10340                return Ok(output);
10341            }
10342            let automatic_ep_device_router = crate::tp::parallel_ep_device_router_enabled()?;
10343            let automatic_ep_q8_act = crate::tp::parallel_ep_q8_act_enabled()?;
10344            let automatic_ep_q8_scope = crate::tp::parallel_ep_q8_scope()?;
10345            crate::tp::parallel_ep_q8_gu_paired_enabled(
10346                automatic_ep_q8_act,
10347                automatic_ep_q8_scope,
10348            )?;
10349            let automatic_ep_q8_active =
10350                automatic_ep_q8_act && t <= crate::tp::NVFP4_EP_Q8_BATCH_CAP;
10351            if automatic_ep_q8_scope.is_some() && !automatic_ep_q8_act {
10352                return Err(
10353                    "MEMRA_PARALLEL_EP_Q8_SCOPE requires MEMRA_PARALLEL_EP_Q8_ACT=1".into(),
10354                );
10355            }
10356            if automatic_ep_q8_act && !automatic_ep_device_router {
10357                return Err(
10358                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires MEMRA_PARALLEL_EP_DEVICE_ROUTER=1".into(),
10359                );
10360            }
10361            if automatic_ep_q8_act && m.step_ep.as_ref().is_none_or(|ep| !ep.nvfp4_device_routes) {
10362                return Err(
10363                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires automatic W4A16 whole-expert EP".into(),
10364                );
10365            }
10366            if t <= crate::tp::NVFP4_EP_DEVICE_ROUTER_BATCH_CAP
10367                && automatic_ep_device_router
10368                && let Some(ep) = &m.step_ep
10369                && ep.nvfp4_device_routes
10370            {
10371                let bank = match &ep.experts {
10372                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
10373                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
10374                        return Err("W4A16 device-routed EP reached an E4M3 expert bank".into());
10375                    }
10376                };
10377                let pairs = t
10378                    .checked_mul(n_used)
10379                    .ok_or("W4A16 device-routed EP pair count overflow")?;
10380                let capacity = crate::tp::NVFP4_EP_DEVICE_BATCH_CAP * n_used;
10381                /// Per-device persistent route scratch: device ordinal -> (armed capacity in
10382                /// pairs, selected-expert rows, route-weight rows). Named because the nested
10383                /// form is unreadable at this depth, not to hide it.
10384                type EpSelwByDevice =
10385                    std::collections::HashMap<usize, (usize, CudaSlice<i32>, CudaSlice<f32>)>;
10386                static EP_SELW: std::sync::Mutex<Option<EpSelwByDevice>> =
10387                    std::sync::Mutex::new(None);
10388                let mut selw = EP_SELW
10389                    .lock()
10390                    .map_err(|_| "automatic EP device-router workspace lock poisoned")?;
10391                let device = e.ctx().ordinal();
10392                let workspaces = selw.get_or_insert_with(Default::default);
10393                if workspaces
10394                    .get(&device)
10395                    .is_none_or(|(cap, ..)| *cap < capacity)
10396                {
10397                    workspaces.insert(
10398                        device,
10399                        (
10400                            capacity,
10401                            e.htod_i32(&vec![0i32; capacity])?,
10402                            e.htod(&vec![0.0f32; capacity])?,
10403                        ),
10404                    );
10405                }
10406                let (_, sel_d, w_d) = workspaces.get_mut(&device).expect("armed above");
10407                let (sf, route_norm) = sigmoid;
10408                e.moe_router_sigmoid_topk_into(
10409                    &logits,
10410                    t,
10411                    n_expert,
10412                    n_used,
10413                    m.active_count(),
10414                    &m.exp_probs_b_dev,
10415                    &m.active_experts_dev,
10416                    sf,
10417                    route_norm,
10418                    sel_d,
10419                    w_d,
10420                )?;
10421                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
10422                crate::moe_sel_dump::refuse_device_only(
10423                    "the automatic W4A16 device-routed EP walk",
10424                )?;
10425                static SHEXP_OV_AUTO: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10426                let shexp_ov = t == 1
10427                    && *SHEXP_OV_AUTO
10428                        .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
10429                let mut ov_issued = false;
10430                let mut output = if shexp_ov {
10431                    ep.runtime
10432                        .run_routed_experts_nvfp4_w4a16_device_routed_prejoin(
10433                            bank,
10434                            e,
10435                            z,
10436                            sel_d,
10437                            w_d,
10438                            t,
10439                            n_used,
10440                            ep.activation_limit,
10441                            || {
10442                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
10443                                Ok(())
10444                            },
10445                        )?
10446                } else {
10447                    ep.runtime.run_routed_experts_nvfp4_w4a16_device_routed(
10448                        bank,
10449                        e,
10450                        z,
10451                        sel_d,
10452                        w_d,
10453                        t,
10454                        n_used,
10455                        ep.activation_limit,
10456                    )?
10457                };
10458                if output.len() != t * n_embd {
10459                    return Err(format!(
10460                        "W4A16 device-routed EP output has {} values, expected \
10461                         {t}x{n_embd}={}",
10462                        output.len(),
10463                        t * n_embd,
10464                    )
10465                    .into());
10466                }
10467                if ov_issued {
10468                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
10469                } else {
10470                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10471                }
10472                static DEVICE_ROUTER_LOGGED: std::sync::atomic::AtomicU64 =
10473                    std::sync::atomic::AtomicU64::new(0);
10474                let layer_bit = 1u64 << (il as u64 % 64);
10475                if DEVICE_ROUTER_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
10476                    & layer_bit
10477                    == 0
10478                {
10479                    eprintln!(
10480                        "[parallel-ep] execute layer={il} tokens={t} devices={:?} \
10481                         router=device expert_transport={} native_p2p={} \
10482                         activation=bf16-rounded accumulation={} output=e-device \
10483                         performance_claim=false (logged once per layer)",
10484                        ep.devices,
10485                        ep.runtime.transport_label(),
10486                        ep.runtime.native_p2p(),
10487                        if automatic_ep_q8_active {
10488                            "token-slot-order-q8"
10489                        } else {
10490                            "token-slot-order"
10491                        },
10492                    );
10493                }
10494                debug_assert!(pairs <= capacity);
10495                return Ok(output);
10496            }
10497            // MEMRA_STEP_TP_TIMING=1: cumulative cost of the host routing seam (the dtoh here
10498            // drains every e-stream op queued since the layer's FFN entry, so this bills the
10499            // router matmul + glue too — the decode-bucket ffn residue decomposes here).
10500            static ROUTE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
10501            static ROUTE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
10502            let route_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
10503            let route_started = route_timing.then(std::time::Instant::now);
10504            let (selected, route_weights, input) = Self::moe_route_sigmoid_with_input(
10505                e,
10506                &logits,
10507                z,
10508                t,
10509                n_embd,
10510                n_expert,
10511                n_used,
10512                m.exp_probs_b.as_deref(),
10513                sigmoid,
10514                m.active_experts.as_deref(),
10515            )?;
10516            if let Some(started) = route_started {
10517                use std::sync::atomic::Ordering;
10518                let ns = ROUTE_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
10519                    + started.elapsed().as_nanos() as u64;
10520                let calls = ROUTE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
10521                if calls.is_multiple_of(430) {
10522                    eprintln!(
10523                        "[moe-route-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
10524                        ns as f64 / 1.0e6,
10525                        ns as f64 / calls as f64 / 1.0e3,
10526                    );
10527                }
10528            }
10529            crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
10530            Self::trace_moe_routes(il, t, &selected, &route_weights)?;
10531            Self::trace_moe_input(e, il, t, n_embd, z)?;
10532            let selected = selected
10533                .iter()
10534                .map(|&expert| expert as usize)
10535                .collect::<Vec<_>>();
10536            if t <= crate::tp::NVFP4_EP_DEVICE_BATCH_CAP
10537                && let Some(ep) = &m.step_ep
10538                && ep.nvfp4_device_routes
10539            {
10540                let bank = match &ep.experts {
10541                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
10542                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
10543                        return Err("W4A16 NVFP4 device EP reached an E4M3 expert bank".into());
10544                    }
10545                };
10546                let mut output = ep.runtime.run_routed_experts_nvfp4_w4a16_device_io(
10547                    bank,
10548                    e,
10549                    z,
10550                    t,
10551                    &selected,
10552                    &route_weights,
10553                    n_used,
10554                    ep.activation_limit,
10555                )?;
10556                if output.len() != t * n_embd {
10557                    return Err(format!(
10558                        "W4A16 NVFP4 EP routed output has {} values, expected \
10559                         {t}x{n_embd}={}",
10560                        output.len(),
10561                        t * n_embd,
10562                    )
10563                    .into());
10564                }
10565                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10566                static W4A16_EP_LOGGED: std::sync::atomic::AtomicU64 =
10567                    std::sync::atomic::AtomicU64::new(0);
10568                let layer_bit = 1u64 << (il as u64 % 64);
10569                if W4A16_EP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
10570                    & layer_bit
10571                    == 0
10572                {
10573                    eprintln!(
10574                        "[step-ep] execute layer={il} tokens={t} devices={:?} \
10575                         expert_transport={} native_p2p={} activation=bf16-rounded \
10576                        accumulation={} output=e-device \
10577                         performance_claim=false (logged once per layer)",
10578                        ep.devices,
10579                        ep.runtime.transport_label(),
10580                        ep.runtime.native_p2p(),
10581                        if t == 1 {
10582                            "owner-grouped-rank-order"
10583                        } else {
10584                            "token-slot-order"
10585                        },
10586                    );
10587                }
10588                return Ok(output);
10589            }
10590            // Device-IO routes (t=1): the layer input goes to the ranks as a device row and the
10591            // combined output comes back as an e-context row — no host round-trip, no host
10592            // stream sync. Program bytes identical to the host-IO twin (dtoh/htod and dtod
10593            // both preserve f32 bits), gated by greedy token identity.
10594            if t == 1
10595                && crate::tp::step_nvfp4_dev_routes_enabled()?
10596                && let Some(tp) = &m.step_tp
10597                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
10598            {
10599                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_device_io(
10600                    bank,
10601                    e,
10602                    z,
10603                    &selected,
10604                    &route_weights,
10605                    n_used,
10606                    tp.activation_limit,
10607                )?;
10608                if output.len() != t * n_embd {
10609                    return Err(format!(
10610                        "Step tp routed output has {} values, expected {t}x{n_embd}",
10611                        output.len()
10612                    )
10613                    .into());
10614                }
10615                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10616                static IO_LOGGED: std::sync::atomic::AtomicU64 =
10617                    std::sync::atomic::AtomicU64::new(0);
10618                let layer_bit = 1u64 << (il as u64 % 64);
10619                if IO_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
10620                    == 0
10621                {
10622                    eprintln!(
10623                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
10624                                 expert_transport={} native_p2p={} activation=host-canonical \
10625                                 accumulation=host-canonical output=e-device io=device \
10626                                 performance_claim=false (logged once per layer)",
10627                        tp.devices,
10628                        tp.runtime.transport_label(),
10629                        tp.runtime.native_p2p(),
10630                    );
10631                }
10632                return Ok(output);
10633            }
10634            let (routed, mode, devices, transport, native_p2p) = if let Some(tp) = &m.step_tp {
10635                (
10636                    match &tp.experts {
10637                        crate::hybrid::StepTpExpertBank::E4m3(bank) => {
10638                            tp.runtime.run_tensor_parallel_routes(
10639                                bank,
10640                                &input,
10641                                t,
10642                                &selected,
10643                                &route_weights,
10644                                n_used,
10645                            )?
10646                        }
10647                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => {
10648                            if t == 1 && crate::tp::step_nvfp4_dev_routes_enabled()? {
10649                                tp.runtime.run_tensor_parallel_routes_nvfp4_device(
10650                                    bank,
10651                                    &input,
10652                                    &selected,
10653                                    &route_weights,
10654                                    n_used,
10655                                    tp.activation_limit,
10656                                )?
10657                            } else {
10658                                tp.runtime.run_tensor_parallel_routes_nvfp4(
10659                                    bank,
10660                                    &input,
10661                                    t,
10662                                    &selected,
10663                                    &route_weights,
10664                                    n_used,
10665                                    tp.activation_limit,
10666                                )?
10667                            }
10668                        }
10669                    },
10670                    "tp",
10671                    &tp.devices,
10672                    tp.runtime.transport_label(),
10673                    tp.runtime.native_p2p(),
10674                )
10675            } else {
10676                let ep = m
10677                    .step_ep
10678                    .as_ref()
10679                    .ok_or("Step distributed runtime has no EP or TP state")?;
10680                (
10681                    match &ep.experts {
10682                        crate::hybrid::StepEpExpertBank::E4m3(bank) => {
10683                            ep.runtime.run_routed_experts(
10684                                bank,
10685                                &input,
10686                                t,
10687                                &selected,
10688                                &route_weights,
10689                                n_used,
10690                                ep.activation_limit,
10691                            )?
10692                        }
10693                        crate::hybrid::StepEpExpertBank::Nvfp4(bank) => {
10694                            ep.runtime.run_routed_experts_nvfp4(
10695                                bank,
10696                                &input,
10697                                t,
10698                                &selected,
10699                                &route_weights,
10700                                n_used,
10701                                ep.activation_limit,
10702                            )?
10703                        }
10704                    },
10705                    if ep.configured_by_tp { "tp-ep" } else { "ep" },
10706                    &ep.devices,
10707                    ep.runtime.transport_label(),
10708                    ep.runtime.native_p2p(),
10709                )
10710            };
10711            if routed.len() != t * n_embd {
10712                return Err(format!(
10713                    "Step {mode} routed output has {} values, expected {t}x{n_embd}",
10714                    routed.len()
10715                )
10716                .into());
10717            }
10718            let mut output = e.htod(&routed)?;
10719            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
10720            // Once per layer per process: the topology contract line is a boot receipt, not a
10721            // per-token trace — 4520 of these per 64-token run measured as real decode wall.
10722            static STEP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
10723            let layer_bit = 1u64 << (il as u64 % 64);
10724            if STEP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
10725                == 0
10726            {
10727                eprintln!(
10728                    "[step-{mode}] execute layer={il} tokens={t} devices={devices:?} \
10729                     expert_transport={transport} native_p2p={native_p2p} \
10730                     activation={} accumulation={} output={} \
10731                     performance_claim=false (logged once per layer)",
10732                    if let Some(ep) = &m.step_ep {
10733                        ep.runtime.expert_activation_label()
10734                    } else {
10735                        "host-canonical"
10736                    },
10737                    if let Some(ep) = &m.step_ep {
10738                        ep.runtime.expert_accumulation_label()
10739                    } else {
10740                        "host-canonical"
10741                    },
10742                    if let Some(ep) = &m.step_ep {
10743                        ep.runtime.expert_output_label()
10744                    } else {
10745                        "host-accumulated"
10746                    },
10747                );
10748                if let Some(ep) = &m.step_ep
10749                    && let Some(limit) = ep.activation_limit
10750                {
10751                    eprintln!(
10752                        "[step-ep-clamp] execute layer={il} tokens={t} routed_clamp={limit} \
10753                             formula=min-silu-times-clamped-up performance_claim=false"
10754                    );
10755                }
10756            }
10757            return Ok(output);
10758        }
10759        if Self::sigmoid_resident_dev_eligible(e, m, cfg, sliding_gated_moe) {
10760            let moe = cfg.moe.as_ref().unwrap();
10761            let n_expert = moe.expert_count as usize;
10762            let n_used = moe.expert_used_count as usize;
10763            let sigmoid = cfg.sigmoid_router().unwrap();
10764            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
10765            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
10766            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
10767        }
10768        // GROUPED MoE PREFILL, sigmoid-router class (glm5_next), MEMRA_MOE_GROUPED_PREFILL
10769        // default ON since 2026-08-29 (owner-accepted flip; =0 rollback seam; receipts on the
10770        // flag helper). The prefill-gap attribution (research/glm53-flash-bringup-20260827/
10771        // prefill-gap-20260829/PREFILL-GAP.md §1.1) measured this arch prefilling every prompt
10772        // token through the decode program: 49 launches per token-layer, ~8.4M launches and
10773        // 4.76 GB of expert-weight VRAM re-reads per token across 42 layers per 4096-token
10774        // chunk, because every batched arm is predicate-denied for sigmoid-router archs. This
10775        // arm is the composition of qualified ingredients: the m-invariant router + sigmoid
10776        // host oracle (routing sel/w BIT-identical to the sequential arm by construction),
10777        // host token-sort by expert (the moe_align_block_size shape), one grouped NVFP4
10778        // tensor-core GEMM per projection (the step37 grouped-prime kernel class via
10779        // `moe_f16_grouped`, generalized to the single-device resident slab), the PRE-clamped
10780        // SwiGLU epilogue and the per-expert weight_scale_2 macro fold the fused-epilogue lane
10781        // gated for this family. Keyed on `prefill` (only the _prefill twin sets it) so decode,
10782        // spec verify and the exact-16 batched-decode tier keep their dispatch class, and on
10783        // `t > MOE_DEV_MAX_T` so t<=16 stays on the per-token program (grouped/pairs prefill
10784        // classes start at 17, same seam as the softmax pairs arm).
10785        // ENGAGEMENT RECEIPT: the announce below prints in BOTH arms (flag on and off), once
10786        // per process, so an A/B grep distinguishes engagement without the line itself being
10787        // an arm-local cost (the step37 engagement-receipt trap: prove the path RAN before
10788        // attributing a number to it).
10789        if prefill && t > MOE_DEV_MAX_T && cfg.sigmoid_router().is_some() && cfg.glm5.is_some() {
10790            // Once per process PER FLAG VALUE (bit 0 = off, bit 1 = on): a server boot prints
10791            // exactly one line, and a gate process that flips the flag shows both arms.
10792            static GPF_ANNOUNCED: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
10793            let enabled = moe_grouped_prefill_enabled();
10794            let bit = 1u8 << u8::from(enabled);
10795            if GPF_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
10796                eprintln!(
10797                    "[moe-grouped-prefill] flag={} t={t} il={il} (announce printed in both \
10798                     arms; engagement is the per-layer execute line + the dispatch counter)",
10799                    if enabled { "on" } else { "off" },
10800                );
10801            }
10802            if enabled
10803                && let Some(out) = Self::moe_ffn_grouped_prefill_sigmoid(e, m, z, t, cfg, il)?
10804            {
10805                return Ok(out);
10806            }
10807        }
10808        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
10809        // current caller into this research arm; the naked default stays on the established path.
10810        if t > 1 && moe_grouped_enabled(cfg, prefill) {
10811            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
10812            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
10813            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
10814            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
10815            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
10816            if std::env::var("MEMRA_MOE_GATE").is_ok() {
10817                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
10818                let g_host = e.dtoh(&grouped_out)?;
10819                let s_host = e.dtoh(&seq_out)?;
10820                let g_bytes: &[u8] = unsafe {
10821                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
10822                };
10823                let s_bytes: &[u8] = unsafe {
10824                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
10825                };
10826                if g_bytes == s_bytes {
10827                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
10828                } else {
10829                    let diffs = g_host
10830                        .iter()
10831                        .zip(s_host.iter())
10832                        .enumerate()
10833                        .filter(|(_, (a, b))| a != b)
10834                        .count();
10835                    let maxdiff = g_host
10836                        .iter()
10837                        .zip(s_host.iter())
10838                        .map(|(a, b)| (a - b).abs())
10839                        .fold(0.0f32, f32::max);
10840                    panic!(
10841                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
10842                        g_host.len()
10843                    );
10844                }
10845            }
10846            return Ok(grouped_out);
10847        }
10848        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block, vrows)
10849    }
10850
10851    /// FNV-1a over the bits, plus the two things a checksum alone cannot say: how many entries
10852    /// are nonzero, and the largest magnitude. Take 5's `nz=16384/16384 absmax=0` is what an
10853    /// all-NaN buffer looks like, and `nz=0` is what a never-written one looks like; neither is
10854    /// distinguishable from "wrong" by a hash.
10855    fn sum_f32(v: &[f32]) -> String {
10856        let mut h: u64 = 0xcbf2_9ce4_8422_2325;
10857        let (mut nz, mut absmax) = (0usize, 0f32);
10858        for x in v {
10859            h ^= x.to_bits() as u64;
10860            h = h.wrapping_mul(0x100_0000_01b3);
10861            if *x != 0.0 {
10862                nz += 1;
10863            }
10864            if x.abs() > absmax {
10865                absmax = x.abs();
10866            }
10867        }
10868        format!("0x{h:016x}/nz{nz}of{}/max{absmax:.4e}", v.len())
10869    }
10870
10871    fn sum_i8(v: &[i8]) -> String {
10872        let mut h: u64 = 0xcbf2_9ce4_8422_2325;
10873        let mut nz = 0usize;
10874        for x in v {
10875            h ^= *x as u8 as u64;
10876            h = h.wrapping_mul(0x100_0000_01b3);
10877            if *x != 0 {
10878                nz += 1;
10879            }
10880        }
10881        format!("0x{h:016x}/nz{nz}of{}", v.len())
10882    }
10883
10884    /// THE ACTIVATION THE CONSUMER ACTUALLY READS, checksummed immediately before its launch, on
10885    /// both arms. Take 9 established that the two arms receive identical `sel`/`w`/macros/strides/
10886    /// limit, so the divergence is inside the consumer call — and the leading hypothesis is that
10887    /// the q8_1 activation the device pair reads is stale or unwritten rather than this token's
10888    /// hidden state. `z` is the shared f32 input: if `z` agrees across arms but `zq`/`zd` do not,
10889    /// the quantize is the seam; if all three agree, the kernels are.
10890    #[allow(clippy::too_many_arguments)] // allow: a diagnostic line's fields are its whole purpose
10891    fn trace_moe_act(
10892        e: &Engine,
10893        arm: &str,
10894        il: u16,
10895        t: usize,
10896        z: &CudaSlice<f32>,
10897        zq: &CudaSlice<i8>,
10898        zd: &CudaSlice<f32>,
10899    ) {
10900        if !crate::glm5_graph_trace_on()
10901            || crate::glm5_graph_capture_open()
10902            || !crate::glm5_trace_take_slot("act", arm, il)
10903        {
10904            return;
10905        }
10906        let (zs, qs, ds) = (e.dtoh(z), e.dtoh_i8(zq), e.dtoh(zd));
10907        match (zs, qs, ds) {
10908            (Ok(zv), Ok(qv), Ok(dv)) => eprintln!(
10909                "[glm5-vrows-act] arm={arm} il={il} t={t} z={} zq={} zd={}",
10910                Self::sum_f32(&zv),
10911                Self::sum_i8(&qv),
10912                Self::sum_f32(&dv),
10913            ),
10914            _ => eprintln!("[glm5-vrows-act] arm={arm} il={il} readback failed"),
10915        }
10916    }
10917
10918    /// The routed-MoE output of ONE layer, on both arms — the other end of the same seam.
10919    fn trace_moe_out(e: &Engine, arm: &str, il: u16, out: &CudaSlice<f32>) {
10920        if !crate::glm5_graph_trace_on()
10921            || crate::glm5_graph_capture_open()
10922            || !crate::glm5_trace_take_slot("out", arm, il)
10923        {
10924            return;
10925        }
10926        match e.dtoh(out) {
10927            Ok(v) => eprintln!(
10928                "[glm5-vrows-out] arm={arm} il={il} out={}",
10929                Self::sum_f32(&v)
10930            ),
10931            Err(err) => eprintln!("[glm5-vrows-out] arm={arm} il={il} readback failed ({err})"),
10932        }
10933    }
10934
10935    /// One line per (arm, routed layer) under `MEMRA_GLM5_GRAPH_TRACE`, printed by BOTH the
10936    /// device-table arm and the host-oracle arm so a box run can diff them field for field.
10937    ///
10938    /// It carries the per-expert MACRO SCALES for the SELECTED experts, not just a `macros=bool`.
10939    /// That is deliberate: `HostExps::macro_scale(e)` is exactly `macros[e]` (model.rs), so the
10940    /// device table kernel's `mac_g[ex]` and the host loop's `macro_scale(ex)` are the same lookup
10941    /// — which the rig gate confirmed — and the only way a difference can still exist is if the
10942    /// two arms are handed different SELECTIONS or different PLANES. Printing the values makes
10943    /// that visible instead of inferred.
10944    #[allow(clippy::too_many_arguments)] // allow: a diagnostic line's fields are its whole purpose
10945    fn dump_moe_t1_inputs(
10946        e: &Engine,
10947        arm: &str,
10948        m: &MoeWeights,
10949        cfg: &ModelConfig,
10950        il: u16,
10951        t: usize,
10952        n_used: usize,
10953        n_expert: usize,
10954        sel: &[u32],
10955        w: &[f32],
10956    ) {
10957        if !crate::glm5_trace_take_slot("shape", arm, il) {
10958            return;
10959        }
10960        let mac = |x: &crate::model::HostExps| -> Vec<f32> {
10961            sel.iter().map(|&ex| x.macro_scale(ex as usize)).collect()
10962        };
10963        eprintln!(
10964            "[glm5-vrows-t1] arm={arm} dev={} il={il} t={t} n_used={n_used} n_pairs={} \
10965             n_expert={n_expert} limit={:?} gu_il={:?} rp={:?} qtypes=({},{},{}) row_bytes=({},{},{}) \
10966             strides=({},{},{}) macros={} sel={sel:?} w={w:?} mac_g={:?} mac_u={:?} mac_d={:?}",
10967            e.ctx().ordinal(),
10968            t * n_used,
10969            cfg.clamp_exp_at(il as u32),
10970            m.dev_exps.as_ref().map(|d| d.gu_il),
10971            m.dev_exps.as_ref().map(|d| d.rp),
10972            m.gate_exps.qtype,
10973            m.up_exps.qtype,
10974            m.down_exps.qtype,
10975            m.gate_exps.row_bytes,
10976            m.up_exps.row_bytes,
10977            m.down_exps.row_bytes,
10978            m.gate_exps.expert_stride,
10979            m.up_exps.expert_stride,
10980            m.down_exps.expert_stride,
10981            m.gate_exps.macros.is_some(),
10982            mac(&m.gate_exps),
10983            mac(&m.up_exps),
10984            mac(&m.down_exps),
10985        );
10986    }
10987
10988    /// Would the T=1 DEVICE-TABLE MoE arm (`vrows_t1_dev`) fire for this layer? The
10989    /// decode-graph door must answer this BEFORE it opens a capture region: a layer that falls
10990    /// through to the host readback would issue a `cuStreamSynchronize` inside the capture and
10991    /// take down the whole request instead of yielding to the eager walk. The conjuncts below
10992    /// are the layer-shaped half of the dispatch predicate (the process-shaped half — traces,
10993    /// observers, the NVMe promotion, `MEMRA_HTOD_DIET` — is checked once by the door itself).
10994    pub(crate) fn glm5_t1_dev_moe_ready(
10995        e: &Engine,
10996        layer: &crate::hybrid::HybridLayer,
10997        cfg: &ModelConfig,
10998        il: usize,
10999    ) -> bool {
11000        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
11001            // A dense FFN carries no router and no selection: nothing to read back, so the
11002            // layer is capture-clean on its own.
11003            return true;
11004        };
11005        let Some(moe) = cfg.moe.as_ref() else {
11006            return false;
11007        };
11008        let slab_local = m
11009            .dev_exps
11010            .as_ref()
11011            .is_some_and(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
11012        slab_local
11013            && m.has_uniform_expert_layout()
11014            && moe_q8_enabled_for_model(cfg, m)
11015            && moe.expert_used_count as usize <= 8
11016            && matches!(cfg.clamp_exp_at(il as u32), Some(SwigluClamp::Pre(l)) if l > 1e-6)
11017            && !crate::cpu_experts::configured()
11018    }
11019
11020    fn sigmoid_resident_dev_eligible(
11021        e: &Engine,
11022        m: &MoeWeights,
11023        cfg: &ModelConfig,
11024        sliding_gated_moe: bool,
11025    ) -> bool {
11026        let Some(moe) = cfg.moe.as_ref() else {
11027            return false;
11028        };
11029        // Cached once per process: this predicate runs per MoE layer per decode step, and five
11030        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
11031        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11032        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
11033            std::env::var("MEMRA_MOE_STATS").is_ok()
11034                || std::env::var("MEMRA_MOE_TRACE").is_ok()
11035                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
11036                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
11037                || std::env::var("MEMRA_MOE_GATE").is_ok()
11038        });
11039        let resident_layout_supported = m.dev_exps.as_ref().is_some_and(|dev| {
11040            if dev.dev != e.ctx().ordinal() {
11041                return false;
11042            }
11043            let q8 = moe_q8_enabled_for_model(cfg, m);
11044            let fp8 = dev.fp8_blk.is_some()
11045                && m.gate_exps.qtype == crate::QT_F8_E4M3_BLK
11046                && m.up_exps.qtype == crate::QT_F8_E4M3_BLK
11047                && m.down_exps.qtype == crate::QT_F8_E4M3_BLK;
11048            q8 || fp8
11049        });
11050        sliding_gated_moe
11051            && sigmoid_router_enabled()
11052            && moe_dev_enabled()
11053            && moe_slab_enabled()
11054            && !observation_mode
11055            && moe.expert_used_count <= 8
11056            && m.has_uniform_expert_layout()
11057            && m.gate_exps.macros.is_none()
11058            && m.up_exps.macros.is_none()
11059            && m.down_exps.macros.is_none()
11060            && !m.has_macros
11061            && resident_layout_supported
11062    }
11063
11064    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
11065    pub(crate) fn moe_ffn_sequential(
11066        e: &Engine,
11067        m: &MoeWeights,
11068        z: &CudaSlice<f32>,
11069        t: usize,
11070        cfg: &ModelConfig,
11071        il: u16,
11072        max_block: usize,
11073    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11074        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block, false)
11075    }
11076
11077    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
11078    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
11079    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
11080    fn moe_router_logits(
11081        e: &Engine,
11082        m: &MoeWeights,
11083        z: &CudaSlice<f32>,
11084        t: usize,
11085        cfg: &ModelConfig,
11086    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11087        if t < PRIME_MIN_T {
11088            // Decode and speculative verify use one fixed per-row reduction program.
11089            if crate::router_kernel_on() {
11090                e.router_gemv(
11091                    m.gate_inp.float_data(),
11092                    z,
11093                    cfg.n_embd as usize,
11094                    m.gate_exps.n_expert,
11095                    t,
11096                )
11097            } else {
11098                e.matmul_decode_exact(&m.gate_inp, z, t)
11099            }
11100        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
11101            e.router_gemv(
11102                m.gate_inp.float_data(),
11103                z,
11104                cfg.n_embd as usize,
11105                m.gate_exps.n_expert,
11106                t,
11107            )
11108        } else {
11109            e.matmul(&m.gate_inp, z, t)
11110        }
11111    }
11112
11113    /// Append the host-visible router selection for one layer/forward when calibration tracing is
11114    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
11115    /// trace is independent of the dispatch optimization selected for the forward.
11116    /// `MEMRA_MOE_WEIGHT_TRACE` is also the co-activation measurement input of
11117    /// LAW:coactivation-expert-placement (lane/glm5-ep-place: rows ride this existing host
11118    /// readback — zero new device syncs; `glm5-tp-gate` arm T holds the ON-identity +
11119    /// row-count bar on the glm5 walks).
11120    fn trace_moe_routes(
11121        il: u16,
11122        t: usize,
11123        sel_all: &[u32],
11124        weights: &[f32],
11125    ) -> Result<(), Box<dyn std::error::Error>> {
11126        use std::io::Write as _;
11127        // MEMRA_MOE_SEL_DUMP (lane/moe-coactivation-20260902): the binary per-token twin of
11128        // the two text taps below. Unlike them it diverts no dispatch (it is in no
11129        // `observe_routes` conjunct) and so sees the served arm's rows; the device-routed
11130        // single-device arms record themselves through `moe_sel_dump::record_device`. One
11131        // OnceLock read when unset.
11132        crate::moe_sel_dump::record_host(il, t, sel_all, weights)?;
11133        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
11134            let mut f = std::fs::OpenOptions::new()
11135                .create(true)
11136                .append(true)
11137                .open(path)?;
11138            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
11139            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
11140        }
11141        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
11142            let mut f = std::fs::OpenOptions::new()
11143                .create(true)
11144                .append(true)
11145                .open(path)?;
11146            let pairs: Vec<String> = sel_all
11147                .iter()
11148                .zip(weights)
11149                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
11150                .collect();
11151            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
11152        }
11153        Ok(())
11154    }
11155
11156    #[allow(clippy::too_many_arguments)]
11157    fn trace_sigmoid_router_logits(
11158        e: &Engine,
11159        il: u16,
11160        t: usize,
11161        n_expert: usize,
11162        n_used: usize,
11163        logits: &CudaSlice<f32>,
11164        m: &MoeWeights,
11165        (scaling_factor, route_norm): (f32, bool),
11166    ) -> Result<(), Box<dyn std::error::Error>> {
11167        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
11168            return Ok(());
11169        }
11170        let logits = e.dtoh(logits)?;
11171        let active: Vec<u8> = m
11172            .active_experts
11173            .as_ref()
11174            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
11175            .unwrap_or_else(|| vec![1; n_expert]);
11176        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
11177        crate::sigrouter_contract::capture_served_logits(
11178            il as u32,
11179            t,
11180            n_expert,
11181            n_used,
11182            scaling_factor,
11183            route_norm,
11184            &active,
11185            &bias,
11186            &logits,
11187        )?;
11188        Ok(())
11189    }
11190
11191    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
11192    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
11193    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
11194    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
11195    fn trace_moe_input(
11196        e: &Engine,
11197        il: u16,
11198        t: usize,
11199        n_embd: usize,
11200        z: &CudaSlice<f32>,
11201    ) -> Result<(), Box<dyn std::error::Error>> {
11202        use std::io::Write as _;
11203        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
11204            return Ok(());
11205        };
11206        let values = active_matrix_values(z.len(), t, n_embd, "MoE input trace activation")?;
11207        let host = e.dtoh_view(&z.slice(0..values))?;
11208        let bytes = unsafe {
11209            std::slice::from_raw_parts(
11210                host.as_ptr().cast::<u8>(),
11211                host.len() * std::mem::size_of::<f32>(),
11212            )
11213        };
11214        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
11215        let mut state = state
11216            .lock()
11217            .map_err(|_| "MoE input trace writer lock is poisoned")?;
11218        if state.is_none() {
11219            let dir = std::path::PathBuf::from(&dir);
11220            std::fs::create_dir_all(&dir)?;
11221            let index = std::fs::OpenOptions::new()
11222                .create(true)
11223                .append(true)
11224                .open(dir.join("index.jsonl"))?;
11225            *state = Some(MoeInputTraceWriter {
11226                dir,
11227                index,
11228                payloads: std::collections::HashMap::new(),
11229            });
11230        }
11231        let writer = state.as_mut().unwrap();
11232        if writer.dir != std::path::Path::new(&dir) {
11233            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
11234        }
11235        let file_name = format!("layer-{il:03}.f32");
11236        if !writer.payloads.contains_key(&il) {
11237            let payload = std::fs::OpenOptions::new()
11238                .create(true)
11239                .append(true)
11240                .open(writer.dir.join(&file_name))?;
11241            let offset = payload.metadata()?.len();
11242            writer.payloads.insert(il, (payload, offset));
11243        }
11244        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
11245        let row_offset = *offset;
11246        payload.write_all(bytes)?;
11247        *offset += bytes.len() as u64;
11248        writeln!(
11249            writer.index,
11250            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
11251             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
11252             \"payload_bytes\":{}}}",
11253            bytes.len()
11254        )?;
11255        Ok(())
11256    }
11257
11258    #[allow(clippy::too_many_arguments)]
11259    #[allow(clippy::too_many_arguments)]
11260    // allow: the parameter list mirrors its moe_ffn_inner caller's dispatch contract
11261    pub(crate) fn moe_ffn_sequential_zq8(
11262        e: &Engine,
11263        m: &MoeWeights,
11264        z: &CudaSlice<f32>,
11265        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11266        t: usize,
11267        cfg: &ModelConfig,
11268        il: u16,
11269        max_block: usize,
11270        vrows: bool,
11271    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11272        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
11273        let moe = cfg.moe.as_ref().unwrap();
11274        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
11275        let n_expert = moe.expert_count as usize; // 256
11276        let n_used = moe.expert_used_count as usize; // 8
11277        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
11278
11279        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
11280        debug_assert_eq!(m.gate_exps.in_f, n_embd);
11281        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
11282        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
11283        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
11284        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
11285
11286        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
11287        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
11288        let lim_exp = cfg.clamp_exp_at(il as u32);
11289        let lim_shexp = cfg.clamp_shexp_at(il as u32);
11290        let use_cache = Engine::moe_cache_enabled();
11291        let uniform_experts = m.has_uniform_expert_layout();
11292        let moe_q8 = uniform_experts && moe_q8_enabled_for_model(cfg, m);
11293        // Experimental secondary backend: complete experts already resident in the SLRU stay on
11294        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
11295        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
11296        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
11297        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
11298        // commands and CI have no llama.cpp or OpenMP dependency.
11299        let cpu_expert_requested = crate::cpu_experts::configured();
11300        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
11301            return Err(std::io::Error::other(
11302                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
11303            )
11304            .into());
11305        }
11306        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
11307        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
11308        // Those backends are each deterministic but are different numeric configurations, so a
11309        // later prefill eviction can change greedy output. Freeze after the first real prefill;
11310        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
11311        // staging below and cannot change backend assignment.
11312        let freeze_cpu_residency = cpu_expert_requested
11313            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
11314        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
11315            .ok()
11316            .and_then(|value| value.parse::<usize>().ok())
11317            .is_some_and(|tokens| tokens > 0);
11318        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
11319            e.freeze_moe_cache();
11320        }
11321        let cache_frozen = use_cache && e.moe_cache_frozen();
11322        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
11323
11324        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
11325        // cannot change logits, selected expert ids, or routing weights.
11326        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
11327        if let Some(sig) = cfg.sigmoid_router() {
11328            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
11329        }
11330
11331        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
11332        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
11333        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
11334        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
11335        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
11336        // per-token host stall that dominated the 35B decode wall after stages 1+2.
11337        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
11338        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
11339        // only difference is where sel/w/pointers are READ from (device instead of params).
11340        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
11341        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
11342        // Any non-resident layer falls through to host routing + the gdec/sequential path.
11343        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
11344        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
11345        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
11346        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
11347        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
11348        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
11349        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
11350        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
11351        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
11352        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
11353        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
11354        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
11355        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
11356        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
11357        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
11358        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
11359        // now rides the dev loop below (same kernels per token as decode); pairs serves real
11360        // prefill (t >= 16, where spec never verifies).
11361        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
11362        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
11363        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
11364        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
11365        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
11366        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
11367        // ride the macro-aware sequential/staged paths below or every expert output is off by
11368        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
11369        let no_exp_macros = m.gate_exps.macros.is_none()
11370            && m.up_exps.macros.is_none()
11371            && m.down_exps.macros.is_none();
11372        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
11373        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
11374        // so it cannot even see the per-layer limit.
11375        if cfg.sigmoid_router().is_none()
11376            && cfg.m3.is_none()
11377            && cfg.hy3.is_none()
11378            && !cfg.swiglu_clamped_at(il as u32)
11379            && no_exp_macros
11380            // > MOE_DEV_MAX_T, not >= PRIME_MIN_T: t==16 is a decode width under the
11381            // exact-16 tier and rides the dev per-token program (see MOE_DEV_MAX_T);
11382            // pairs serves real prefill from 17 up.
11383            && t > MOE_DEV_MAX_T
11384            && m.dev_exps.is_some()
11385            && moe_q8_enabled_for_model(cfg, m)
11386            && std::env::var("MEMRA_MOE_PAIRS")
11387                .map(|v| v != "0")
11388                .unwrap_or(true)
11389            && std::env::var("MEMRA_MOE_STATS").is_err()
11390        {
11391            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
11392        }
11393
11394        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
11395        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
11396        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
11397        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
11398        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
11399        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
11400        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
11401        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
11402        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
11403        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
11404        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
11405        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
11406        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
11407        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
11408        // Keyed off sigmoid_router() so arch #4 is denied by construction.
11409        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
11410        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
11411        let dev_ok = uniform_experts
11412            && cfg.sigmoid_router().is_none()
11413            && cfg.m3.is_none()
11414            && cfg.hy3.is_none()
11415            && !cfg.swiglu_clamped_at(il as u32);
11416        // Observation modes must route through the host-visible selection below. Otherwise a fully
11417        // resident layer returns through device dispatch before its trace/stats row is recorded,
11418        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
11419        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
11420            || std::env::var("MEMRA_MOE_TRACE").is_ok()
11421            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
11422            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
11423        if dev_ok
11424            && t <= MOE_DEV_MAX_T
11425            && m.dev_exps.is_some()
11426            && n_used <= 8
11427            && moe_dev_enabled()
11428            && !observe_routes
11429        {
11430            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
11431        }
11432        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
11433            let row_ok = e.with_moe_cache(max_block, |c, eng| {
11434                if moe_prewarm_enabled() {
11435                    c.prewarm_layer(il, m, eng)?;
11436                }
11437                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
11438            })?;
11439            if row_ok {
11440                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
11441            }
11442        }
11443
11444        // SLAB-LOCAL RESIDENT ARM bases, hoisted above the router (lane/glm5-moe-loc door D):
11445        // whether the layer can run the DEVICE vrows table build decides HOW it routes, and
11446        // that has to be settled before the router runs. Pure immutable pointer reads with no
11447        // side effects, so the hoist changes nothing for any other arm; the full rationale for
11448        // the arm itself is at the `slab_fused_may_fire` predicate below.
11449        let slab_local = m
11450            .dev_exps
11451            .as_ref()
11452            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
11453        let slab_bases = slab_local.map(|d| {
11454            use cudarc::driver::DevicePtr;
11455            let s = e.stream();
11456            let (pg, _g0) = d.gate.device_ptr(&s);
11457            let (pu, _g1) = d.up.device_ptr(&s);
11458            let (pd, _g2) = d.down.device_ptr(&s);
11459            (pg, pu, pd)
11460        });
11461        // memra#147: the slab those bases name is split-plane; its readers are told so.
11462        let slab_rp = slab_local.is_some_and(|d| d.rp);
11463        // HOISTED above the router (lane/b200-glm5-graph-20260902): `promote_worker_h2d` reads
11464        // the HOST selection, so the t=1 device-table arm below must be able to deny itself when
11465        // the NVMe worker promotion is live. It was safe to compute this after the router while
11466        // door D required `t >= 2` (the promotion requires `t == 1`); the decode-graph door
11467        // lifts that exclusion, so the conjunct has to exist before the arm is chosen.
11468        let worker_disk_prefetch =
11469            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
11470        let promote_worker_h2d =
11471            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
11472        // DOOR `MEMRA_GLM5_DECODE_GRAPH` (lane/b200-glm5-graph-20260902, default ON since
11473        // 2026-09-04): the T=1
11474        // DECODE twin of door D. On the glm5_next serving shape the per-MoE-layer selection is
11475        // read back through a pinned stage plus a full `cuStreamSynchronize`
11476        // (`Engine::moe_router_sigmoid_topk_host`) for the sole purpose of letting the HOST
11477        // compute `base + ex*stride` — 42 device-wide drains per token, and the reason the T=1
11478        // walk cannot be captured as a CUDA graph at all (a capture region admits no sync and no
11479        // pageable HtoD). This arm keeps the selection on device and builds the same pointer and
11480        // scale tables with `moe_vrows_tables_from_sel`, then runs the verify-rows kernel pair at
11481        // `n_pairs = n_used`.
11482        //
11483        // BIT IDENTITY, and why it is a claim about ONE program: the rows twins
11484        // (`moe_gate_up_preclamp8_q8_rows` / `moe_down8_fma_q8_rows`) are the per-row form of the
11485        // fused epilogue, which is itself the per-token form of the sequential
11486        // `qmatvec_expert_q8` + `ffn_act_lim` + `axpy_into` chain — same g-strided dots, same
11487        // pre-clamped SwiGLU expression, same slot-ordered `__fmaf_rn` accumulation, same macro
11488        // folds in the same places. The table VALUES are term-for-term the host loop's (exact
11489        // integer `base + ex*stride`, the same macro planes, one IEEE-754 product for
11490        // `w * macro_down`). At t=1 this is the single-row case of the claim the vrest lane
11491        // already gated at t>=2, and `glm5_decode_graph_gate` re-proves it per token per layer.
11492        //
11493        // FAIL-CLOSED: every host-visible consumer of the selection must be disarmed, exactly as
11494        // door D requires, PLUS `promote_worker_h2d` (which door D could ignore because it
11495        // requires t == 1 and door D required t >= 2).
11496        //
11497        // KEYED ON THE OPEN CAPTURE, NOT THE DOOR ENV (2026-09-03, box takes 4-11). This arm
11498        // exists for ONE reason: a host sel/w readback is illegal inside a stream-capture
11499        // region, so a captured T=1 walk has to build its tables on device. Outside a capture
11500        // region nothing needs it, and keying it on `glm5_decode_graph_on()` made the door
11501        // change the program on paths where it captures NOTHING — the eager fall-through, a
11502        // latched stage, a refused stage, and the `MEMRA_GLM5_GRAPH_TRACE` split walk. That
11503        // broke the door's own contract ("every refusal falls through byte-identically") and it
11504        // is what box take 5 actually caught: the 1.356879e19 blow-up printed on an
11505        // `eager-run` trace line, which `hyper_range_decode` only reaches when the graph arm
11506        // was NOT taken. Eleven box takes read that as a capture/replay defect; it was a walk
11507        // running the device-table arm with no graph in sight.
11508        //
11509        // `glm5_graph_capture_open()` is set for the duration of `capture_one`'s recording and
11510        // clear everywhere else, so the captured body gets the device tables it requires and
11511        // every other path gets the shipped host-oracle program, unchanged.
11512        let vrows_t1_dev = t == 1
11513            // `MEMRA_GLM5_VROWS_T1_DEV` (default OFF, gate harness) forces the arm on with NO
11514            // capture and NO graph anywhere in the walk. It is the bisect cell eleven box takes
11515            // never had: door ON conflated "device-table MoE at T=1" with "capture and replay",
11516            // so a mismatch could not be attributed. With this the two questions are asked one
11517            // at a time, and the answer to the first one is a plain decode run.
11518            && (crate::glm5_graph_capture_open() || crate::glm5_vrows_t1_dev_forced())
11519            // BISECT (MEMRA_GLM5_GRAPH_HOST_MOE, gate harness): stand the device-table arm down
11520            // while leaving the door on, so a box run separates the door's two enablers. The
11521            // capture refuses by name when this is set — a host readback cannot live inside a
11522            // capture region — which is the point: it isolates one enabler at a time.
11523            && !crate::glm5_graph_host_moe()
11524            && !promote_worker_h2d
11525            && sigmoid_router_enabled()
11526            && cfg.sigmoid_router().is_some()
11527            && !observe_routes
11528            && !memra_reference::hidden_trace::enabled()
11529            && !crate::moesd::capture_active();
11530        // DOOR D (`MEMRA_MOE_VROWS_DEV_TABLES`, default OFF): route WITHOUT the pinned sel/w
11531        // readback and build the pair's pointer/scale tables on device instead. On the serving
11532        // shape the host table build is the selection's ONLY consumer, and it costs a full
11533        // `cuStreamSynchronize` + 2 DtoH + 2 pageable HtoD + 2 host Vecs per MoE layer-call —
11534        // 42 device-wide drains, 84 DtoH and 84 HtoD per ship round (the decode-gap
11535        // attribution's "43 cuStreamSynchronize/token ... the per-layer router-admission sync
11536        // structure", and 44.6% of the unattributed 71.6 HtoD calls/token).
11537        //
11538        // The extra conjuncts beyond `vrows_fires` (asserted equal at the dispatch) are exactly
11539        // the host-visible consumers of `sel_all` between here and there, each of which would
11540        // silently read an empty selection: `moesd::record_host_routes`, `hidden_trace`,
11541        // `MEMRA_MOE_TRACE`/`MEMRA_MOE_STATS`/the other `observe_routes` modes. Plus
11542        // `sigmoid_router_enabled()`, because `MEMRA_SIG_ROUTER=0` is a full-logit HOST oracle
11543        // with no device selection to read. `promote_worker_h2d` USED to need no conjunct here
11544        // (it requires t == 1 and door D required t >= 2); the T=1 decode-graph arm above lifts
11545        // that exclusion and carries the conjunct itself. Any miss falls closed to the host
11546        // readback.
11547        let vrows_dev = ((vrows && t >= 2 && crate::moe_vrows_dev_tables_on()) || vrows_t1_dev)
11548            && slab_bases.is_some()
11549            && moe_q8
11550            && uniform_experts
11551            && n_used <= 8
11552            && cfg.sigmoid_router().is_some()
11553            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
11554            && !cpu_hybrid
11555            && sigmoid_router_enabled()
11556            && !observe_routes
11557            && !memra_reference::hidden_trace::enabled()
11558            && !crate::moesd::capture_active();
11559        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
11560        // Door D adds a fourth arm returning the router's DEVICE sel/w with no readback.
11561        let mut sel_dev: Option<(CudaSlice<i32>, CudaSlice<f32>)> = None;
11562        let (sel_all, w_all, routed_cpu_input) = if vrows_dev {
11563            let (sf, route_norm) = cfg
11564                .sigmoid_router()
11565                .expect("vrows_dev carries cfg.sigmoid_router().is_some()");
11566            sel_dev = Some(e.moe_router_sigmoid_topk(
11567                &logits,
11568                t,
11569                n_expert,
11570                n_used,
11571                m.active_count(),
11572                &m.exp_probs_b_dev,
11573                &m.active_experts_dev,
11574                sf,
11575                route_norm,
11576            )?);
11577            crate::MOE_VROWS_ROUTER_SYNCS_AVOIDED
11578                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11579            // GATE INSTRUMENT (MEMRA_GLM5_GRAPH_SEL_LEDGER, never a serving flag): capture-legal
11580            // D2D of this layer's device selection into a pre-armed persistent slot, so the
11581            // decode-graph gate can assert the device arm picks the same experts and weights the
11582            // host oracle reads back. No sync, no DtoH, no allocation — see glm5_sel_ledger.rs.
11583            if let Some((si, sw)) = sel_dev.as_ref() {
11584                // MEMRA_MOE_SEL_DUMP: the device-table arm's own rows (one DtoH pair per
11585                // layer-call, diagnostic only; unset = one OnceLock read). The host twin
11586                // below hands `trace_moe_routes` an empty selection, so this is the only
11587                // record this arm makes.
11588                //
11589                // COMPOSITION with the decode-graph door: that DtoH pair cannot run inside a
11590                // capture region, so `glm5_decode_graph_refusal` refuses the door BY NAME while
11591                // this dump is armed. Without that refusal the pair would be RECORDED and not
11592                // executed, the dump would silently write stale rows, and the door would look
11593                // fine (2026-09-03, when the two lanes met in a rebase).
11594                crate::moe_sel_dump::record_device(e, il, t, n_used, si, sw)?;
11595                // PRE-ARM OUTSIDE THE CAPTURE, HERE. `record_device` skips a layer that has no
11596                // persistent slot (it may never allocate inside a capture region), and the only
11597                // other `prearm` caller is `glm5_capture_stage` — so any cell that runs the
11598                // device arm WITHOUT capturing recorded zero device rows and the gate reported
11599                // `VACUOUS: the selection ledger recorded no rows on one of the arms`. That is
11600                // the instrument failing, not the door: box take 12's NO_CAPTURE run died on it.
11601                // This call is skipped while a capture is open (allocation is illegal there) and
11602                // is a no-op once the slot exists, so the captured path is unchanged.
11603                if !crate::glm5_graph_capture_open() {
11604                    crate::glm5_sel_ledger::prearm(e, il, n_used)?;
11605                }
11606                crate::glm5_sel_ledger::record_device(e, il, si, sw)?;
11607            }
11608            // MEMRA_GLM5_GRAPH_TRACE: dump this arm's inputs for the first two routed layers so
11609            // the DEVICE arm and the HOST arm can be diffed line for line on the box. Only
11610            // outside a capture region — the readback below is illegal inside one.
11611            if crate::glm5_graph_trace_on() && !crate::glm5_graph_capture_open() {
11612                let (sel_h, w_h) = match sel_dev.as_ref() {
11613                    Some((si, sw)) => (
11614                        e.dtoh_i32(si)?
11615                            .iter()
11616                            .map(|&x| x as u32)
11617                            .collect::<Vec<_>>(),
11618                        e.dtoh(sw)?,
11619                    ),
11620                    None => (Vec::new(), Vec::new()),
11621                };
11622                Self::dump_moe_t1_inputs(
11623                    e, "device", m, cfg, il, t, n_used, n_expert, &sel_h, &w_h,
11624                );
11625            }
11626            (Vec::new(), Vec::new(), None)
11627        } else if let Some(sig) = cfg.sigmoid_router() {
11628            if cpu_hybrid {
11629                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
11630                    e,
11631                    &logits,
11632                    z,
11633                    t,
11634                    n_embd,
11635                    n_expert,
11636                    n_used,
11637                    m.exp_probs_b.as_deref(),
11638                    sig,
11639                    m.active_experts.as_deref(),
11640                )?;
11641                (sel, w, Some(input))
11642            } else {
11643                let (sel, w) =
11644                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
11645                (sel, w, None)
11646            }
11647        } else {
11648            let (sel, w) =
11649                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
11650            (sel, w, None)
11651        };
11652        // The HOST arm's twin of the dump above (same trace flag, same first two routed layers).
11653        // Run 7B could not be diffed against run A because only the device arm printed; with both
11654        // arms printing, the first differing field IS the answer.
11655        if crate::glm5_graph_trace_on() && !crate::glm5_graph_capture_open() && !sel_all.is_empty()
11656        {
11657            let last = (t - 1) * n_used;
11658            Self::dump_moe_t1_inputs(
11659                e,
11660                "host",
11661                m,
11662                cfg,
11663                il,
11664                t,
11665                n_used,
11666                n_expert,
11667                &sel_all[last..last + n_used],
11668                &w_all[last..last + n_used],
11669            );
11670        }
11671        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
11672        // The host-oracle twin of the ledger record above (gate instrument only).
11673        if !sel_all.is_empty() && crate::glm5_sel_ledger::armed() {
11674            let last = (t - 1) * n_used;
11675            crate::glm5_sel_ledger::record_host(
11676                e.ctx().ordinal(),
11677                il,
11678                &sel_all[last..last + n_used],
11679                &w_all[last..last + n_used],
11680            );
11681        }
11682        if memra_reference::hidden_trace::enabled() {
11683            memra_reference::hidden_trace::emit_last_row(
11684                "router",
11685                il as i64,
11686                t,
11687                n_expert,
11688                &e.dtoh(&logits)?,
11689            );
11690            let last = (t - 1) * n_used;
11691            let mut route = Vec::with_capacity(n_used * 2);
11692            for slot in 0..n_used {
11693                route.push(sel_all[last + slot] as f32);
11694                route.push(w_all[last + slot]);
11695            }
11696            memra_reference::hidden_trace::emit_last_row("route", il as i64, 1, n_used * 2, &route);
11697        }
11698
11699        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
11700        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
11701        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
11702        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
11703        Self::trace_moe_input(e, il, t, n_embd, z)?;
11704
11705        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
11706        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
11707        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
11708        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
11709        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
11710        // wait for each pending block, so later copies can overlap the earlier expert kernels while
11711        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
11712        // T=1; batched forwards can have token-local consumers still in flight between selections.
11713        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
11714        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
11715        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
11716        if promote_worker_h2d {
11717            let mut selected_blocks = Vec::with_capacity(n_used * 3);
11718            for &ex in sel_all.iter().take(n_used) {
11719                let ex = ex as u16;
11720                selected_blocks.extend([
11721                    BlockId::new(il, PROJ_GATE, ex),
11722                    BlockId::new(il, PROJ_UP, ex),
11723                    BlockId::new(il, PROJ_DOWN, ex),
11724                ]);
11725            }
11726            for &ex in sel_all.iter().take(n_used) {
11727                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
11728            }
11729            e.with_moe_cache(max_block, |cache, eng| {
11730                cache.promote_worker_reads_at_safe_boundary(
11731                    &selected_blocks,
11732                    &selected_blocks,
11733                    eng,
11734                )?;
11735                Ok(())
11736            })?;
11737        }
11738
11739        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
11740        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
11741        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
11742            let mut cnt = vec![0u32; n_expert];
11743            for &s in sel_all.iter() {
11744                cnt[s as usize] += 1;
11745            }
11746            let total = sel_all.len() as f64;
11747            let mut h = 0.0f64;
11748            let mut active = 0usize;
11749            for &c in &cnt {
11750                if c > 0 {
11751                    active += 1;
11752                    let p = c as f64 / total;
11753                    h -= p * p.log2();
11754                }
11755            }
11756            let maxc = cnt.iter().copied().max().unwrap_or(0);
11757            println!(
11758                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
11759                il,
11760                t,
11761                sel_all.len(),
11762                active,
11763                n_expert,
11764                h,
11765                (n_expert as f64).log2(),
11766                total / active.max(1) as f64,
11767                maxc
11768            );
11769        }
11770
11771        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
11772        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
11773        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
11774        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
11775        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
11776        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
11777        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
11778        // zeroed-then-accumulated exactly as before (fallback).
11779        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
11780        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
11781        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
11782        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
11783        let gdec_may_fire = uniform_experts
11784            && use_cache
11785            && n_used <= 8
11786            && gdec_enabled()
11787            && !cfg.swiglu_clamped_at(il as u32);
11788        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
11789        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
11790        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
11791        // archs the slabs were uploaded but never read, and every expert went through the
11792        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
11793        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
11794        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
11795        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
11796        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
11797        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
11798        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
11799        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
11800        // strictly worse than staging); under PP-2 without the prime walker this admits
11801        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
11802        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
11803        // `slab_local` / `slab_bases` are bound ABOVE the router: door D
11804        // (`MEMRA_MOE_VROWS_DEV_TABLES`) has to pre-decide the routing arm, and this is the
11805        // predicate it needs. Pure immutable pointer reads, so the hoist is behaviour-neutral.
11806        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
11807        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
11808        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
11809        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
11810        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
11811        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
11812        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
11813        // all-resident tokens, staged loop for misses), which is a dispatch-class
11814        // comparison, not a provenance one.
11815        let slab_fused_may_fire = slab_bases.is_some()
11816            && n_used <= 8
11817            && gdec_enabled()
11818            && !cfg.swiglu_clamped_at(il as u32)
11819            && cfg.m3.is_none()
11820            && no_exp_macros
11821            && moe_q8;
11822        // FUSED MoE EPILOGUE (lane/glm53-epilogue 2026-08-28, MEMRA_MOE_FUSED_EPI default OFF).
11823        // The arm glm5_next is denied by every other predicate in this function. It runs the SAME
11824        // launch pair shape as gdec — one gate/up kernel, one down/FMA kernel per token — but
11825        // with the three things this arch actually needs, none of which any existing fused
11826        // epilogue has:
11827        //   * the SIGMOID noaux_tc router's sel/w (host-routed above; the fused softmax router
11828        //     `moe_router_topk` that pairs/dev use would pick different experts — the M3
11829        //     gate-MISMATCH 74602-vs-92 lesson);
11830        //   * the PRE-clamped SwiGLU epilogue `silu(min(g,l)) * clamp(u,±l)` — step35's POST form
11831        //     is a different, plausible-but-wrong program (`fused_post_limit`);
11832        //   * the per-expert NVFP4 `weight_scale_2` macro fold, gate/up through the kernel's
11833        //     gs/us and down through the routing weight, exactly as `ffn_act_lim` + `axpy_into`
11834        //     do it in the sequential loop.
11835        // UNLIKE gdec it does NOT require the layer to be already-resident: it ADMITS the
11836        // 3*n_used selected blocks through the same `dispatch_source` the sequential loop uses
11837        // (hit = no copy, miss = the identical H2D into a slot) and only then collects the fixed
11838        // slot addresses. The staged bytes are unchanged — the §B.3 provenance property — so the
11839        // arm engages at any miss rate instead of gdec's P(all resident). It needs the cache to
11840        // hold 3*n_used blocks at once; `moe_fused_epi_token_q8` returns false when it cannot and
11841        // the token falls through to the sequential loop below.
11842        // TWO PROVENANCES, ONE LAUNCH PATH (slab arm added 2026-08-28). The SLRU arm below keys
11843        // on `slab_local.is_none()`; the SLAB arm keys on the slab existing. They differ ONLY in
11844        // where the eight expert pointers come from and both call `moe_fused_epi_launch`, so the
11845        // macro fold, the clamp and the kernel pair cannot drift apart between them.
11846        //
11847        // The slab arm is not an optimization, it is the arm that matters. Full two-card expert
11848        // residency makes `dev_exps` present on every stage engine, which makes `slab_local`
11849        // `Some`, which under the original predicate DENIED the fused epilogue outright — the
11850        // measured A/B would have read 0 dispatches and looked like "no effect". The residency
11851        // config is the serving config now, so the slab provenance is the one the product runs.
11852        //
11853        // It is also the SIMPLER arm: a slab holds every expert by construction, so there is no
11854        // admission, no eviction, no pass-2 re-verification and no fall-through. The SLRU arm's
11855        // capacity floor and re-check exist only because admission can move a slot.
11856        let fused_epi_common = n_used <= 8
11857            && moe_q8
11858            && cfg.m3.is_none()
11859            && cfg.sigmoid_router().is_some()
11860            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
11861            && moe_fused_epi_enabled();
11862        let fused_epi_may_fire = fused_epi_common
11863            && uniform_experts
11864            && use_cache
11865            && cache_dispatch
11866            && slab_local.is_none();
11867        let fused_epi_slab_may_fire = fused_epi_common && slab_bases.is_some();
11868        // VERIFY-ROWS BATCHED ROUTED-EXPERT ARM (lane/glm5-vrest, 2026-08-31; rides
11869        // `MEMRA_GLM5_VERIFY_BATCH` — only the verify walk's batched arm passes `vrows`).
11870        // ONE launch pair covers ALL t x n_used routed pairs (the fused-epilogue kernels'
11871        // verify-rows twins) instead of the per-(token,expert) loop's ~49 launches per
11872        // token-layer — the flip-reprice cell-2 vrest wall (9.46 ms/row marginal at K=3).
11873        // Bit identity per row vs the sequential chain is the bar and it is structural:
11874        // routing is the SAME host invocation above; per-pair dots are qmatvec_expert_q8's
11875        // g-strided order; the epilogue is swiglu_preclamped_mul_scaled_f32's expression
11876        // with the per-expert macro fold exactly where ffn_act_lim/axpy_into fold it; the
11877        // down accumulation is the slot-ordered __fmaf_rn chain (the gdec-gated class).
11878        // Gated by glm5_verify_batch_gpu (kernel pair vs sequential chain + swapped-pair
11879        // and dropped-macro reds) and the glm5_tparallel_verify_gpu walk battery on the
11880        // NVFP4+macro serving expert class. Fail-closed: any unqualified shape falls
11881        // through to the unchanged loop below. Same slab-only scope as the fused epilogue
11882        // (the serving config's provenance); n_used<=8 mirrors its cap.
11883        let vrows_fires = ((vrows && t >= 2) || vrows_t1_dev)
11884            && slab_bases.is_some()
11885            && moe_q8
11886            && uniform_experts
11887            && n_used <= 8
11888            && cfg.sigmoid_router().is_some()
11889            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
11890            && !cpu_hybrid;
11891        // WHY THE ARM DID NOT FIRE, named rather than inferred (2026-09-03). Twelve box takes
11892        // could not tell "the device-table arm ran and was wrong" from "the arm never fired and
11893        // the layer took the host-readback loop", because the only evidence was a `arm=host`
11894        // label that BOTH the sequential loop and the verify-rows host arm can print. One line
11895        // per denying layer, under the door or the trace flag only, settles it on the next run.
11896        if t == 1
11897            && !vrows_fires
11898            && (crate::glm5_decode_graph_on() || crate::glm5_graph_trace_on())
11899            && crate::glm5_trace_take_slot("deny", "t1", il)
11900        {
11901            eprintln!(
11902                "[glm5-vrows-t1-deny] dev={} il={il} capture_open={} t1_dev={vrows_t1_dev} \
11903                 forced={} slab_bases={} moe_q8={moe_q8} uniform={uniform_experts} \
11904                 n_used={n_used} sigmoid_cfg={} pre_clamp={} cpu_hybrid={cpu_hybrid} \
11905                 promote_worker_h2d={promote_worker_h2d} observe_routes={observe_routes} \
11906                 sig_router_env={} — the T=1 device-table MoE arm stood down and this layer \
11907                 takes the host-readback path. Benign with capture_open=false (an eager gap \
11908                 layer); with capture_open=true it is the wrong-answer class the router guard \
11909                 refuses, and the first denied conjunct on the line is the reason.",
11910                e.ctx().ordinal(),
11911                crate::glm5_graph_capture_open(),
11912                crate::glm5_vrows_t1_dev_forced(),
11913                slab_bases.is_some(),
11914                cfg.sigmoid_router().is_some(),
11915                matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6),
11916                sigmoid_router_enabled(),
11917            );
11918        }
11919        // moe_out memset elision: EVERY full-row-overwrite arm (gdec, slab fused, fused epilogue,
11920        // verify-rows) allocates uninit; a token that falls through to any accumulating loop
11921        // zeroes its own row. The fused epilogue's `moe_down8_fma_q8` fully overwrites `dst[o]`,
11922        // same as gdec's; `moe_down8_fma_q8_rows` fully overwrites every row.
11923        let mut moe_out = if gdec_may_fire
11924            || slab_fused_may_fire
11925            || fused_epi_may_fire
11926            || fused_epi_slab_may_fire
11927            || vrows_fires
11928        {
11929            e.uninit(t * n_embd)?
11930        } else {
11931            e.zeros(t * n_embd)?
11932        };
11933        // The router readback above already established a host boundary. Copy each small-t hidden
11934        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
11935        let cpu_input = if cpu_hybrid {
11936            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
11937        } else {
11938            None
11939        };
11940
11941        // Door D's predicate was evaluated at the router, BEFORE `vrows_fires` existed. If the
11942        // two ever disagreed, the layer would have skipped its readback and then dispatched a
11943        // per-row arm holding an EMPTY host selection — a silent wrong-answer class. The two
11944        // predicates share every conjunct by construction; assert it rather than hope.
11945        if vrows_dev && !vrows_fires {
11946            return Err(
11947                "MEMRA_MOE_VROWS_DEV_TABLES routed device-only but the verify-rows arm did not \
11948                 fire: the door-D and vrows_fires predicates disagree"
11949                    .into(),
11950            );
11951        }
11952        if vrows_fires {
11953            let Some(SwigluClamp::Pre(limit)) = lim_exp else {
11954                return Err(
11955                    "verify-rows MoE arm fired without a live PRE clamp: the predicate and \
11956                     the dispatch disagree"
11957                        .into(),
11958                );
11959            };
11960            let bases = slab_bases.expect("vrows_fires carries slab_bases.is_some()");
11961            let sel = match sel_dev.as_ref() {
11962                Some((si, sw)) => VrowsSel::Dev(si, sw),
11963                None => VrowsSel::Host(&sel_all, &w_all),
11964            };
11965            Self::moe_vrows_pairs_q8(
11966                e,
11967                m,
11968                z,
11969                sel,
11970                il,
11971                bases,
11972                slab_rp,
11973                t,
11974                n_embd,
11975                n_ff_exp,
11976                n_used,
11977                limit,
11978                &mut moe_out,
11979            )?;
11980            if memra_reference::hidden_trace::enabled() {
11981                memra_reference::hidden_trace::emit_last_row(
11982                    "routed",
11983                    il as i64,
11984                    t,
11985                    n_embd,
11986                    &e.dtoh(&moe_out)?,
11987                );
11988            }
11989            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11990            return Ok(moe_out);
11991        }
11992
11993        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
11994        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
11995        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
11996        // measured ~123 memsets/token of the decode wall).
11997        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
11998        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
11999        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
12000        let mut scratch_g: Option<CudaSlice<u8>> = None;
12001        let mut scratch_u: Option<CudaSlice<u8>> = None;
12002        let mut scratch_d: Option<CudaSlice<u8>> = None;
12003        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
12004        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
12005
12006        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
12007        // the copy stream before launching the current expert's compute. Pending slots stay invisible
12008        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
12009        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
12010        let page_window = moe_page_prefetch_window();
12011
12012        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
12013        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
12014        for tok in 0..t {
12015            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
12016            let w = &w_all[tok * n_used..(tok + 1) * n_used];
12017            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
12018            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12019
12020            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
12021            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
12022            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
12023            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
12024            // memcpy, zero admission, so no slot can move under the collected pointers) — any
12025            // miss falls through to the sequential loop below, which admits as before. In steady
12026            // state on a fully-resident rig every token-layer takes the grouped path.
12027            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
12028            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
12029            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
12030            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
12031            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
12032            // per-expert macro-scales the fused kernels don't fold — those fall through too.
12033            let no_macros = m.gate_exps.macros.is_none()
12034                && m.up_exps.macros.is_none()
12035                && m.down_exps.macros.is_none();
12036            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
12037            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
12038            // with pointers computed from the resident slab base + ex*stride instead of
12039            // collected SLRU slot addresses. No cache lock, no residency predicate — the
12040            // slab holds every expert by construction, so this arm never falls through
12041            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
12042            // staging both die). Bit-identity class: pointer provenance only, the same
12043            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
12044            // slab exists it is strictly better (no lock, no miss).
12045            if slab_fused_may_fire {
12046                let (pg, pu, pd) = slab_bases.unwrap();
12047                let mut gp = [0u64; 8];
12048                let mut up = [0u64; 8];
12049                let mut dp = [0u64; 8];
12050                for (j, &ex) in sel.iter().enumerate() {
12051                    let ex = ex as usize;
12052                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
12053                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
12054                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
12055                }
12056                let mut wv = [0f32; 8];
12057                wv[..n_used].copy_from_slice(w);
12058                if tok_q8.is_none() {
12059                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12060                }
12061                let (zq, zd) = tok_q8.as_ref().unwrap();
12062                Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12063                let act = e.moe_gate_up_silu8_q8(
12064                    crate::WPtr8(gp),
12065                    crate::WPtr8(up),
12066                    zq,
12067                    zd,
12068                    n_embd,
12069                    n_ff_exp,
12070                    n_used,
12071                    m.gate_exps.qtype,
12072                    m.up_exps.qtype,
12073                    m.gate_exps.row_bytes,
12074                    m.up_exps.row_bytes,
12075                )?;
12076                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
12077                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12078                e.moe_down8_fma_q8(
12079                    crate::WPtr8(dp),
12080                    crate::F32x8(wv),
12081                    &aq2,
12082                    &ad2,
12083                    &mut dst,
12084                    n_ff_exp,
12085                    n_embd,
12086                    n_used,
12087                    m.down_exps.qtype,
12088                    m.down_exps.row_bytes,
12089                )?;
12090                continue;
12091            }
12092            // FUSED MoE EPILOGUE, SLAB PROVENANCE. Ordered first: when a local slab exists it is
12093            // strictly better than anything the SLRU can offer — every expert is present by
12094            // construction, so there is no admission, no eviction and no fall-through. This is
12095            // the arm the two-card residency serving config actually runs.
12096            if fused_epi_slab_may_fire {
12097                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
12098                    return Err(
12099                        "fused MoE epilogue (slab) fired without a live PRE clamp: the \
12100                                predicate and the dispatch disagree"
12101                            .into(),
12102                    );
12103                };
12104                let (pg, pu, pd) = slab_bases.unwrap();
12105                let mut g = [0u64; 8];
12106                let mut u = [0u64; 8];
12107                let mut d = [0u64; 8];
12108                for (j, &ex) in sel.iter().enumerate() {
12109                    let ex = ex as usize;
12110                    g[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
12111                    u[j] = pu + (ex * m.up_exps.expert_stride) as u64;
12112                    d[j] = pd + (ex * m.down_exps.expert_stride) as u64;
12113                }
12114                if tok_q8.is_none() {
12115                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12116                }
12117                let (zq, zd) = tok_q8.as_ref().unwrap();
12118                Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12119                Self::moe_fused_epi_launch(
12120                    e,
12121                    m,
12122                    zq,
12123                    zd,
12124                    sel,
12125                    w,
12126                    g,
12127                    u,
12128                    d,
12129                    &mut moe_out,
12130                    tok,
12131                    n_embd,
12132                    n_ff_exp,
12133                    n_used,
12134                    limit,
12135                    slab_rp,
12136                )?;
12137                continue;
12138            }
12139            // FUSED MoE EPILOGUE, SLRU PROVENANCE. Ordered before gdec (which this arch never
12140            // reaches anyway: `gdec_may_fire` carries `!swiglu_clamped_at`). A `false` return
12141            // means the cache could not hold 3*n_used blocks at once — the token falls through to
12142            // the sequential loop, which zeroes its own row below.
12143            if fused_epi_may_fire {
12144                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
12145                    return Err(
12146                        "fused MoE epilogue fired without a live PRE clamp: the predicate and \
12147                         the dispatch disagree"
12148                            .into(),
12149                    );
12150                };
12151                if tok_q8.is_none() {
12152                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12153                }
12154                let (zq, zd) = tok_q8.as_ref().unwrap();
12155                Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12156                if Self::moe_fused_epi_token_q8(
12157                    e,
12158                    m,
12159                    il,
12160                    max_block,
12161                    zq,
12162                    zd,
12163                    sel,
12164                    w,
12165                    &mut moe_out,
12166                    tok,
12167                    n_embd,
12168                    n_ff_exp,
12169                    n_used,
12170                    limit,
12171                )? {
12172                    continue;
12173                }
12174            }
12175            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
12176                if tok_q8.is_none() {
12177                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12178                }
12179                let (zq, zd) = tok_q8.as_ref().unwrap();
12180                Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12181                if Self::moe_gdec_token_q8(
12182                    e,
12183                    m,
12184                    il,
12185                    max_block,
12186                    zq,
12187                    zd,
12188                    sel,
12189                    w,
12190                    &mut moe_out,
12191                    tok,
12192                    n_embd,
12193                    n_ff_exp,
12194                    n_used,
12195                )? {
12196                    continue;
12197                }
12198            } else if gdec_may_fire
12199                && cfg.m3.is_none()
12200                && no_macros
12201                && Self::moe_gdec_token(
12202                    e,
12203                    m,
12204                    il,
12205                    max_block,
12206                    &zt,
12207                    sel,
12208                    w,
12209                    &mut moe_out,
12210                    tok,
12211                    n_embd,
12212                    n_ff_exp,
12213                    n_used,
12214                )?
12215            {
12216                continue;
12217            }
12218
12219            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
12220            // slab pair could fire. This token fell through to a sequential axpy loop, which
12221            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
12222            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
12223            // has no fallible predicate), included for the allocation invariant's symmetry.
12224            if gdec_may_fire || slab_fused_may_fire || fused_epi_may_fire || fused_epi_slab_may_fire
12225            {
12226                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12227                e.memset_zeros_view(&mut row)?;
12228            }
12229
12230            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
12231            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
12232            // stall this path exists to remove, while mixing projections would require another
12233            // activation round-trip. Weight addresses remain valid until this worker is joined at
12234            // the bottom of the token scope.
12235            let mut cpu_mask = vec![false; sel.len()];
12236            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
12237                let gpu_resident = if use_cache {
12238                    e.with_moe_cache(max_block, |cache, _| {
12239                        Ok(sel
12240                            .iter()
12241                            .map(|&expert| {
12242                                let expert = expert as u16;
12243                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
12244                                    .into_iter()
12245                                    .filter(|&projection| {
12246                                        cache
12247                                            .resident(BlockId::new(il, projection, expert))
12248                                            .is_some()
12249                                    })
12250                                    .count()
12251                            })
12252                            .collect::<Vec<_>>())
12253                    })?
12254                } else {
12255                    vec![0; sel.len()]
12256                };
12257                let mut cpu_selected = Vec::new();
12258                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
12259                    if gpu_resident[index] != 3 {
12260                        cpu_mask[index] = true;
12261                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
12262                        let expert = expert as usize;
12263                        cpu_selected.push((expert, route_weight));
12264                    }
12265                }
12266                if crate::cpu_experts::predictor_enabled() {
12267                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
12268                    // from this layer's MoE input and prefetches predicted-and-missing
12269                    // experts into the companion RAM cache. Never blocks this thread.
12270                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
12271                    crate::cpu_experts::predictor_submit(il, row);
12272                }
12273                if cpu_selected.is_empty() {
12274                    None
12275                } else {
12276                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
12277                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
12278                        .map_err(std::io::Error::other)?;
12279                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
12280                }
12281            } else {
12282                None
12283            };
12284
12285            let worker_window = worker_disk_prefetch
12286                .then(worker_prefetch_window)
12287                .unwrap_or(0);
12288            for (j, &ex) in sel.iter().enumerate() {
12289                if cpu_mask[j] {
12290                    continue;
12291                }
12292                let ex = ex as usize;
12293                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
12294                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
12295                // fused form) and macro-carrying artifacts — still have their bytes in the
12296                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
12297                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
12298                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
12299                if let Some(d) = slab_local {
12300                    let gl = m.gate_exps.expert_layout(ex);
12301                    let ul = m.up_exps.expert_layout(ex);
12302                    let dl = m.down_exps.expert_layout(ex);
12303                    let (g0, u0, d0) = (
12304                        ex * m.gate_exps.expert_stride,
12305                        ex * m.up_exps.expert_stride,
12306                        ex * m.down_exps.expert_stride,
12307                    );
12308                    let (gate, up) = if moe_q8 {
12309                        if tok_q8.is_none() {
12310                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12311                        }
12312                        let (zq, zd) = tok_q8.as_ref().unwrap();
12313                        Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12314                        (
12315                            e.qmatvec_expert_q8(
12316                                &d.gate,
12317                                g0..g0 + gl.len,
12318                                zq,
12319                                zd,
12320                                1,
12321                                m.gate_exps.in_f,
12322                                m.gate_exps.out_f,
12323                                gl.qtype,
12324                                gl.row_bytes,
12325                            )?,
12326                            e.qmatvec_expert_q8(
12327                                &d.up,
12328                                u0..u0 + ul.len,
12329                                zq,
12330                                zd,
12331                                1,
12332                                m.up_exps.in_f,
12333                                m.up_exps.out_f,
12334                                ul.qtype,
12335                                ul.row_bytes,
12336                            )?,
12337                        )
12338                    } else {
12339                        (
12340                            m.qmatvec_view(
12341                                e,
12342                                &d.gate,
12343                                g0..g0 + gl.len,
12344                                &zt,
12345                                1,
12346                                m.gate_exps.in_f,
12347                                m.gate_exps.out_f,
12348                                gl.qtype,
12349                                gl.row_bytes,
12350                            )?,
12351                            m.qmatvec_view(
12352                                e,
12353                                &d.up,
12354                                u0..u0 + ul.len,
12355                                &zt,
12356                                1,
12357                                m.up_exps.in_f,
12358                                m.up_exps.out_f,
12359                                ul.qtype,
12360                                ul.row_bytes,
12361                            )?,
12362                        )
12363                    };
12364                    let mut act = e.uninit(n_ff_exp)?;
12365                    Self::ffn_act_lim(
12366                        e,
12367                        cfg,
12368                        &gate,
12369                        &up,
12370                        m.gate_exps.macro_scale(ex),
12371                        m.up_exps.macro_scale(ex),
12372                        lim_exp,
12373                        &mut act,
12374                        n_ff_exp,
12375                    )?;
12376                    let y = if moe_q8 {
12377                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
12378                        e.qmatvec_expert_q8(
12379                            &d.down,
12380                            d0..d0 + dl.len,
12381                            &aq2,
12382                            &ad2,
12383                            1,
12384                            m.down_exps.in_f,
12385                            m.down_exps.out_f,
12386                            dl.qtype,
12387                            dl.row_bytes,
12388                        )?
12389                    } else {
12390                        let actv = act.slice(0..n_ff_exp);
12391                        m.qmatvec_view(
12392                            e,
12393                            &d.down,
12394                            d0..d0 + dl.len,
12395                            &actv,
12396                            1,
12397                            m.down_exps.in_f,
12398                            m.down_exps.out_f,
12399                            dl.qtype,
12400                            dl.row_bytes,
12401                        )?
12402                    };
12403                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12404                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
12405                    continue;
12406                }
12407                for next in page_prefetch_positions(j, sel.len(), page_window) {
12408                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
12409                }
12410                let keep = [
12411                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
12412                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
12413                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
12414                ];
12415                if worker_disk_prefetch && worker_window > 0 {
12416                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
12417                        Self::moe_prefetch_disk_expert(
12418                            e,
12419                            il,
12420                            sel[next] as usize,
12421                            m,
12422                            max_block,
12423                            &keep,
12424                        )?;
12425                    }
12426                } else if cache_dispatch
12427                    && !cpu_hybrid
12428                    && moe_prefetch_enabled()
12429                    && j + 1 < sel.len()
12430                {
12431                    let next = sel[j + 1] as usize;
12432                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
12433                }
12434                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
12435                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
12436                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
12437                    // layouts stay on the metadata-aware f32 path.
12438                    if (gate_q8 || up_q8) && tok_q8.is_none() {
12439                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
12440                    }
12441                    let gate = if gate_q8 {
12442                        let (zq, zd) = tok_q8.as_ref().unwrap();
12443                        Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12444                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
12445                    } else {
12446                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
12447                    };
12448                    let up = if up_q8 {
12449                        let (zq, zd) = tok_q8.as_ref().unwrap();
12450                        Self::trace_moe_act(e, "host", il, t, z, zq, zd);
12451                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
12452                    } else {
12453                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
12454                    };
12455                    let mut act = e.uninit(n_ff_exp)?;
12456                    Self::ffn_act_lim(
12457                        e,
12458                        cfg,
12459                        &gate,
12460                        &up,
12461                        m.gate_exps.macro_scale(ex),
12462                        m.up_exps.macro_scale(ex),
12463                        lim_exp,
12464                        &mut act,
12465                        n_ff_exp,
12466                    )?;
12467                    let y = if down_q8 {
12468                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
12469                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
12470                    } else {
12471                        let actv = act.slice(0..n_ff_exp);
12472                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
12473                    };
12474                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12475                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
12476                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
12477                } else if cache_dispatch {
12478                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
12479                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
12480                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
12481                    // only difference between HIT and MISS is whether the memcpy_htod ran.
12482                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
12483                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
12484                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
12485                    Self::ffn_act_lim(
12486                        e,
12487                        cfg,
12488                        &gate,
12489                        &up,
12490                        m.gate_exps.macro_scale(ex),
12491                        m.up_exps.macro_scale(ex),
12492                        lim_exp,
12493                        &mut act,
12494                        n_ff_exp,
12495                    )?;
12496                    let actv = act.slice(0..n_ff_exp);
12497                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
12498                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12499                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
12500                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
12501                } else if cache_frozen {
12502                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
12503                    // first prime. Reuse every fixed resident projection directly and stage only a
12504                    // true miss through the ordinary scratch slot. This preserves the established
12505                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
12506                    let gate = Self::moe_frozen_gemm(
12507                        e,
12508                        il,
12509                        PROJ_GATE,
12510                        ex,
12511                        m,
12512                        max_block,
12513                        &zt,
12514                        &mut scratch_g,
12515                        g_len,
12516                    )?;
12517                    let up = Self::moe_frozen_gemm(
12518                        e,
12519                        il,
12520                        PROJ_UP,
12521                        ex,
12522                        m,
12523                        max_block,
12524                        &zt,
12525                        &mut scratch_u,
12526                        u_len,
12527                    )?;
12528                    let mut act = e.uninit(n_ff_exp)?;
12529                    Self::ffn_act_lim(
12530                        e,
12531                        cfg,
12532                        &gate,
12533                        &up,
12534                        m.gate_exps.macro_scale(ex),
12535                        m.up_exps.macro_scale(ex),
12536                        lim_exp,
12537                        &mut act,
12538                        n_ff_exp,
12539                    )?;
12540                    let actv = act.slice(0..n_ff_exp);
12541                    let y = Self::moe_frozen_gemm(
12542                        e,
12543                        il,
12544                        PROJ_DOWN,
12545                        ex,
12546                        m,
12547                        max_block,
12548                        &actv,
12549                        &mut scratch_d,
12550                        d_len,
12551                    )?;
12552                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12553                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
12554                } else {
12555                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
12556                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
12557                    // fully overwrites the byte range the GEMM reads).
12558                    if scratch_g.is_none() {
12559                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
12560                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
12561                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
12562                    }
12563                    let (sg, su, sd) = (
12564                        scratch_g.as_mut().unwrap(),
12565                        scratch_u.as_mut().unwrap(),
12566                        scratch_d.as_mut().unwrap(),
12567                    );
12568                    let gl = m.gate_exps.expert_layout(ex);
12569                    let ul = m.up_exps.expert_layout(ex);
12570                    let dl = m.down_exps.expert_layout(ex);
12571                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
12572                    let gate = m.qmatvec_view(
12573                        e,
12574                        sg,
12575                        0..gl.len,
12576                        &zt,
12577                        1,
12578                        m.gate_exps.in_f,
12579                        m.gate_exps.out_f,
12580                        gl.qtype,
12581                        gl.row_bytes,
12582                    )?;
12583
12584                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
12585                    let up = m.qmatvec_view(
12586                        e,
12587                        su,
12588                        0..ul.len,
12589                        &zt,
12590                        1,
12591                        m.up_exps.in_f,
12592                        m.up_exps.out_f,
12593                        ul.qtype,
12594                        ul.row_bytes,
12595                    )?;
12596
12597                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
12598                    Self::ffn_act_lim(
12599                        e,
12600                        cfg,
12601                        &gate,
12602                        &up,
12603                        m.gate_exps.macro_scale(ex),
12604                        m.up_exps.macro_scale(ex),
12605                        lim_exp,
12606                        &mut act,
12607                        n_ff_exp,
12608                    )?;
12609
12610                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
12611                    let actv = act.slice(0..n_ff_exp);
12612                    let y = m.qmatvec_view(
12613                        e,
12614                        sd,
12615                        0..dl.len,
12616                        &actv,
12617                        1,
12618                        m.down_exps.in_f,
12619                        m.down_exps.out_f,
12620                        dl.qtype,
12621                        dl.row_bytes,
12622                    )?;
12623
12624                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12625                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
12626                }
12627            }
12628            if let Some(worker) = cpu_worker {
12629                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
12630                let cpu_output = e.htod(&cpu_output)?;
12631                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12632                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
12633            }
12634            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
12635                for (j, &ex) in sel.iter().enumerate() {
12636                    if cpu_mask[j] {
12637                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
12638                    }
12639                }
12640            }
12641        }
12642
12643        if memra_reference::hidden_trace::enabled() {
12644            memra_reference::hidden_trace::emit_last_row(
12645                "routed",
12646                il as i64,
12647                t,
12648                n_embd,
12649                &e.dtoh(&moe_out)?,
12650            );
12651        }
12652
12653        Self::trace_moe_out(e, "host", il, &moe_out);
12654        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
12655
12656        Ok(moe_out)
12657    }
12658
12659    /// The glm5 TP-2 EP walk (`MEMRA_GLM5_TP`): one MoE layer's routed-expert FFN over
12660    /// whole-expert contiguous halves. The ROUTER is the unchanged root-side program
12661    /// (`moe_router_logits` + `moe_route_sigmoid_cfg` — bit-identical selection by
12662    /// construction); each rank computes its owned slots' UNWEIGHTED expert rows through
12663    /// the sequential per-expert program (gate/up qmatvec + `ffn_act_lim` + down qmatvec —
12664    /// per-expert-independent dots), the peer's rows return host-canonically, and root
12665    /// applies the slot-ordered `axpy` accumulation chain — the same rounded-operation
12666    /// sequence the plain sequential walk applies. The ROOT-owned shared expert then adds
12667    /// through the extracted `moe_shexp_add`, verbatim.
12668    ///
12669    /// Three transport arms behind ONE routing (lane/glm5-ep-diet; sel/w are shared so a
12670    /// dispatch change can never change selection):
12671    ///   * `MEMRA_GLM5_EP_GROUPED_PRIME` (prefill shapes only): per-rank grouped-GEMM prime
12672    ///     over the rank slabs — the plain grouped-prefill program split by ownership.
12673    ///     Falls closed to the arms below whenever the plain arm's conjuncts do not hold.
12674    ///   * `MEMRA_GLM5_EP_DIET`: the v1 walk's kernels and combine chain with dieted data
12675    ///     movement — one bulk fan-out, zero per-slot host round-trips, one combine launch.
12676    ///     Decode-byte-identical to v1 by construction.
12677    ///   * default: the v1 per-slot host-canonical walk, byte-for-byte.
12678    #[allow(clippy::too_many_arguments)]
12679    fn moe_ffn_glm5_ep(
12680        e: &Engine,
12681        m: &MoeWeights,
12682        ep: &crate::glm5_tp::Glm5EpExps,
12683        z: &CudaSlice<f32>,
12684        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
12685        t: usize,
12686        cfg: &ModelConfig,
12687        il: u16,
12688        prefill: bool,
12689    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12690        let moe = cfg
12691            .moe
12692            .as_ref()
12693            .ok_or("glm5 EP execution requires MoE model metadata")?;
12694        let n_embd = cfg.n_embd as usize;
12695        let n_expert = moe.expert_count as usize;
12696        let n_used = moe.expert_used_count as usize;
12697        let n_ff_exp = moe.expert_ff_length as usize;
12698        let sig = cfg
12699            .sigmoid_router()
12700            .ok_or("glm5 EP execution requires the sigmoid router")?;
12701        let lim_exp = cfg.clamp_exp_at(il as u32);
12702        let lim_shexp = cfg.clamp_shexp_at(il as u32);
12703        if ep.slabs.iter().map(|s| s.n_experts).sum::<usize>() != n_expert {
12704            return Err(format!(
12705                "glm5 EP slabs cover {:?} experts, model declares {n_expert}",
12706                ep.slabs.iter().map(|s| s.n_experts).collect::<Vec<_>>()
12707            )
12708            .into());
12709        }
12710        let rt = &ep.rt;
12711        let ranks = ep.ranks();
12712
12713        // Root router, unchanged program (selection bit-identical to the sequential arm).
12714        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
12715        let (sel_all, w_all) =
12716            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
12717        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
12718        // The route trace taps ride this walk too (sel/w are already host-side here);
12719        // the EP walk must never be a blind spot for the co-activation measurement.
12720        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
12721
12722        // EP grouped prime (`MEMRA_GLM5_EP_GROUPED_PRIME`, default OFF): keyed exactly like
12723        // the plain grouped-prefill arm (`prefill` + t above the per-token tier) and honoring
12724        // the family rollback (`MEMRA_MOE_GROUPED_PREFILL=0` kills it too). The announce
12725        // prints once per process PER FLAG VALUE, in both arms, so an A/B grep distinguishes
12726        // engagement without the line being an arm-local cost.
12727        if prefill && t > MOE_DEV_MAX_T {
12728            static EPGP_ANNOUNCED: std::sync::atomic::AtomicU8 =
12729                std::sync::atomic::AtomicU8::new(0);
12730            let enabled = crate::ep_grouped_prime_on() && moe_grouped_prefill_enabled();
12731            let bit = 1u8 << u8::from(enabled);
12732            if EPGP_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
12733                eprintln!(
12734                    "[glm5-ep-grouped-prime] flag={} t={t} il={il} (announce printed in both \
12735                     arms; engagement is the dispatch counter + per-layer execute line)",
12736                    if enabled { "on" } else { "off" },
12737                );
12738            }
12739            if enabled
12740                && let Some(mut out) =
12741                    Self::moe_ffn_glm5_ep_grouped_prime(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?
12742            {
12743                Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
12744                return Ok(out);
12745            }
12746        }
12747
12748        // spec x TP composition receipt (the #80 review's confirmed perf-shape finding):
12749        // at verify widths (1 < t < prefill) the EP walk PREEMPTS the batched vrows MoE
12750        // pair — a composed verify round pays the sequential per-(token,expert) walk per
12751        // MoE layer, re-inheriting the vrest wall the vrows lane removed on the unsharded
12752        // shape. Announced once so a composed-shape battery cannot mistake a flat
12753        // MOE_VROWS_DISPATCHES counter for a wiring bug; the EP-aware vrows arm is the
12754        // named lever (composition-20260901/box/CELLS.md).
12755        if t > 1 && !prefill {
12756            static EP_VERIFY_MARKED: std::sync::atomic::AtomicBool =
12757                std::sync::atomic::AtomicBool::new(false);
12758            if !EP_VERIFY_MARKED.swap(true, std::sync::atomic::Ordering::Relaxed) {
12759                eprintln!(
12760                    "[glm5-tp-ep] verify rows ride the SEQUENTIAL EP walk (t={t}): the \
12761                     batched vrows MoE pair is preempted by EP; the EP-aware vrows arm is \
12762                     the named lever performance_claim=false"
12763                );
12764            }
12765        }
12766        // EP dispatch diet (`MEMRA_GLM5_EP_DIET`, default OFF): same kernels, same combine
12767        // chain, restructured movement. Read per call — `=0`/unset restores the v1 walk.
12768        if crate::ep_diet_on() {
12769            let mut out = Self::moe_ffn_glm5_ep_diet(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?;
12770            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
12771            return Ok(out);
12772        }
12773
12774        let mut moe_out = e.zeros(t * n_embd)?;
12775        use crate::tp_transport::TpTransport as TpXport;
12776        let hop = ep.rt.hop(e);
12777        // Peer replicas of `z`, one arm each (lane/glm5-tp-transport):
12778        //   host-canonical — v1's EXACT pattern: one draining `dtoh` of the whole block here,
12779        //     then one row `htod` per token PER PEER RANK inside the loop (at two ranks that
12780        //     is byte- and hop-identical to v1). Preserved hop-for-hop so
12781        //     `MEMRA_GLM5_TP_TRANSPORT=0` reproduces the banked v1 walk, not a faster cousin.
12782        //   peer-pull — ONE device copy of the whole `[t, n_embd]` block per peer rank; rows
12783        //     are sliced out of it. Same bytes on the peers, the host uploads and drains
12784        //     removed.
12785        let z_host = match hop.transport {
12786            TpXport::HostCanonical => Some(crate::tp_transport::host_stage_block(
12787                &hop,
12788                0,
12789                z,
12790                t * n_embd,
12791            )?),
12792            TpXport::PeerPull => None,
12793        };
12794        let z_peer_bulks = match hop.transport {
12795            TpXport::PeerPull => Some(crate::tp_transport::fanout_f32(&hop, z, t * n_embd)?),
12796            TpXport::HostCanonical => None,
12797        };
12798        for tok in 0..t {
12799            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
12800            let w = &w_all[tok * n_used..(tok + 1) * n_used];
12801            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
12802            // The host-canonical arm materializes this token's peer rows here (v1's
12803            // per-token `htod`, one per peer rank); the peer-pull arm slices them out of the
12804            // bulk blocks. The holders must outlive the views, hence the two-step.
12805            let z_peer_row_holders: Option<Vec<CudaSlice<f32>>> = match &z_host {
12806                Some(h) => {
12807                    let mut rows = Vec::with_capacity(ranks - 1);
12808                    for r in 1..ranks {
12809                        rows.push(crate::tp_transport::host_row_to(
12810                            &hop,
12811                            r,
12812                            &h[tok * n_embd..(tok + 1) * n_embd],
12813                        )?);
12814                    }
12815                    Some(rows)
12816                }
12817                None => None,
12818            };
12819            // Per slot, in ROUTER SLOT ORDER: compute the UNWEIGHTED expert row on its
12820            // owner, then fmaf-accumulate on root — the plain walk's exact chain.
12821            for (j, &ex) in sel.iter().enumerate() {
12822                let ex = ex as usize;
12823                let owner = ep.owner(ex);
12824                if owner != 0 {
12825                    // Engagement counter FIRST (a red skip still counts as ROUTED).
12826                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES
12827                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12828                    // Gate red arm: dropping the peers' slots MUST diverge — the
12829                    // non-vacuity proof that the peer ranks contribute real expert work.
12830                    if matches!(
12831                        crate::glm5_tp::gate_red(),
12832                        Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
12833                    ) {
12834                        continue;
12835                    }
12836                }
12837                let zin_holder;
12838                let (dev, slab, zin) = if owner == 0 {
12839                    (e, &ep.slabs[0], &zt)
12840                } else {
12841                    zin_holder = match (&z_peer_row_holders, &z_peer_bulks) {
12842                        (Some(rows), _) => rows[owner - 1].slice(0..n_embd),
12843                        (None, Some(bulks)) => {
12844                            bulks[owner - 1].slice(tok * n_embd..(tok + 1) * n_embd)
12845                        }
12846                        (None, None) => {
12847                            return Err(
12848                                "glm5 EP: neither transport arm staged the peer activation".into(),
12849                            );
12850                        }
12851                    };
12852                    (
12853                        crate::glm5_tp::rank_engine(e, rt, owner),
12854                        &ep.slabs[owner],
12855                        &zin_holder,
12856                    )
12857                };
12858                // Placement indirection: the owner's slab packs its experts in
12859                // ascending-id order; `local_of` is the slot (identical to
12860                // `ex - first_expert` under the even split).
12861                let local = ep.local_of[ex] as usize;
12862                let gl = m.gate_exps.expert_stride;
12863                let ul = m.up_exps.expert_stride;
12864                let dl = m.down_exps.expert_stride;
12865                let gate = dev.qmatvec_view(
12866                    &slab.gate,
12867                    local * gl..(local + 1) * gl,
12868                    zin,
12869                    1,
12870                    m.gate_exps.in_f,
12871                    m.gate_exps.out_f,
12872                    m.gate_exps.qtype,
12873                    m.gate_exps.row_bytes,
12874                )?;
12875                let up = dev.qmatvec_view(
12876                    &slab.up,
12877                    local * ul..(local + 1) * ul,
12878                    zin,
12879                    1,
12880                    m.up_exps.in_f,
12881                    m.up_exps.out_f,
12882                    m.up_exps.qtype,
12883                    m.up_exps.row_bytes,
12884                )?;
12885                let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
12886                Self::ffn_act_lim(
12887                    dev,
12888                    cfg,
12889                    &gate,
12890                    &up,
12891                    m.gate_exps.macro_scale(ex),
12892                    m.up_exps.macro_scale(ex),
12893                    lim_exp,
12894                    &mut act,
12895                    n_ff_exp,
12896                )?;
12897                let actv = act.slice(0..n_ff_exp);
12898                let y = dev.qmatvec_view(
12899                    &slab.down,
12900                    local * dl..(local + 1) * dl,
12901                    &actv,
12902                    1,
12903                    m.down_exps.in_f,
12904                    m.down_exps.out_f,
12905                    m.down_exps.qtype,
12906                    m.down_exps.row_bytes,
12907                )?;
12908                // The owner's row returns through the armed transport; the root row stays
12909                // put. The slot-ordered axpy below is the ONE cross-rank arithmetic site and
12910                // it reproduces the sequential walk's accumulate chain operation for
12911                // operation — unchanged by which transport delivered the row.
12912                let y_root = if owner == 0 {
12913                    y
12914                } else {
12915                    crate::tp_transport::return_row_to_root(&hop, owner, &y, n_embd)?
12916                };
12917                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12918                e.axpy_into(
12919                    &y_root,
12920                    w[j] * m.down_exps.macro_scale(ex),
12921                    &mut dst,
12922                    n_embd,
12923                )?;
12924            }
12925        }
12926
12927        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
12928        Ok(moe_out)
12929    }
12930
12931    /// The DIETED glm5 EP walk (`MEMRA_GLM5_EP_DIET`, lane/glm5-ep-diet): the v1 walk's
12932    /// per-slot expert kernels and its exact slot-ordered combine chain, with the data
12933    /// movement restructured in whole groups (the tp2-battery's measured 13-18 ms/token v1
12934    /// join+dispatch tax, attributed to per-token host fan-out x42 layers + ~4-5 sync
12935    /// peer-slot round-trips/layer + interleaved host-blocked issue):
12936    ///
12937    ///   1. ONE bulk peer z fan-out per layer-call ([t, n_embd] in one upload; SKIPPED
12938    ///      entirely when the call routed no peer-owned expert — the placement-map
12939    ///      multiplier: a single-rank layer-call moves zero activation bytes off root).
12940    ///   2. Peer-owned rows compute back-to-back on the peer stream into a compact block
12941    ///      (issue order cannot change bytes: every row is an independent per-expert
12942    ///      program; the combine order below is fixed by the id table, not by issue).
12943    ///   3. Root-owned rows compute on the root stream, un-blocked by peer returns.
12944    ///   4. ONE bulk peer return (peer DtoH + root HtoD of the compact block) replaces the
12945    ///      per-slot round-trip dribble.
12946    ///   5. ONE `moe_pairs_scatter` launch applies the per-token slot-ordered fmaf chain —
12947    ///      the kernel header carries the byte-identity contract vs the zeros +
12948    ///      sequential-`axpy_f32` chain this replaces, and the weights are the SAME host
12949    ///      fold (`w * macro_scale(ex)`) v1 passed per launch.
12950    ///
12951    /// BYTE-IDENTICAL to the v1 walk (and to plain, wherever v1 is) by construction: same
12952    /// kernels over the same bytes; copies (dtod / bulk DtoH+HtoD) preserve bits; the one
12953    /// arithmetic site keeps its exact chain. Transport stays HOST-CANONICAL: the two bulk
12954    /// hops of steps 1 and 4 are the named native-P2P swap points for the box arc (the
12955    /// `MEMRA_STEP_TP_BULK_P2P` precedent: peer copies 61,452 -> ~21/layer on step; the
12956    /// glm5 seam inherits `configure_native_p2p` but does NOT wire it on the rig — the
12957    /// same-device dual-context emulation has no real peer transport to qualify).
12958    #[allow(clippy::too_many_arguments)]
12959    fn moe_ffn_glm5_ep_diet(
12960        e: &Engine,
12961        m: &MoeWeights,
12962        ep: &crate::glm5_tp::Glm5EpExps,
12963        z: &CudaSlice<f32>,
12964        sel_all: &[u32],
12965        w_all: &[f32],
12966        t: usize,
12967        cfg: &ModelConfig,
12968        il: u16,
12969    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12970        use std::sync::atomic::Ordering;
12971        let moe = cfg
12972            .moe
12973            .as_ref()
12974            .ok_or("glm5 EP execution requires MoE model metadata")?;
12975        let n_embd = cfg.n_embd as usize;
12976        let n_used = moe.expert_used_count as usize;
12977        let n_ff_exp = moe.expert_ff_length as usize;
12978        let lim_exp = cfg.clamp_exp_at(il as u32);
12979        let rt = &ep.rt;
12980        let n_pairs = t * n_used;
12981        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
12982            return Err("glm5 EP diet geometry".into());
12983        }
12984        let red_skip_peer = matches!(
12985            crate::glm5_tp::gate_red(),
12986            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
12987        );
12988
12989        // Slab-position table: root-owned pairs pack the slab head in pair order; each peer
12990        // rank's pairs pack a contiguous tail segment (so every rank's bulk return is ONE
12991        // contiguous upload). `ids[p]` is pair p's slab row; the scatter walks ids in slot
12992        // order per token, which is what pins the combine chain to v1's regardless of
12993        // packing. At two ranks this is byte-for-byte the original head/tail split.
12994        let ranks = ep.ranks();
12995        let mut per_rank = vec![0usize; ranks];
12996        for &s in sel_all.iter().take(n_pairs) {
12997            let ex = s as usize;
12998            if ex >= ep.owner_of.len() {
12999                return Err(format!("glm5 EP diet: selection {ex} outside the bank").into());
13000            }
13001            per_rank[ep.owner(ex)] += 1;
13002        }
13003        let mut base = vec![0usize; ranks];
13004        for r in 1..ranks {
13005            base[r] = base[r - 1] + per_rank[r - 1];
13006        }
13007        let mut ids = vec![0i32; n_pairs];
13008        {
13009            let mut k = vec![0usize; ranks];
13010            for (p, id) in ids.iter_mut().enumerate() {
13011                let r = ep.owner(sel_all[p] as usize);
13012                *id = (base[r] + k[r]) as i32;
13013                k[r] += 1;
13014            }
13015        }
13016
13017        crate::glm5_tp::GLM5_EP_DIET_DISPATCHES.fetch_add(1, Ordering::Relaxed);
13018        for r in 1..ranks {
13019            crate::glm5_tp::GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.fetch_add(
13020                if per_rank[r] > 0 {
13021                    (t - 1) as u64
13022                } else {
13023                    t as u64
13024                },
13025                Ordering::Relaxed,
13026            );
13027        }
13028        static EP_DIET_MARKED: std::sync::atomic::AtomicBool =
13029            std::sync::atomic::AtomicBool::new(false);
13030        if !EP_DIET_MARKED.swap(true, Ordering::Relaxed) {
13031            eprintln!(
13032                "[glm5-ep-diet] engaged: bulk fan-out + compact peer staging + single \
13033                 slot-ordered scatter combine; per-slot host round-trips removed \
13034                 transport={} performance_claim=false",
13035                ep.rt.transport.name(),
13036            );
13037        }
13038
13039        // Per-slot expert program, shared verbatim with the v1 walk (same kernels, same
13040        // argument order): gate/up qmatvec + ffn_act_lim + down qmatvec on the OWNING rank.
13041        let expert_row = |dev: &Engine,
13042                          slab: &crate::glm5_tp::EpRankSlab,
13043                          zin: &cudarc::driver::CudaView<f32>,
13044                          ex: usize|
13045         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13046            let local = ep.local_of[ex] as usize;
13047            let gl = m.gate_exps.expert_stride;
13048            let ul = m.up_exps.expert_stride;
13049            let dl = m.down_exps.expert_stride;
13050            let gate = dev.qmatvec_view(
13051                &slab.gate,
13052                local * gl..(local + 1) * gl,
13053                zin,
13054                1,
13055                m.gate_exps.in_f,
13056                m.gate_exps.out_f,
13057                m.gate_exps.qtype,
13058                m.gate_exps.row_bytes,
13059            )?;
13060            let up = dev.qmatvec_view(
13061                &slab.up,
13062                local * ul..(local + 1) * ul,
13063                zin,
13064                1,
13065                m.up_exps.in_f,
13066                m.up_exps.out_f,
13067                m.up_exps.qtype,
13068                m.up_exps.row_bytes,
13069            )?;
13070            let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
13071            Self::ffn_act_lim(
13072                dev,
13073                cfg,
13074                &gate,
13075                &up,
13076                m.gate_exps.macro_scale(ex),
13077                m.up_exps.macro_scale(ex),
13078                lim_exp,
13079                &mut act,
13080                n_ff_exp,
13081            )?;
13082            let actv = act.slice(0..n_ff_exp);
13083            dev.qmatvec_view(
13084                &slab.down,
13085                local * dl..(local + 1) * dl,
13086                &actv,
13087                1,
13088                m.down_exps.in_f,
13089                m.down_exps.out_f,
13090                m.down_exps.qtype,
13091                m.down_exps.row_bytes,
13092            )
13093        };
13094
13095        // Pass 1 — PEERS: one bulk fan-out per pair-owning rank, then every owned row
13096        // back-to-back on that rank's stream into its compact block. No host boundary until
13097        // the bulk returns.
13098        let hop = ep.rt.hop(e);
13099        let mut y_peer_blks: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
13100        for r in 1..ranks {
13101            if per_rank[r] == 0 {
13102                continue;
13103            }
13104            let dev = crate::glm5_tp::rank_engine(e, rt, r);
13105            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only
13106            // (a rank with zero owned pairs moves zero activation bytes off root).
13107            let z_r = crate::tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
13108            // Under the skip-peer-combine red the block stays ZERO for skipped rows (a red
13109            // must drop the peer contribution loudly, never multiply garbage into the chain).
13110            let mut blk = if red_skip_peer {
13111                dev.zeros(per_rank[r] * n_embd)?
13112            } else {
13113                dev.uninit(per_rank[r] * n_embd)?
13114            };
13115            let mut k = 0usize;
13116            for tok in 0..t {
13117                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
13118                for &ex in sel.iter() {
13119                    let ex = ex as usize;
13120                    if ep.owner(ex) != r {
13121                        continue;
13122                    }
13123                    // Engagement counters FIRST (a red skip still counts as ROUTED).
13124                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(1, Ordering::Relaxed);
13125                    crate::glm5_tp::GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED
13126                        .fetch_add(1, Ordering::Relaxed);
13127                    if red_skip_peer {
13128                        k += 1;
13129                        continue;
13130                    }
13131                    let zt_r = z_r.slice(tok * n_embd..(tok + 1) * n_embd);
13132                    let y = expert_row(dev, &ep.slabs[r], &zt_r, ex)?;
13133                    dev.copy_into(&mut blk, k * n_embd, &y, n_embd)?;
13134                    k += 1;
13135                }
13136            }
13137            y_peer_blks[r] = Some(blk);
13138        }
13139
13140        // Pass 2 — ROOT: every root-owned row into the slab head, never blocked on a peer
13141        // return (the v1 walk interleaved root issue behind per-slot peer syncs).
13142        let mut y_all = e.uninit(n_pairs * n_embd)?;
13143        {
13144            let mut k = 0usize;
13145            for tok in 0..t {
13146                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
13147                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13148                for &ex in sel.iter() {
13149                    let ex = ex as usize;
13150                    if ep.owner(ex) != 0 {
13151                        continue;
13152                    }
13153                    let y = expert_row(e, &ep.slabs[0], &zt, ex)?;
13154                    e.copy_into(&mut y_all, k * n_embd, &y, n_embd)?;
13155                    k += 1;
13156                }
13157            }
13158        }
13159
13160        // SWAP POINT 2 (bulk returns) — Pass 3: ONE rank->root block move into each rank's
13161        // tail segment. On host-canonical each is the ONE draining peer sync of that rank's
13162        // layer-call share, exactly as before; on peer-pull each is one event-ordered device
13163        // copy and no host boundary at all.
13164        for r in 1..ranks {
13165            if let Some(blk) = &y_peer_blks[r] {
13166                crate::tp_transport::return_block_to_root(
13167                    &hop,
13168                    r,
13169                    blk,
13170                    &mut y_all,
13171                    base[r] * n_embd,
13172                    per_rank[r] * n_embd,
13173                )?;
13174                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
13175            }
13176        }
13177
13178        // Pass 4 — ONE combine launch. Weights are v1's exact host fold, placed at slab
13179        // positions; the scatter walks each token's pairs in SLOT order (ids[p], p pair-major),
13180        // reproducing zeros + n_used sequential axpy_f32 per the kernel's bit contract.
13181        let mut wd = vec![0f32; n_pairs];
13182        for ((&id, &w), &s) in ids.iter().zip(w_all.iter()).zip(sel_all.iter()) {
13183            wd[id as usize] = w * m.down_exps.macro_scale(s as usize);
13184        }
13185        let pw = e.htod(&wd)?;
13186        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
13187        let toff_d = e.htod_i32(&toff)?;
13188        let ids_d = e.htod_i32(&ids)?;
13189        let mut moe_out = e.uninit(t * n_embd)?; // the scatter fully overwrites
13190        e.moe_pairs_scatter(&y_all, &pw, &toff_d, &ids_d, &mut moe_out, t, n_embd)?;
13191        Ok(moe_out)
13192    }
13193
13194    /// The EP GROUPED PRIME (`MEMRA_GLM5_EP_GROUPED_PRIME`, lane/glm5-ep-diet): the plain
13195    /// walk's grouped-prefill program (`moe_ffn_grouped_prefill_sigmoid`, default ON on the
13196    /// serving artifact — 85 -> 616-639 tok/s prefill in its box A/B) split by expert
13197    /// ownership. Per rank: expert-major CSR over the rank's OWNED (token, expert) pairs,
13198    /// one grouped f16 GEMM per projection over the rank's resident EP slab (pointer tables
13199    /// minted at arm time), the PRE-clamped SwiGLU epilogue, the per-expert macro folds,
13200    /// and the slot-ordered per-token scatter — all composed from the SAME Engine calls the
13201    /// plain arm makes, so per-expert GEMM bytes match the plain grouped arm's (grouping is
13202    /// per expert, and an expert's token rows all live on its owner). The ONE new
13203    /// reassociation is the per-token partial add (root chain + peer chain instead of one
13204    /// 8-term chain) — band-gated on minted NVFP4 slabs (`glm5_ep_diet_doors_gpu`), never
13205    /// claimed byte.
13206    ///
13207    /// Returns `Ok(None)` — fall closed to the sequential EP walk — whenever the plain
13208    /// grouped arm's own admission would (f16g-ineligible qtypes, bank/top-k shape, no
13209    /// sigmoid clamp form). The rig fixture's Q8_0 bank therefore ALWAYS falls closed;
13210    /// `glm5-tp-gate`'s grouped arm proves exactly that (dispatch counter pinned 0, walk
13211    /// bytes unchanged).
13212    #[allow(clippy::too_many_arguments)]
13213    fn moe_ffn_glm5_ep_grouped_prime(
13214        e: &Engine,
13215        m: &MoeWeights,
13216        ep: &crate::glm5_tp::Glm5EpExps,
13217        z: &CudaSlice<f32>,
13218        sel_all: &[u32],
13219        w_all: &[f32],
13220        t: usize,
13221        cfg: &ModelConfig,
13222        il: u16,
13223    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13224        use std::sync::atomic::Ordering;
13225        let moe = cfg
13226            .moe
13227            .as_ref()
13228            .ok_or("glm5 EP grouped prime requires MoE model metadata")?;
13229        let n_embd = cfg.n_embd as usize;
13230        let n_expert = moe.expert_count as usize;
13231        let n_used = moe.expert_used_count as usize;
13232        let n_ff_exp = moe.expert_ff_length as usize;
13233        // The plain grouped arm's admission, mirrored term for term (fall closed, never a
13234        // new admission class). MEMRA_MOE_GATE is the sequential byte-identity oracle; this
13235        // arm is a band class and must not shadow that comparison.
13236        if crate::moe_f16g_mode() == 0 || std::env::var("MEMRA_MOE_GATE").is_ok() {
13237            return Ok(None);
13238        }
13239        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
13240            && f16g_proj_ok(m.up_exps.qtype, n_embd)
13241            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
13242        {
13243            return Ok(None);
13244        }
13245        if n_expert > 512 || n_used == 0 || n_used > 8 {
13246            return Ok(None);
13247        }
13248        let lim_exp = cfg.clamp_exp_at(il as u32);
13249        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
13250            return Err(
13251                "EP grouped prime is qualified for the PRE-clamped SwiGLU form only; \
13252                 a POST-clamp layer must ride the sequential arm"
13253                    .into(),
13254            );
13255        }
13256        let n_pairs = t * n_used;
13257        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
13258            return Err("EP grouped prime geometry".into());
13259        }
13260        let rt = &ep.rt;
13261        let red_skip_peer = matches!(
13262            crate::glm5_tp::gate_red(),
13263            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
13264        );
13265
13266        // One rank's whole grouped program: CSR over OWNED pairs -> grouped gate/up GEMMs ->
13267        // macro folds -> PRE-clamped epilogue -> grouped down GEMM -> CSR->local permute ->
13268        // slot-ordered scatter into the rank partial [t, n_embd] (empty token windows write
13269        // 0.0 — the scatter fully overwrites, so partials add cleanly on root).
13270        let rank_pass = |dev: &Engine,
13271                         rank: u8,
13272                         ptr_row: &CudaSlice<u64>,
13273                         z_dev: &CudaSlice<f32>|
13274         -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
13275            // Expert-major CSR restricted to this rank, local pair index l in ascending
13276            // global-pair order (so per-token slot order == ascending l).
13277            let mut buckets_l: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
13278            let mut local_tok = Vec::new(); // token of local pair l
13279            let mut local_ex = Vec::new(); // expert of local pair l (macro folds)
13280            let mut local_wd = Vec::new(); // v1's exact weight fold, at local positions
13281            let mut local_count_per_tok = vec![0i32; t];
13282            for p in 0..n_pairs {
13283                let ex = sel_all[p] as usize;
13284                if ex >= n_expert {
13285                    return Err(format!("EP grouped prime selection {ex} >= {n_expert}").into());
13286                }
13287                if ep.owner(ex) != rank as usize {
13288                    continue;
13289                }
13290                let l = local_tok.len() as i32;
13291                buckets_l[ex].push(l);
13292                let tok = p / n_used;
13293                local_tok.push(tok as i32);
13294                local_ex.push(ex);
13295                local_wd.push(w_all[p] * m.down_exps.macro_scale(ex));
13296                local_count_per_tok[tok] += 1;
13297            }
13298            let n_owned = local_tok.len();
13299            if n_owned == 0 {
13300                return Ok(None);
13301            }
13302            let mut ex_ids: Vec<i32> = Vec::new();
13303            let mut ex_off: Vec<i32> = vec![0];
13304            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_owned); // local l, CSR order
13305            let mut csr_tok: Vec<i32> = Vec::with_capacity(n_owned);
13306            for (e_id, b) in buckets_l.iter().enumerate() {
13307                if !b.is_empty() {
13308                    ex_ids.push(e_id as i32);
13309                    for &l in b {
13310                        ex_pairs.push(l);
13311                        csr_tok.push(local_tok[l as usize]);
13312                    }
13313                    ex_off.push(ex_pairs.len() as i32);
13314                }
13315            }
13316            let n_active = ex_ids.len();
13317            if n_active == 0 || n_active > 512 {
13318                return Err(format!("EP grouped prime n_active {n_active} outside 1..=512").into());
13319            }
13320
13321            let exi = dev.htod_i32(&ex_ids)?;
13322            let exo = dev.htod_i32(&ex_off)?;
13323            let exp_d = dev.htod_i32(&ex_pairs)?;
13324            let csr_tok_d = dev.htod_i32(&csr_tok)?;
13325
13326            // GATE/UP grouped GEMMs over the rank slab, CSR order end to end.
13327            let (z16, zs) = dev.moe_f16g_act(z_dev, Some(&csr_tok_d), n_embd, n_owned)?;
13328            let mut g = dev.moe_f16_grouped(
13329                ptr_row,
13330                0,
13331                n_expert,
13332                &exi,
13333                &ex_off,
13334                &exo,
13335                &z16,
13336                &zs,
13337                n_embd,
13338                n_ff_exp,
13339                n_active,
13340                n_owned,
13341                m.gate_exps.qtype,
13342                m.gate_exps.row_bytes,
13343            )?;
13344            if m.gate_exps.macros.is_some() {
13345                let mg: Vec<f32> = ex_pairs
13346                    .iter()
13347                    .map(|&l| m.gate_exps.macro_scale(local_ex[l as usize]))
13348                    .collect();
13349                let mg_d = dev.htod(&mg)?;
13350                dev.scale_rows(&mut g, &mg_d, n_ff_exp, n_owned)?;
13351            }
13352            let mut u = dev.moe_f16_grouped(
13353                ptr_row,
13354                1,
13355                n_expert,
13356                &exi,
13357                &ex_off,
13358                &exo,
13359                &z16,
13360                &zs,
13361                n_embd,
13362                n_ff_exp,
13363                n_active,
13364                n_owned,
13365                m.up_exps.qtype,
13366                m.up_exps.row_bytes,
13367            )?;
13368            if m.up_exps.macros.is_some() {
13369                let mu: Vec<f32> = ex_pairs
13370                    .iter()
13371                    .map(|&l| m.up_exps.macro_scale(local_ex[l as usize]))
13372                    .collect();
13373                let mu_d = dev.htod(&mu)?;
13374                dev.scale_rows(&mut u, &mu_d, n_ff_exp, n_owned)?;
13375            }
13376
13377            // Epilogue: PRE-clamped SwiGLU (POST refused above), plain-silu pair otherwise.
13378            let act = match lim_exp {
13379                Some(SwigluClamp::Pre(limit)) => {
13380                    let mut a = dev.uninit(n_owned * n_ff_exp)?;
13381                    dev.swiglu_preclamped_mul_scaled(
13382                        &g,
13383                        &u,
13384                        1.0,
13385                        1.0,
13386                        limit,
13387                        &mut a,
13388                        n_owned * n_ff_exp,
13389                    )?;
13390                    a
13391                }
13392                None => dev.moe_pairs_silu_mul(&g, &u, n_owned * n_ff_exp)?,
13393                Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
13394            };
13395
13396            // DOWN grouped GEMM, permute CSR -> local pair order, slot-ordered scatter.
13397            let (a16, a_s) = dev.moe_f16g_act(&act, None, n_ff_exp, n_owned)?;
13398            let d_csr = dev.moe_f16_grouped(
13399                ptr_row,
13400                2,
13401                n_expert,
13402                &exi,
13403                &ex_off,
13404                &exo,
13405                &a16,
13406                &a_s,
13407                n_ff_exp,
13408                n_embd,
13409                n_active,
13410                n_owned,
13411                m.down_exps.qtype,
13412                m.down_exps.row_bytes,
13413            )?;
13414            let y_local = dev.rows_permute(&d_csr, &exp_d, n_owned, n_embd)?;
13415            let mut toff: Vec<i32> = Vec::with_capacity(t + 1);
13416            let mut acc = 0i32;
13417            toff.push(0);
13418            for &c in &local_count_per_tok {
13419                acc += c;
13420                toff.push(acc);
13421            }
13422            let tids: Vec<i32> = (0..n_owned as i32).collect();
13423            let pw = dev.htod(&local_wd)?;
13424            let toff_d = dev.htod_i32(&toff)?;
13425            let tids_d = dev.htod_i32(&tids)?;
13426            let mut partial = dev.uninit(t * n_embd)?; // scatter fully overwrites
13427            dev.moe_pairs_scatter(&y_local, &pw, &toff_d, &tids_d, &mut partial, t, n_embd)?;
13428            Ok(Some(partial))
13429        };
13430
13431        // Peer passes first (their GEMMs overlap root's), each on its own runtime binding —
13432        // the grouped-MoE FFI follows the RUNTIME device, not cudarc's pushed context
13433        // (`bind_runtime_device`'s contract). Engagement counters count ROUTED peer pairs
13434        // before any red skip, exactly like the sequential walk.
13435        let ranks = ep.ranks();
13436        let n_peer_pairs = sel_all
13437            .iter()
13438            .take(n_pairs)
13439            .filter(|&&ex| ep.owner(ex as usize) != 0)
13440            .count() as u64;
13441        crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(n_peer_pairs, Ordering::Relaxed);
13442        let hop = ep.rt.hop(e);
13443        let mut peer_partials: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
13444        for r in 1..ranks {
13445            let rank_owns_pairs = sel_all
13446                .iter()
13447                .take(n_pairs)
13448                .any(|&ex| ep.owner(ex as usize) == r);
13449            if !rank_owns_pairs {
13450                continue;
13451            }
13452            let dev = crate::glm5_tp::rank_engine(e, rt, r);
13453            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only.
13454            let z_r = crate::tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
13455            dev.bind_runtime_device(dev.ctx().ordinal() as i32)?;
13456            let res = rank_pass(dev, r as u8, &ep.ptr_rows[r], &z_r);
13457            e.bind_runtime_device(e.ctx().ordinal() as i32)?;
13458            peer_partials[r] = res?;
13459        }
13460        let root_partial = rank_pass(e, 0, &ep.ptr_rows[0], z)?;
13461
13462        // Root combine: root partial + bulk-returned peer partials (SWAP POINT 2, the named
13463        // transport shape). One partial add per contributing rank — the same reassociation
13464        // class the two-rank arm band-gated (root chain + per-rank chains instead of one
13465        // 8-term chain), never claimed byte. The skip-peer-combine red drops every peer
13466        // partial AFTER counting — the loud non-vacuity arm.
13467        let mut out = match root_partial {
13468            Some(p) => p,
13469            None => e.zeros(t * n_embd)?,
13470        };
13471        for r in 1..ranks {
13472            if let Some(pp) = &peer_partials[r]
13473                && !red_skip_peer
13474            {
13475                let pp_root = crate::tp_transport::return_row_to_root(&hop, r, pp, t * n_embd)?;
13476                let mut dst = out.slice_mut(0..t * n_embd);
13477                e.axpy_into(&pp_root, 1.0, &mut dst, t * n_embd)?;
13478                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
13479            }
13480        }
13481        crate::glm5_tp::GLM5_EP_GROUPED_PRIME_DISPATCHES.fetch_add(1, Ordering::Relaxed);
13482        static EPGP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
13483        let layer_bit = 1u64 << (il as u64 % 64);
13484        if EPGP_LOGGED.fetch_or(layer_bit, Ordering::Relaxed) & layer_bit == 0 {
13485            eprintln!(
13486                "[glm5-ep-grouped-prime] execute layer={il} tokens={t} \
13487                 provenance=ep-rank-slabs router=sigmoid-host-oracle epilogue=pre-clamped \
13488                 combine=rank-partial-add transport={} performance_claim=false \
13489                 (logged once per layer)",
13490                hop.transport.name(),
13491            );
13492        }
13493        Ok(Some(out))
13494    }
13495
13496    /// Step 3 of the MoE body — SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z —
13497    /// qwen35moe only. OLMoE and most vanilla MoE have NO shared expert (the shexp tensors
13498    /// are absent / `None`); skip it then. gate_inp_shexp is OPTIONAL: qwen35moe gates the
13499    /// shared expert (sigmoid(gate_inp) x sh); MiniMax-M3 (DeepSeek-V3 class) has NO shexp
13500    /// gate — the shared expert adds directly. (Extracted verbatim from the sequential body
13501    /// so the glm5 EP-2 walk adds the ROOT-owned shared expert through the identical
13502    /// program.)
13503    #[allow(clippy::too_many_arguments)]
13504    fn moe_shexp_add(
13505        e: &Engine,
13506        m: &MoeWeights,
13507        z: &CudaSlice<f32>,
13508        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
13509        t: usize,
13510        cfg: &ModelConfig,
13511        lim_shexp: Option<memra_gguf::config::SwigluClamp>,
13512        moe_out: &mut CudaSlice<f32>,
13513    ) -> Result<(), Box<dyn std::error::Error>> {
13514        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13515            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13516        {
13517            let n_embd = cfg.n_embd as usize;
13518            let n_ff_sh = gate_shexp.out_features(); // 512
13519            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
13520            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
13521            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
13522            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
13523            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
13524            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
13525            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
13526            let verify_t = t > 1 && t < PRIME_MIN_T;
13527            let (sg_gate, sg_up) = if t == 1 {
13528                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
13529            } else if verify_t {
13530                (
13531                    e.matmul_decode_exact(gate_shexp, z, t)?,
13532                    e.matmul_decode_exact(up_shexp, z, t)?,
13533                )
13534            } else {
13535                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
13536            };
13537            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
13538            Self::ffn_act_lim(
13539                e,
13540                cfg,
13541                &sg_gate,
13542                &sg_up,
13543                1.0,
13544                1.0,
13545                lim_shexp,
13546                &mut sa,
13547                t * n_ff_sh,
13548            )?;
13549            let sh = if verify_t {
13550                e.matmul_decode_exact(down_shexp, &sa, t)?
13551            } else {
13552                e.matmul(down_shexp, &sa, t)?
13553            }; // [T, n_embd]
13554
13555            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
13556            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
13557            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
13558            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
13559            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
13560            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
13561            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
13562            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
13563            // expert's contribution into every token's residual, so under cross-request
13564            // concat prefill a session's hidden state depended on its co-arrivals' token
13565            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
13566            // DOOR H (`MEMRA_GLM5_HTOD_DIET`): glm5 has no `ffn_gate_inp_shexp`, so this is the
13567            // LIVE arm on the serving artifact and it re-uploaded a constant `vec![1.0f32; t]`
13568            // on every MoE layer-call — 42 pageable HtoD per ship round (26.9% of the round's
13569            // 156). The resident ones buffer feeds the SAME `add_scaled_rows_f32` kernel the
13570            // same 1.0 values, so the arms are bit-identical.
13571            if m.gate_inp_shexp.is_none() && crate::htod_diet_on() {
13572                crate::HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13573                e.add_scaled_rows_ones(&sh, moe_out, n_embd, t)?;
13574                return Ok(());
13575            }
13576            let g = match &m.gate_inp_shexp {
13577                Some(gate_inp_shexp) => {
13578                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
13579                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13580                    } else {
13581                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13582                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
13583                        e.sigmoid(&gs, &mut g, t)?;
13584                        g
13585                    }
13586                }
13587                None => e.htod(&vec![1.0f32; t])?,
13588            };
13589            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
13590            e.add_scaled_rows(&sh, &g, moe_out, n_embd, t)?;
13591        }
13592        Ok(())
13593    }
13594
13595    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
13596    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
13597    pub fn stage1_h2d_per_token(&self) -> u64 {
13598        use crate::hybrid::Ffn;
13599        let n_used = self
13600            .cfg
13601            .moe
13602            .as_ref()
13603            .map(|m| m.expert_used_count as u64)
13604            .unwrap_or(0);
13605        let mut bytes = 0u64;
13606        for l in self.layers.iter() {
13607            if let Ffn::Moe(m) = &l.ffn {
13608                bytes += n_used
13609                    * (m.gate_exps.max_expert_bytes()
13610                        + m.up_exps.max_expert_bytes()
13611                        + m.down_exps.max_expert_bytes()) as u64;
13612            }
13613        }
13614        bytes
13615    }
13616
13617    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
13618    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
13619    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
13620    pub(crate) fn max_moe_block(&self) -> usize {
13621        use crate::hybrid::Ffn;
13622        let mut mx = 0usize;
13623        let mut scan = |ffn: &Ffn| {
13624            if let Ffn::Moe(m) = ffn {
13625                mx = mx
13626                    .max(m.gate_exps.max_expert_bytes())
13627                    .max(m.up_exps.max_expert_bytes())
13628                    .max(m.down_exps.max_expert_bytes());
13629            }
13630        };
13631        for l in self.layers.iter() {
13632            scan(&l.ffn);
13633        }
13634        if let Some(mtp) = self.mtp.as_ref() {
13635            scan(&mtp.ffn);
13636        }
13637        mx
13638    }
13639
13640    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
13641    /// but have no bytes and therefore consume no residency slot.
13642    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
13643        use crate::hybrid::Ffn;
13644        let mut sizes = Vec::new();
13645        let mut scan = |ffn: &Ffn| {
13646            let Ffn::Moe(m) = ffn else { return };
13647            for ex in 0..m.gate_exps.n_expert {
13648                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
13649                    continue;
13650                }
13651                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
13652                    let len = exps.expert_layout(ex).len;
13653                    if len > 0 {
13654                        sizes.push(len);
13655                    }
13656                }
13657            }
13658        };
13659        for layer in &self.layers {
13660            scan(&layer.ffn);
13661        }
13662        if let Some(mtp) = &self.mtp {
13663            scan(&mtp.ffn);
13664        }
13665        sizes
13666    }
13667
13668    /// Persist the frozen residency set so a later process can restage it directly and skip
13669    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
13670    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
13671    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
13672    /// post-freeze argmax gate still validates the serving assignment.
13673    pub fn save_cpu_expert_residency_profile(
13674        &self,
13675        e: &Engine,
13676        path: &std::path::Path,
13677    ) -> Result<(), Box<dyn std::error::Error>> {
13678        let Some(ids) = e.export_moe_residency() else {
13679            return Err("no MoE residency cache to persist".into());
13680        };
13681        let mut body = format!(
13682            "memra-freeze-profile v1 max_block={} blocks={}\n",
13683            self.max_moe_block(),
13684            ids.len()
13685        );
13686        for (layer, proj, ex) in &ids {
13687            body.push_str(&format!("{layer} {proj} {ex}\n"));
13688        }
13689        let tmp = path.with_extension("tmp");
13690        std::fs::write(&tmp, body)?;
13691        std::fs::rename(&tmp, path)?;
13692        println!(
13693            "[moe-cache] freeze profile saved: {} blocks -> {}",
13694            ids.len(),
13695            path.display()
13696        );
13697        Ok(())
13698    }
13699
13700    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
13701    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
13702    /// missing or its header does not match this model's slot geometry.
13703    pub fn restore_cpu_expert_residency_profile(
13704        &self,
13705        e: &Engine,
13706        path: &std::path::Path,
13707    ) -> Result<bool, Box<dyn std::error::Error>> {
13708        use crate::hybrid::Ffn;
13709        use crate::moe_cache::BlockId;
13710        let Ok(content) = std::fs::read_to_string(path) else {
13711            return Ok(false);
13712        };
13713        let mut lines = content.lines();
13714        let Some(header) = lines.next() else {
13715            return Ok(false);
13716        };
13717        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
13718        if !header.starts_with(&expected) {
13719            println!(
13720                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
13721                path.display()
13722            );
13723            return Ok(false);
13724        }
13725        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
13726            std::collections::HashMap::new();
13727        for line in lines {
13728            let mut fields = line.split_whitespace();
13729            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
13730            else {
13731                continue;
13732            };
13733            let (Ok(layer), Ok(proj), Ok(ex)) =
13734                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
13735            else {
13736                continue;
13737            };
13738            by_layer
13739                .entry(layer)
13740                .or_default()
13741                .push(BlockId::new(layer, proj, ex));
13742        }
13743        let requested: usize = by_layer.values().map(Vec::len).sum();
13744        if requested == 0 {
13745            return Ok(false);
13746        }
13747        let max_block = self.max_moe_block();
13748        let mut restaged = 0usize;
13749        let mut stage_layer =
13750            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
13751                let Ffn::Moe(m) = ffn else { return Ok(()) };
13752                let Some(ids) = by_layer.get(&layer_index) else {
13753                    return Ok(());
13754                };
13755                e.with_moe_cache(max_block, |cache, eng| {
13756                    for id in ids {
13757                        if cache.restage_block(*id, m, eng)? {
13758                            restaged += 1;
13759                        }
13760                    }
13761                    Ok(())
13762                })
13763            };
13764        for (index, layer) in self.layers.iter().enumerate() {
13765            stage_layer(index as u16, &layer.ffn)?;
13766        }
13767        if let Some(mtp) = self.mtp.as_ref() {
13768            stage_layer(u16::MAX, &mtp.ffn)?;
13769        }
13770        e.freeze_moe_cache();
13771        println!(
13772            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
13773            path.display()
13774        );
13775        Ok(true)
13776    }
13777
13778    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
13779    pub fn freeze_cpu_expert_residency(
13780        &self,
13781        e: &Engine,
13782    ) -> Result<(), Box<dyn std::error::Error>> {
13783        e.freeze_moe_cache();
13784        Ok(())
13785    }
13786
13787    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
13788    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
13789    /// the model's activation exactly.
13790    ///
13791    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
13792    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
13793    /// form for anything that can land on a clamped layer.
13794    pub fn ffn_act(
13795        e: &Engine,
13796        cfg: &ModelConfig,
13797        gate: &CudaSlice<f32>,
13798        up: &CudaSlice<f32>,
13799        act: &mut CudaSlice<f32>,
13800        n: usize,
13801    ) -> Result<(), Box<dyn std::error::Error>> {
13802        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
13803    }
13804
13805    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
13806    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
13807    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
13808    #[allow(clippy::too_many_arguments)]
13809    pub(crate) fn ffn_act_scaled(
13810        e: &Engine,
13811        cfg: &ModelConfig,
13812        gate: &CudaSlice<f32>,
13813        up: &CudaSlice<f32>,
13814        gs: f32,
13815        us: f32,
13816        act: &mut CudaSlice<f32>,
13817        n: usize,
13818    ) -> Result<(), Box<dyn std::error::Error>> {
13819        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
13820    }
13821
13822    /// ffn_act_scaled + a PER-LAYER clamped SwiGLU. `limit`:
13823    ///   * `None`          -> the unclamped dispatch (every arch with no live clamp).
13824    ///   * `Some(Post(l))` -> step35: `min(silu(gate*gs), l) * clamp(up*us, +-l)`
13825    ///     (llama-graph.cpp:2146/1751, non-DEEPSEEK4 branch).
13826    ///   * `Some(Pre(l))`  -> glm5_next: `silu(min(gate*gs, l)) * clamp(up*us, +-l)`.
13827    ///     Callers source it from `cfg.clamp_exp_at(il)` (routed experts) or `cfg.clamp_shexp_at(il)`
13828    ///     (shared expert / dense MLP) — on step35 the two arrays are SEPARATE and a layer can have
13829    ///     one without the other. The `> 1e-6` eps gate lives in the accessors, so a `Some` here is
13830    ///     already known live. The match is exhaustive so a new clamp form cannot default to either
13831    ///     existing one.
13832    #[allow(clippy::too_many_arguments)]
13833    pub(crate) fn ffn_act_lim(
13834        e: &Engine,
13835        cfg: &ModelConfig,
13836        gate: &CudaSlice<f32>,
13837        up: &CudaSlice<f32>,
13838        gs: f32,
13839        us: f32,
13840        limit: Option<SwigluClamp>,
13841        act: &mut CudaSlice<f32>,
13842        n: usize,
13843    ) -> Result<(), Box<dyn std::error::Error>> {
13844        if let Some(m3) = cfg.m3.as_ref() {
13845            debug_assert!(
13846                limit.is_none(),
13847                "m3 swigluoai and the step35/glm5_next clamps are different archs"
13848            );
13849            return e.swigluoai_mul_scaled(
13850                gate,
13851                up,
13852                gs,
13853                us,
13854                m3.swiglu_alpha,
13855                m3.swiglu_limit,
13856                act,
13857                n,
13858            );
13859        }
13860        match limit {
13861            Some(SwigluClamp::Post(l)) => {
13862                return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
13863            }
13864            Some(SwigluClamp::Pre(l)) => {
13865                return e.swiglu_preclamped_mul_scaled(gate, up, gs, us, l, act, n);
13866            }
13867            None => {}
13868        }
13869        if gs == 1.0 && us == 1.0 {
13870            return e.silu_mul(gate, up, act, n);
13871        }
13872        e.silu_mul_scaled(gate, up, gs, us, act, n)
13873    }
13874
13875    /// The bare POST limit for the fused kernels whose epilogue HARDCODES step35's form
13876    /// (`matvec_bf16_dual_silu` / `_rows`, qmatvec.cu:10676). `Ok` = the kernel may run;
13877    /// `Err(())` = glm5_next's PRE form, which has no fused twin, and the caller MUST return its
13878    /// not-handled value so the layer falls through to the unfused `ffn_act_lim` seam. Feeding a
13879    /// PRE limit to a POST epilogue compiles, runs, and returns plausible-but-wrong logits.
13880    fn fused_post_limit(lim: Option<SwigluClamp>) -> Result<Option<f32>, ()> {
13881        match lim {
13882            None => Ok(None),
13883            Some(SwigluClamp::Post(l)) => Ok(Some(l)),
13884            Some(SwigluClamp::Pre(_)) => Err(()),
13885        }
13886    }
13887
13888    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
13889    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
13890    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
13891    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
13892    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
13893    fn moe_route(
13894        e: &Engine,
13895        logits: &CudaSlice<f32>,
13896        t: usize,
13897        n_expert: usize,
13898        n_used: usize,
13899    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
13900        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
13901    }
13902
13903    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
13904    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
13905    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
13906    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
13907    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
13908    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
13909    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
13910    #[allow(clippy::too_many_arguments)]
13911    fn moe_route_sigmoid_cfg(
13912        e: &Engine,
13913        logits: &CudaSlice<f32>,
13914        t: usize,
13915        n_expert: usize,
13916        n_used: usize,
13917        m: &MoeWeights,
13918        (sf, route_norm): (f32, bool),
13919    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
13920        if sigmoid_router_enabled() {
13921            return e.moe_router_sigmoid_topk_host(
13922                logits,
13923                t,
13924                n_expert,
13925                n_used,
13926                m.active_count(),
13927                &m.exp_probs_b_dev,
13928                &m.active_experts_dev,
13929                sf,
13930                route_norm,
13931            );
13932        }
13933        let lg = e.dtoh(logits)?;
13934        Self::moe_route_sigmoid_host(
13935            &lg,
13936            t,
13937            n_expert,
13938            n_used,
13939            m.exp_probs_b.as_deref(),
13940            sf,
13941            route_norm,
13942            m.active_experts.as_deref(),
13943        )
13944    }
13945
13946    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
13947    /// the existing softmax device kernel has no mask input.
13948    #[allow(clippy::excessive_precision)] // allow: literal kept verbatim from the reference/measured value
13949    fn moe_route_cfg(
13950        e: &Engine,
13951        logits: &CudaSlice<f32>,
13952        t: usize,
13953        n_expert: usize,
13954        n_used: usize,
13955        active: Option<&[bool]>,
13956    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
13957        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
13958        // rollback) via the single-sync pinned readback — softmax arch only.
13959        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
13960            return e.moe_router_topk_host(logits, t, n_expert, n_used);
13961        }
13962        // Host oracle (the §D bit-identity reference).
13963        let lg = e.dtoh(logits)?; // [T*n_expert] host
13964        let mut sel = vec![0u32; t * n_used];
13965        let mut w_out = vec![0f32; t * n_used];
13966        for tok in 0..t {
13967            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
13968            // softmax over ALL n_expert (stable: subtract max)
13969            let maxl = row
13970                .iter()
13971                .enumerate()
13972                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
13973                .map(|(_, &x)| x)
13974                .fold(f32::NEG_INFINITY, f32::max);
13975            let mut probs = vec![0f32; n_expert];
13976            let mut den = 0f32;
13977            for i in 0..n_expert {
13978                if active.is_some_and(|mask| !mask[i]) {
13979                    continue;
13980                }
13981                let x = (row[i] - maxl).exp();
13982                probs[i] = x;
13983                den += x;
13984            }
13985            for p in probs.iter_mut() {
13986                *p /= den;
13987            }
13988            // stable DESC sort: prob DESC, ascending-index tiebreak.
13989            let mut idx: Vec<usize> = (0..n_expert)
13990                .filter(|&i| active.is_none_or(|mask| mask[i]))
13991                .collect();
13992            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
13993            let sl = &idx[..n_used];
13994            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
13995            let mut ws: f32 = wv.iter().sum();
13996            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
13997            for x in wv.iter_mut() {
13998                *x /= ws;
13999            }
14000            for j in 0..n_used {
14001                sel[tok * n_used + j] = sl[j] as u32;
14002                w_out[tok * n_used + j] = wv[j];
14003            }
14004        }
14005        Ok((sel, w_out))
14006    }
14007
14008    #[allow(clippy::too_many_arguments)]
14009    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
14010    fn moe_route_sigmoid_with_input(
14011        e: &Engine,
14012        logits: &CudaSlice<f32>,
14013        input: &CudaSlice<f32>,
14014        t: usize,
14015        in_features: usize,
14016        n_expert: usize,
14017        n_used: usize,
14018        bias: Option<&[f32]>,
14019        (sf, route_norm): (f32, bool),
14020        active: Option<&[bool]>,
14021    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
14022        let logit_values =
14023            active_matrix_values(logits.len(), t, n_expert, "sigmoid router logits")?;
14024        let input_values =
14025            active_matrix_values(input.len(), t, in_features, "sigmoid router input")?;
14026        let (lg, input) = e.dtoh_pair_views(
14027            &logits.slice(0..logit_values),
14028            &input.slice(0..input_values),
14029        )?;
14030        let (sel, w) =
14031            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
14032        Ok((sel, w, input))
14033    }
14034
14035    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
14036    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
14037    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
14038    /// active mask, prebuilt projection descriptors) so no model reference escapes.
14039    pub fn start_moe_prefetch_predictor(
14040        &self,
14041        e: &Engine,
14042        cfg: &ModelConfig,
14043    ) -> Result<(), Box<dyn std::error::Error>> {
14044        use crate::hybrid::Ffn;
14045        let Some(sig) = cfg.sigmoid_router() else {
14046            return Err("prefetch predictor requires a sigmoid-router arch".into());
14047        };
14048        let resident: std::collections::HashSet<(u16, u8, u16)> = e
14049            .export_moe_residency()
14050            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
14051            .into_iter()
14052            .collect();
14053        let mut layers = Vec::new();
14054        for (index, layer) in self.layers.iter().enumerate() {
14055            let Ffn::Moe(m) = &layer.ffn else { continue };
14056            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
14057                continue;
14058            };
14059            let router = e.dtoh(data)?;
14060            let n_expert = m.gate_exps.n_expert;
14061            let n_embd = m.gate_exps.in_f;
14062            if router.len() != n_embd * n_expert {
14063                continue;
14064            }
14065            let build = |exps: &crate::model::HostExps| {
14066                (0..n_expert)
14067                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
14068                    .collect::<Vec<_>>()
14069            };
14070            layers.push((
14071                index as u16,
14072                crate::cpu_experts::PredictLayerInit {
14073                    router,
14074                    bias: m.exp_probs_b.clone(),
14075                    active: m.active_experts.clone(),
14076                    n_embd,
14077                    n_used: cfg
14078                        .moe
14079                        .as_ref()
14080                        .map(|moe| moe.expert_used_count as usize)
14081                        .ok_or("prefetch predictor requires MoE config")?,
14082                    sig,
14083                    weights_n_expert: n_expert,
14084                    gate: build(&m.gate_exps),
14085                    up: build(&m.up_exps),
14086                    down: build(&m.down_exps),
14087                },
14088            ));
14089        }
14090        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
14091    }
14092
14093    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
14094    /// selection math to the rollback runtime, applied to host-computed logits.
14095    #[allow(clippy::too_many_arguments)]
14096    pub fn moe_route_sigmoid_host_public(
14097        logits: &[f32],
14098        t: usize,
14099        n_expert: usize,
14100        n_used: usize,
14101        bias: Option<&[f32]>,
14102        sf: f32,
14103        route_norm: bool,
14104        active: Option<&[bool]>,
14105    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
14106        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
14107    }
14108
14109    #[allow(clippy::too_many_arguments)]
14110    fn moe_route_sigmoid_host(
14111        lg: &[f32],
14112        t: usize,
14113        n_expert: usize,
14114        n_used: usize,
14115        bias: Option<&[f32]>,
14116        sf: f32,
14117        route_norm: bool,
14118        active: Option<&[bool]>,
14119    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
14120        let active_count = active
14121            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
14122            .unwrap_or(n_expert);
14123        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
14124        if lg.len() != t * n_expert {
14125            return Err(format!(
14126                "sigmoid router logits length mismatch: got {}, expected {}",
14127                lg.len(),
14128                t * n_expert,
14129            )
14130            .into());
14131        }
14132        let mut sel = vec![0u32; t * n_used];
14133        let mut w_out = vec![0f32; t * n_used];
14134        for tok in 0..t {
14135            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
14136            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
14137            // selection score = sigmoid + bias; weight = plain sigmoid.
14138            let selsc: Vec<f32> = match bias {
14139                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
14140                None => scores.clone(),
14141            };
14142            let mut idx: Vec<usize> = (0..n_expert)
14143                .filter(|&i| active.is_none_or(|mask| mask[i]))
14144                .collect();
14145            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
14146            let sl = &idx[..n_used];
14147            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
14148            if route_norm {
14149                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
14150                for x in wv.iter_mut() {
14151                    *x = *x / ws * sf;
14152                }
14153            } else {
14154                for x in wv.iter_mut() {
14155                    *x *= sf;
14156                }
14157            }
14158            for j in 0..n_used {
14159                sel[tok * n_used + j] = sl[j] as u32;
14160                w_out[tok * n_used + j] = wv[j];
14161            }
14162        }
14163        Ok((sel, w_out))
14164    }
14165
14166    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
14167    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
14168    /// macro-scaled experts, and observation modes are denied by the caller.
14169    #[allow(clippy::too_many_arguments)]
14170    fn moe_ffn_sigmoid_dev(
14171        e: &Engine,
14172        m: &MoeWeights,
14173        z: &CudaSlice<f32>,
14174        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
14175        logits: &CudaSlice<f32>,
14176        t: usize,
14177        cfg: &ModelConfig,
14178        il: u16,
14179        (scaling_factor, route_norm): (f32, bool),
14180    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14181        crate::moe_rp_refuse(
14182            m.dev_exps.as_ref().is_some_and(|d| d.rp),
14183            "moe_ffn_sigmoid_dev",
14184        )?; // memra#147: no split-plane arm here
14185        let moe = cfg.moe.as_ref().unwrap();
14186        let n_embd = cfg.n_embd as usize;
14187        let n_expert = moe.expert_count as usize;
14188        let n_used = moe.expert_used_count as usize;
14189        let n_ff_exp = moe.expert_ff_length as usize;
14190        let dev = m.dev_exps.as_ref().unwrap();
14191        debug_assert_eq!(dev.dev, e.ctx().ordinal());
14192        debug_assert!(m.has_uniform_expert_layout());
14193        debug_assert!(!m.has_macros);
14194
14195        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
14196            logits,
14197            t,
14198            n_expert,
14199            n_used,
14200            m.active_count(),
14201            &m.exp_probs_b_dev,
14202            &m.active_experts_dev,
14203            scaling_factor,
14204            route_norm,
14205        )?;
14206        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
14207        crate::moe_sel_dump::record_device(e, il, t, n_used, &sel_d, &w_d)?;
14208        if let Some(fp8) = dev.fp8_blk.as_ref() {
14209            debug_assert_eq!(m.gate_exps.qtype, crate::QT_F8_E4M3_BLK);
14210            debug_assert_eq!(m.up_exps.qtype, crate::QT_F8_E4M3_BLK);
14211            debug_assert_eq!(m.down_exps.qtype, crate::QT_F8_E4M3_BLK);
14212            debug_assert_eq!(fp8.gate.rows, m.gate_exps.out_f.div_ceil(128));
14213            debug_assert_eq!(fp8.up.rows, m.up_exps.out_f.div_ceil(128));
14214            debug_assert_eq!(fp8.down.rows, m.down_exps.out_f.div_ceil(128));
14215
14216            // Official Step-3.7 FP8 uses dynamic per-token/per-128 E4M3
14217            // activations with block-128 E4M3 weights. This deliberately
14218            // simple resident reference is the correctness oracle for later
14219            // grouped and TP/EP execution. MEMRA_ST_E4M3=0 chooses the
14220            // load-time Q8 diagnostic representation, so one process never
14221            // crosses between numerical programs.
14222            let selected = e.dtoh_i32(&sel_d)?;
14223            let route_weights = e.dtoh(&w_d)?;
14224            let mut moe_out = e.zeros(t * n_embd)?;
14225            for tok in 0..t {
14226                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
14227                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14228                for j in 0..n_used {
14229                    let pair = tok * n_used + j;
14230                    let expert = selected[pair] as usize;
14231                    let gate = Self::moe_resident_fp8_e4m3(
14232                        e,
14233                        &m.gate_exps,
14234                        &dev.gate,
14235                        &fp8.gate,
14236                        expert,
14237                        &zt,
14238                        1,
14239                    )?;
14240                    let up = Self::moe_resident_fp8_e4m3(
14241                        e, &m.up_exps, &dev.up, &fp8.up, expert, &zt, 1,
14242                    )?;
14243                    let mut act = e.uninit(n_ff_exp)?;
14244                    Self::ffn_act_lim(
14245                        e,
14246                        cfg,
14247                        &gate,
14248                        &up,
14249                        1.0,
14250                        1.0,
14251                        cfg.clamp_exp_at(il as u32),
14252                        &mut act,
14253                        n_ff_exp,
14254                    )?;
14255                    let act = act.slice(0..n_ff_exp);
14256                    let down = Self::moe_resident_fp8_e4m3(
14257                        e,
14258                        &m.down_exps,
14259                        &dev.down,
14260                        &fp8.down,
14261                        expert,
14262                        &act,
14263                        1,
14264                    )?;
14265                    e.axpy_into(&down, route_weights[pair], &mut dst, n_embd)?;
14266                }
14267            }
14268            if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
14269                eprintln!(
14270                    "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} \
14271                     native=fp8blk-w8a8-e4m3-reference clamp={}",
14272                    cfg.clamp_exp_at(il as u32).is_some(),
14273                );
14274            }
14275            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
14276            return Ok(moe_out);
14277        }
14278        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
14279            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14280            (combined, combined)
14281        } else {
14282            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14283        };
14284        let (zq, zd) = match (t, zq8) {
14285            (1, Some((q, d))) => (q.clone(), d.clone()),
14286            _ => e.quantize_q8_1(z, t, n_embd)?,
14287        };
14288        let n_pairs = t * n_used;
14289        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
14290            // The final Step layers retain the established separate gate/up -> clamp -> down
14291            // arithmetic. Pair rows are derived from token position; selected expert ids and
14292            // routing weights remain the device router's buffers throughout.
14293            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
14294            let pair_tok_d = e.htod_i32(&pair_tok)?;
14295            let gate = e.moe_pairs_matvec_q8(
14296                &dev.ptr_row,
14297                0,
14298                &pair_tok_d,
14299                &sel_d,
14300                &zq,
14301                &zd,
14302                n_embd,
14303                n_ff_exp,
14304                n_expert,
14305                n_pairs,
14306                m.gate_exps.qtype,
14307                gate_row_bytes,
14308            )?;
14309            let up = e.moe_pairs_matvec_q8(
14310                &dev.ptr_row,
14311                1,
14312                &pair_tok_d,
14313                &sel_d,
14314                &zq,
14315                &zd,
14316                n_embd,
14317                n_ff_exp,
14318                n_expert,
14319                n_pairs,
14320                m.up_exps.qtype,
14321                up_row_bytes,
14322            )?;
14323            let mut act = e.uninit(n_pairs * n_ff_exp)?;
14324            Self::ffn_act_lim(
14325                e,
14326                cfg,
14327                &gate,
14328                &up,
14329                1.0,
14330                1.0,
14331                cfg.clamp_exp_at(il as u32),
14332                &mut act,
14333                n_pairs * n_ff_exp,
14334            )?;
14335            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
14336            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
14337            let pair_self_d = e.htod_i32(&pair_self)?;
14338            let down = e.moe_pairs_matvec_q8(
14339                &dev.ptr_row,
14340                2,
14341                &pair_self_d,
14342                &sel_d,
14343                &aq2,
14344                &ad2,
14345                n_ff_exp,
14346                n_embd,
14347                n_expert,
14348                n_pairs,
14349                m.down_exps.qtype,
14350                m.down_exps.row_bytes,
14351            )?;
14352            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14353            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
14354            let tok_off_d = e.htod_i32(&tok_off)?;
14355            let tok_ids_d = e.htod_i32(&tok_ids)?;
14356            let mut output = e.uninit(t * n_embd)?;
14357            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
14358            output
14359        } else {
14360            let act = e.moe_gate_up_silu8_dev_q8_rows(
14361                &dev.ptr_row,
14362                &sel_d,
14363                &zq,
14364                &zd,
14365                t,
14366                n_embd,
14367                n_ff_exp,
14368                n_used,
14369                n_expert,
14370                m.gate_exps.qtype,
14371                m.up_exps.qtype,
14372                gate_row_bytes,
14373                up_row_bytes,
14374                &m.dev_macros,
14375            )?;
14376            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
14377            let mut output = e.uninit(t * n_embd)?;
14378            e.moe_down8_fma_dev_q8_rows_g(
14379                &dev.ptr_row,
14380                &sel_d,
14381                &w_d,
14382                &aq2,
14383                &ad2,
14384                &mut output,
14385                t,
14386                n_ff_exp,
14387                n_embd,
14388                n_used,
14389                n_expert,
14390                m.down_exps.qtype,
14391                m.down_exps.row_bytes,
14392            )?;
14393            output
14394        };
14395
14396        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
14397            eprintln!(
14398                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
14399                cfg.clamp_exp_at(il as u32).is_some(),
14400                dev.gu_il,
14401            );
14402        }
14403        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
14404        Ok(moe_out)
14405    }
14406
14407    #[allow(clippy::too_many_arguments)]
14408    fn moe_resident_fp8_e4m3(
14409        e: &Engine,
14410        exps: &crate::model::HostExps,
14411        bytes: &CudaSlice<u8>,
14412        scales: &crate::hybrid::DevExpertFp8ProjectionScales,
14413        expert: usize,
14414        x: &cudarc::driver::CudaView<f32>,
14415        m: usize,
14416    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14417        let layout = exps.expert_layout(expert);
14418        debug_assert_eq!(layout.qtype, crate::QT_F8_E4M3_BLK);
14419        debug_assert_eq!(scales.rows * scales.cols, scales.expert_stride);
14420        let byte_start = expert * exps.expert_stride;
14421        let scale_start = expert * scales.expert_stride;
14422        let weight = bytes.slice(byte_start..byte_start + layout.len);
14423        let scale = scales
14424            .scales
14425            .slice(scale_start..scale_start + scales.expert_stride);
14426        e.qmatvec_mmq_fp8_blk_view(&weight, &scale, x, m, exps.in_f, exps.out_f)
14427    }
14428
14429    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
14430    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
14431    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
14432    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
14433    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
14434    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
14435    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
14436    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
14437    fn moe_ffn_pairs(
14438        e: &Engine,
14439        m: &MoeWeights,
14440        z: &CudaSlice<f32>,
14441        logits: &CudaSlice<f32>,
14442        t: usize,
14443        cfg: &ModelConfig,
14444    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14445        crate::moe_rp_refuse(m.dev_exps.as_ref().is_some_and(|d| d.rp), "moe_ffn_pairs")?; // memra#147: no split-plane arm here
14446        let moe = cfg.moe.as_ref().unwrap();
14447        let n_embd = cfg.n_embd as usize;
14448        let n_expert = moe.expert_count as usize;
14449        let n_used = moe.expert_used_count as usize;
14450        let n_ff_exp = moe.expert_ff_length as usize;
14451        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
14452        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
14453        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
14454        // that forgets the gate fails loudly in debug instead of returning wrong logits.
14455        debug_assert!(
14456            !cfg.swiglu_clamped_anywhere(),
14457            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
14458        );
14459        let dev = m.dev_exps.as_ref().unwrap();
14460        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
14461        let (rbg_d, rbu_d) = if dev.gu_il {
14462            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14463            (sxx, sxx)
14464        } else {
14465            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14466        };
14467
14468        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
14469        let n_pairs = t * n_used;
14470        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
14471        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
14472        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
14473        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
14474        let pair_w: Vec<f32> = w_all.clone();
14475        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14476        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
14477        let pt = e.htod_i32(&pair_tok)?;
14478        let px = e.htod_i32(&pair_ex)?;
14479        let pw = e.htod(&pair_w)?;
14480        let toff = e.htod_i32(&tok_off)?;
14481        let tids = e.htod_i32(&tok_ids)?;
14482
14483        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
14484        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
14485        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
14486        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14487        for p in 0..n_pairs {
14488            by_ex[pair_ex[p] as usize].push(p as i32);
14489        }
14490        let mut ex_ids: Vec<i32> = Vec::new();
14491        let mut ex_off: Vec<i32> = vec![0];
14492        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
14493        for (ex, list) in by_ex.iter().enumerate() {
14494            if list.is_empty() {
14495                continue;
14496            }
14497            ex_ids.push(ex as i32);
14498            ex_pairs.extend_from_slice(list);
14499            ex_off.push(ex_pairs.len() as i32);
14500        }
14501        let n_active = ex_ids.len();
14502        let exi = e.htod_i32(&ex_ids)?;
14503        let exo = e.htod_i32(&ex_off)?;
14504        let exp_d = e.htod_i32(&ex_pairs)?;
14505        let _ = &px; // pair-major twin keeps it; em path uses CSR
14506
14507        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
14508        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
14509        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
14510        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
14511        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
14512        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
14513        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
14514        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
14515        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
14516        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
14517        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
14518        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
14519        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
14520        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
14521        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
14522        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
14523        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
14524        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
14525        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
14526        let mma_t = *MMA_T.get_or_init(|| {
14527            std::env::var("MEMRA_MOE_MMA_T")
14528                .ok()
14529                .and_then(|v| v.parse().ok())
14530                .unwrap_or(16)
14531        });
14532        let use_mma = std::env::var("MEMRA_MOE_MMA")
14533            .map(|v| v != "0")
14534            .unwrap_or(true)
14535            && t >= mma_t
14536            && q8_expert_dec_supported(m.gate_exps.qtype)
14537            && q8_expert_dec_supported(m.up_exps.qtype)
14538            && q8_expert_dec_supported(m.down_exps.qtype)
14539            && n_embd.is_multiple_of(256)
14540            && n_ff_exp.is_multiple_of(256);
14541        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
14542        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
14543        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
14544        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
14545        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
14546        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
14547        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
14548        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
14549        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
14550        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
14551        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
14552        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
14553        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
14554        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
14555        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
14556        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
14557            && q8_expert_dec_supported(m.up_exps.qtype)
14558            && q8_expert_dec_supported(m.down_exps.qtype)
14559            && n_embd.is_multiple_of(256)
14560            && n_ff_exp.is_multiple_of(256);
14561        let f16g_mode = crate::moe_f16g_mode();
14562        let f16g = f16g_mode != 0
14563            && t >= mma_t
14564            && (f16g_mode != 3 || !mma_capable)
14565            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
14566            && f16g_proj_ok(m.up_exps.qtype, n_embd)
14567            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
14568        if use_mma || f16g {
14569            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
14570            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
14571            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
14572            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
14573            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
14574            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
14575            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
14576            let y_down = if f16g {
14577                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
14578                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
14579                // permute at the very end back to pair-id order for the scatter.
14580                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
14581                let csr_tok_d = e.htod_i32(&csr_tok)?;
14582                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
14583                let g_csr = e.moe_f16_grouped(
14584                    &dev.ptr_row,
14585                    0,
14586                    n_expert,
14587                    &exi,
14588                    &ex_off,
14589                    &exo,
14590                    &z_f16,
14591                    &z_s,
14592                    n_embd,
14593                    n_ff_exp,
14594                    n_active,
14595                    n_pairs,
14596                    m.gate_exps.qtype,
14597                    rbg_d,
14598                )?;
14599                let u_csr = e.moe_f16_grouped(
14600                    &dev.ptr_row,
14601                    1,
14602                    n_expert,
14603                    &exi,
14604                    &ex_off,
14605                    &exo,
14606                    &z_f16,
14607                    &z_s,
14608                    n_embd,
14609                    n_ff_exp,
14610                    n_active,
14611                    n_pairs,
14612                    m.up_exps.qtype,
14613                    rbu_d,
14614                )?;
14615                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
14616                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
14617                let d_csr = e.moe_f16_grouped(
14618                    &dev.ptr_row,
14619                    2,
14620                    n_expert,
14621                    &exi,
14622                    &ex_off,
14623                    &exo,
14624                    &a_f16,
14625                    &a_s,
14626                    n_ff_exp,
14627                    n_embd,
14628                    n_active,
14629                    n_pairs,
14630                    m.down_exps.qtype,
14631                    m.down_exps.row_bytes,
14632                )?;
14633                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
14634            } else {
14635                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
14636                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
14637                let gate = e.mmq_iq_experts(
14638                    &dev.ptr_row,
14639                    0,
14640                    n_expert,
14641                    &exi,
14642                    &exo,
14643                    &exp_d,
14644                    &pt,
14645                    &z_scr,
14646                    n_embd,
14647                    n_ff_exp,
14648                    n_active,
14649                    n_pairs,
14650                    t,
14651                    m.gate_exps.qtype,
14652                    rbg_d,
14653                )?;
14654                let up = e.mmq_iq_experts(
14655                    &dev.ptr_row,
14656                    1,
14657                    n_expert,
14658                    &exi,
14659                    &exo,
14660                    &exp_d,
14661                    &pt,
14662                    &z_scr,
14663                    n_embd,
14664                    n_ff_exp,
14665                    n_active,
14666                    n_pairs,
14667                    t,
14668                    m.up_exps.qtype,
14669                    rbu_d,
14670                )?;
14671                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
14672                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
14673                // registers and writes ONLY the quantized scratch — the two-pass chain
14674                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
14675                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
14676                let a_scr = if crate::moe_fuse_actq_on() {
14677                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
14678                } else {
14679                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
14680                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
14681                };
14682                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
14683                let pself = e.htod_i32(&pair_self)?;
14684                e.mmq_iq_experts(
14685                    &dev.ptr_row,
14686                    2,
14687                    n_expert,
14688                    &exi,
14689                    &exo,
14690                    &exp_d,
14691                    &pself,
14692                    &a_scr,
14693                    n_ff_exp,
14694                    n_embd,
14695                    n_active,
14696                    n_pairs,
14697                    n_pairs,
14698                    m.down_exps.qtype,
14699                    m.down_exps.row_bytes,
14700                )?
14701            };
14702            let mut moe_out = e.uninit(t * n_embd)?;
14703            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
14704            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
14705                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
14706            {
14707                let n_ff_sh = gate_shexp.out_features();
14708                let sg_gate = e.matmul(gate_shexp, z, t)?;
14709                let sg_up = e.matmul(up_shexp, z, t)?;
14710                let mut sa = e.uninit(t * n_ff_sh)?;
14711                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
14712                let sh = e.matmul(down_shexp, &sa, t)?;
14713                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
14714                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
14715                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
14716                // i.e. the one real prefill actually takes on a resident-expert MoE model,
14717                // so the concat-prime isolation fix has to land here as well.
14718                let g = match &m.gate_inp_shexp {
14719                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
14720                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
14721                    }
14722                    Some(gate_inp_shexp) => {
14723                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
14724                        let mut g = e.uninit(t)?;
14725                        e.sigmoid(&gs, &mut g, t)?;
14726                        g
14727                    }
14728                    None => e.htod(&vec![1.0f32; t])?,
14729                };
14730                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
14731            }
14732            return Ok(moe_out);
14733        }
14734
14735        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
14736        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
14737        let dec = std::env::var("MEMRA_MOE_DEC")
14738            .map(|v| v != "0")
14739            .unwrap_or(true);
14740        let matvec = |proj,
14741                      exi: &_,
14742                      exo: &_,
14743                      exp_d: &_,
14744                      pt: &_,
14745                      aq: &_,
14746                      ad: &_,
14747                      inf,
14748                      outf,
14749                      qtype,
14750                      rb|
14751         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14752            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
14753            let dec = dec && q8_expert_dec_supported(qtype);
14754            if dec {
14755                e.moe_pairs_matvec_q8_dec(
14756                    &dev.ptr_row,
14757                    proj,
14758                    exi,
14759                    exo,
14760                    exp_d,
14761                    pt,
14762                    aq,
14763                    ad,
14764                    inf,
14765                    outf,
14766                    n_expert,
14767                    n_active,
14768                    n_pairs,
14769                    qtype,
14770                    rb,
14771                )
14772            } else {
14773                e.moe_pairs_matvec_q8_em(
14774                    &dev.ptr_row,
14775                    proj,
14776                    exi,
14777                    exo,
14778                    exp_d,
14779                    pt,
14780                    aq,
14781                    ad,
14782                    inf,
14783                    outf,
14784                    n_expert,
14785                    n_active,
14786                    n_pairs,
14787                    qtype,
14788                    rb,
14789                )
14790            }
14791        };
14792        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14793        let gate = matvec(
14794            0,
14795            &exi,
14796            &exo,
14797            &exp_d,
14798            &pt,
14799            &zq,
14800            &zd,
14801            n_embd,
14802            n_ff_exp,
14803            m.gate_exps.qtype,
14804            rbg_d,
14805        )?;
14806        let up = matvec(
14807            1,
14808            &exi,
14809            &exo,
14810            &exp_d,
14811            &pt,
14812            &zq,
14813            &zd,
14814            n_embd,
14815            n_ff_exp,
14816            m.up_exps.qtype,
14817            rbu_d,
14818        )?;
14819        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
14820        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
14821        // down consumes PAIR-major activation rows: pair_tok = identity.
14822        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
14823        let pself = e.htod_i32(&pair_self)?;
14824        let y_down = matvec(
14825            2,
14826            &exi,
14827            &exo,
14828            &exp_d,
14829            &pself,
14830            &aq2,
14831            &ad2,
14832            n_ff_exp,
14833            n_embd,
14834            m.down_exps.qtype,
14835            m.down_exps.row_bytes,
14836        )?;
14837        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
14838        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
14839
14840        // SHARED EXPERT epilogue — same as the other paths.
14841        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
14842        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
14843        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
14844            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
14845        {
14846            let n_ff_sh = gate_shexp.out_features();
14847            // These decode-exact forms are required by the new Step resident arm. Keep the
14848            // established grouped shared-expert program for every other architecture: widening
14849            // this to Gemma changed its speculative acceptance despite green argmax gates.
14850            let step_exact = true;
14851            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
14852            let (sg_gate, sg_up) = if step_exact && t == 1 {
14853                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, None)?
14854            } else if verify_t {
14855                let mut fused = None;
14856                if crate::spec::spec_fused_t()
14857                    && (2..=4).contains(&t)
14858                    && e.uses_q8_1_fast(gate_shexp)
14859                    && e.uses_q8_1_fast(up_shexp)
14860                {
14861                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14862                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
14863                }
14864                match fused {
14865                    Some(pair) => pair,
14866                    None => (
14867                        e.matmul_decode_exact(gate_shexp, z, t)?,
14868                        e.matmul_decode_exact(up_shexp, z, t)?,
14869                    ),
14870                }
14871            } else {
14872                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
14873            };
14874            let mut sa = e.uninit(t * n_ff_sh)?;
14875            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
14876            let sh = if verify_t {
14877                e.matmul_decode_exact(down_shexp, &sa, t)?
14878            } else {
14879                e.matmul(down_shexp, &sa, t)?
14880            };
14881            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
14882            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
14883            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
14884            // dispatch choice cannot change bits.
14885            let g = match &m.gate_inp_shexp {
14886                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
14887                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
14888                }
14889                Some(gate_inp_shexp) => {
14890                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
14891                    let mut g = e.uninit(t)?;
14892                    e.sigmoid(&gs, &mut g, t)?;
14893                    g
14894                }
14895                None => e.htod(&vec![1.0f32; t])?,
14896            };
14897            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
14898        }
14899        Ok(moe_out)
14900    }
14901
14902    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
14903    #[allow(clippy::too_many_arguments)]
14904    #[allow(clippy::too_many_arguments)]
14905    fn moe_ffn_dev(
14906        e: &Engine,
14907        m: &MoeWeights,
14908        z: &CudaSlice<f32>,
14909        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
14910        logits: &CudaSlice<f32>,
14911        t: usize,
14912        cfg: &ModelConfig,
14913        il: u16,
14914        max_block: usize,
14915    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14916        crate::moe_rp_refuse(m.dev_exps.as_ref().is_some_and(|d| d.rp), "moe_ffn_dev")?; // memra#147: no split-plane arm here
14917        let moe = cfg.moe.as_ref().unwrap();
14918        let n_embd = cfg.n_embd as usize;
14919        let n_expert = moe.expert_count as usize;
14920        let n_used = moe.expert_used_count as usize;
14921        let n_ff_exp = moe.expert_ff_length as usize;
14922        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
14923        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
14924        // clamped layers; assert both so a future caller that skips the gate fails loudly.
14925        debug_assert!(
14926            cfg.sigmoid_router().is_none(),
14927            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
14928        );
14929        debug_assert!(
14930            !cfg.swiglu_clamped_at(il as u32),
14931            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
14932        );
14933
14934        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
14935        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
14936        // MEMRA_MOE_SEL_DUMP: this arm's selection never returns to the host (softmax
14937        // device-routed decode/verify, MEMRA_MOE_DEV default ON for any fully-resident
14938        // non-sigmoid MoE) — refuse an armed dump by name rather than silently dropping
14939        // every record on it, the same law the other device-only walks hold.
14940        crate::moe_sel_dump::refuse_device_only(
14941            "the softmax device-routed decode arm (moe_ffn_dev)",
14942        )?;
14943        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
14944        // skipped entirely for macro-free experts (every k-quant GGUF).
14945        if m.has_macros {
14946            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
14947        }
14948
14949        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
14950        let mut moe_out = e.uninit(t * n_embd)?;
14951
14952        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
14953        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
14954        if let Some(dev) = m.dev_exps.as_ref() {
14955            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
14956            // the combined stride; up's base is offset in the ptr table. Down unchanged.
14957            let (rbg_d, rbu_d) = if dev.gu_il {
14958                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14959                (sxx, sxx)
14960            } else {
14961                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14962            };
14963            let q8 = moe_q8_enabled_for_model(cfg, m);
14964            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
14965            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
14966            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
14967            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
14968            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
14969            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
14970            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
14971            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
14972            let rows_arm = q8
14973                && t > 1
14974                && crate::spec::spec_m2()
14975                && n_ff_exp == 512
14976                && n_used <= 8
14977                && std::env::var("MEMRA_MOE_DEVQ8_GU")
14978                    .map(|v| v.is_empty() || v == "v")
14979                    .unwrap_or(true)
14980                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
14981                    .map(|v| v.is_empty() || v == "w8h2v")
14982                    .unwrap_or(true);
14983            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
14984            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
14985            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
14986            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
14987            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
14988            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
14989            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
14990            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
14991            let csr_mode = std::env::var("MEMRA_MOE_CSR")
14992                .ok()
14993                .and_then(|v| v.parse::<i32>().ok())
14994                .unwrap_or(1);
14995            // NVFP4 admission REVERTED 2026-08-21 (lane/samplat, decode-batch-gate2 find):
14996            // the csr_nvfp4 kernel drifts last-ULP vs the rows program (11041/32768 ACT
14997            // elements at t=8) and the drift is BATCH-COMPOSITION-DEPENDENT — gate2 (B=8 vs
14998            // isolated) FAILED on the ornith15 artifact, the one-numeric-program law's batch
14999            // axis. Three chain-pinning attempts did not close it (receipts,
15000            // research/samplat-20260821/); a source-verbatim per-pair helper form IS
15001            // bit-identical but loses the dedup win (-3% vs rows). NVFP4 stays on the rows
15002            // twins until a cached form passes gate2 + the =2 byte-compare at t=8. The
15003            // increment-1 qualification hole: =2 ran across run-spec (solo verify shapes),
15004            // never decode-batch-gate at B=8 on the MoE model itself.
15005            // MEMRA_MOE_CSR_NVFP4=1 (lane/orndecode, DIAGNOSTIC PROBE ONLY): re-admits NVFP4
15006            // to the CSR arm and widens it to the exact-16 decode widths, so gate2 B=12/16 +
15007            // the =2 byte-compare can re-adjudicate the cached form at the widths where the
15008            // serial dev loop hurts most (B=16 tick: 1280 launches/step). The v0.100.1
15009            // de-admission verdict above stands until those gates are GREEN on the MoE
15010            // artifact; this door must never default on.
15011            let csr_nvfp4_probe = std::env::var("MEMRA_MOE_CSR_NVFP4").as_deref() == Ok("1");
15012            let csr_qt = |qt: i32| {
15013                qt == crate::QT_IQ4_XS
15014                    || qt == crate::QT_IQ3_S
15015                    || (csr_nvfp4_probe && qt == crate::QT_NVFP4)
15016            };
15017            let csr_t_max = if csr_nvfp4_probe { MOE_DEV_MAX_T } else { 10 };
15018            let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
15019            let csr_arm = rows_arm
15020                && csr_mode > 0
15021                && t <= csr_t_max
15022                && csr_uniform
15023                && csr_qt(m.gate_exps.qtype)
15024                && csr_qt(m.up_exps.qtype)
15025                && csr_qt(m.down_exps.qtype);
15026            if csr_arm {
15027                if csr_mode == 2 {
15028                    static ENGAGED: std::sync::Once = std::sync::Once::new();
15029                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
15030                }
15031                let n_pairs = t * n_used;
15032                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
15033                let act = e.moe_gate_up_silu8_dev_q8_csr(
15034                    &dev.ptr_row,
15035                    &sel_d,
15036                    &zq,
15037                    &zd,
15038                    n_pairs,
15039                    n_embd,
15040                    n_ff_exp,
15041                    n_used,
15042                    n_expert,
15043                    m.gate_exps.qtype,
15044                    m.up_exps.qtype,
15045                    rbg_d,
15046                    rbu_d,
15047                )?;
15048                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
15049                // down stays on the _rows twin — BOTH CSR down variants measured negative
15050                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
15051                // 16-group rows have too little decode to amortize any dedup structure.
15052                e.moe_down8_fma_dev_q8_rows(
15053                    &dev.ptr_row,
15054                    &sel_d,
15055                    &w_d,
15056                    &aq2,
15057                    &ad2,
15058                    &mut moe_out,
15059                    t,
15060                    n_ff_exp,
15061                    n_embd,
15062                    n_used,
15063                    n_expert,
15064                    m.down_exps.qtype,
15065                    m.down_exps.row_bytes,
15066                )?;
15067                if csr_mode == 2 {
15068                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
15069                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
15070                        &dev.ptr_row,
15071                        &sel_d,
15072                        &zq,
15073                        &zd,
15074                        t,
15075                        n_embd,
15076                        n_ff_exp,
15077                        n_used,
15078                        n_expert,
15079                        m.gate_exps.qtype,
15080                        m.up_exps.qtype,
15081                        rbg_d,
15082                        rbu_d,
15083                        &m.dev_macros,
15084                    )?;
15085                    let mut out_r = e.uninit(t * n_embd)?;
15086                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
15087                    e.moe_down8_fma_dev_q8_rows(
15088                        &dev.ptr_row,
15089                        &sel_d,
15090                        &w_d,
15091                        &aq2r,
15092                        &ad2r,
15093                        &mut out_r,
15094                        t,
15095                        n_ff_exp,
15096                        n_embd,
15097                        n_used,
15098                        n_expert,
15099                        m.down_exps.qtype,
15100                        m.down_exps.row_bytes,
15101                    )?;
15102                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
15103                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
15104                    let ba = a1
15105                        .iter()
15106                        .zip(&a2)
15107                        .filter(|(x, y)| x.to_bits() != y.to_bits())
15108                        .count();
15109                    let bo = o1
15110                        .iter()
15111                        .zip(&o2)
15112                        .filter(|(x, y)| x.to_bits() != y.to_bits())
15113                        .count();
15114                    if ba + bo > 0 {
15115                        eprintln!(
15116                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
15117                            a1.len(),
15118                            o1.len()
15119                        );
15120                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
15121                        let sel_h = e.dtoh_i32(&sel_d)?;
15122                        let mut shown = 0;
15123                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
15124                            if x.to_bits() != y.to_bits() && shown < 4 {
15125                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
15126                                let ex = sel_h[p];
15127                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
15128                                eprintln!(
15129                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
15130                                );
15131                                shown += 1;
15132                            }
15133                        }
15134                        std::process::exit(3);
15135                    }
15136                }
15137            } else if rows_arm {
15138                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
15139                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
15140                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
15141                    use std::sync::atomic::{AtomicU64, Ordering};
15142                    static PAIRS: AtomicU64 = AtomicU64::new(0);
15143                    static UNIQ: AtomicU64 = AtomicU64::new(0);
15144                    static CALLS: AtomicU64 = AtomicU64::new(0);
15145                    let sel_h = e.dtoh_i32(&sel_d)?;
15146                    let mut u: Vec<i32> = sel_h.clone();
15147                    u.sort_unstable();
15148                    u.dedup();
15149                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
15150                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
15151                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
15152                    if c.is_multiple_of(480) {
15153                        let p = PAIRS.load(Ordering::Relaxed);
15154                        let q = UNIQ.load(Ordering::Relaxed);
15155                        eprintln!(
15156                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
15157                            q as f64 / p as f64
15158                        );
15159                    }
15160                }
15161                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
15162                let act = e.moe_gate_up_silu8_dev_q8_rows(
15163                    &dev.ptr_row,
15164                    &sel_d,
15165                    &zq,
15166                    &zd,
15167                    t,
15168                    n_embd,
15169                    n_ff_exp,
15170                    n_used,
15171                    n_expert,
15172                    m.gate_exps.qtype,
15173                    m.up_exps.qtype,
15174                    rbg_d,
15175                    rbu_d,
15176                    &m.dev_macros,
15177                )?;
15178                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
15179                e.moe_down8_fma_dev_q8_rows(
15180                    &dev.ptr_row,
15181                    &sel_d,
15182                    &w_d,
15183                    &aq2,
15184                    &ad2,
15185                    &mut moe_out,
15186                    t,
15187                    n_ff_exp,
15188                    n_embd,
15189                    n_used,
15190                    n_expert,
15191                    m.down_exps.qtype,
15192                    m.down_exps.row_bytes,
15193                )?;
15194            } else {
15195                for tok in 0..t {
15196                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
15197                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
15198                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
15199                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
15200                    if q8 {
15201                        let (zq, zd) = match (t, zq8) {
15202                            (1, Some((q, d))) => (q.clone(), d.clone()),
15203                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
15204                        };
15205                        let act = e.moe_gate_up_silu8_dev_q8(
15206                            &dev.ptr_row,
15207                            &selt,
15208                            &zq,
15209                            &zd,
15210                            n_embd,
15211                            n_ff_exp,
15212                            n_used,
15213                            n_expert,
15214                            m.gate_exps.qtype,
15215                            m.up_exps.qtype,
15216                            rbg_d,
15217                            rbu_d,
15218                            &m.dev_macros,
15219                        )?;
15220                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
15221                        e.moe_down8_fma_dev_q8(
15222                            &dev.ptr_row,
15223                            &selt,
15224                            &wt,
15225                            &aq2,
15226                            &ad2,
15227                            &mut dst,
15228                            n_ff_exp,
15229                            n_embd,
15230                            n_used,
15231                            n_expert,
15232                            m.down_exps.qtype,
15233                            m.down_exps.row_bytes,
15234                        )?;
15235                    } else {
15236                        let act = e.moe_gate_up_silu8_dev(
15237                            &dev.ptr_row,
15238                            &selt,
15239                            &zt,
15240                            n_embd,
15241                            n_ff_exp,
15242                            n_used,
15243                            n_expert,
15244                            m.gate_exps.qtype,
15245                            m.up_exps.qtype,
15246                            rbg_d,
15247                            rbu_d,
15248                            &m.dev_macros,
15249                        )?;
15250                        e.moe_down8_fma_dev(
15251                            &dev.ptr_row,
15252                            &selt,
15253                            &wt,
15254                            &act,
15255                            &mut dst,
15256                            n_ff_exp,
15257                            n_embd,
15258                            n_used,
15259                            n_expert,
15260                            m.down_exps.qtype,
15261                            m.down_exps.row_bytes,
15262                        )?;
15263                    }
15264                }
15265            }
15266        } else {
15267            // Launch under the cache lock: the row borrow lives as long as the closure, and the
15268            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
15269            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
15270            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
15271            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
15272            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
15273            let q8 = moe_q8_enabled_for_model(cfg, m);
15274            e.with_moe_cache(max_block, |c, eng| {
15275                let row = c
15276                    .layer_dev_row(il, n_expert, eng)?
15277                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
15278                for tok in 0..t {
15279                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
15280                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
15281                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
15282                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
15283                    if q8 {
15284                        let (zq, zd) = match (t, zq8) {
15285                            (1, Some((q, d))) => (q.clone(), d.clone()),
15286                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
15287                        };
15288                        let act = eng.moe_gate_up_silu8_dev_q8(
15289                            row,
15290                            &selt,
15291                            &zq,
15292                            &zd,
15293                            n_embd,
15294                            n_ff_exp,
15295                            n_used,
15296                            n_expert,
15297                            m.gate_exps.qtype,
15298                            m.up_exps.qtype,
15299                            m.gate_exps.row_bytes,
15300                            m.up_exps.row_bytes,
15301                            &m.dev_macros,
15302                        )?;
15303                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
15304                        eng.moe_down8_fma_dev_q8(
15305                            row,
15306                            &selt,
15307                            &wt,
15308                            &aq2,
15309                            &ad2,
15310                            &mut dst,
15311                            n_ff_exp,
15312                            n_embd,
15313                            n_used,
15314                            n_expert,
15315                            m.down_exps.qtype,
15316                            m.down_exps.row_bytes,
15317                        )?;
15318                    } else {
15319                        let act = eng.moe_gate_up_silu8_dev(
15320                            row,
15321                            &selt,
15322                            &zt,
15323                            n_embd,
15324                            n_ff_exp,
15325                            n_used,
15326                            n_expert,
15327                            m.gate_exps.qtype,
15328                            m.up_exps.qtype,
15329                            m.gate_exps.row_bytes,
15330                            m.up_exps.row_bytes,
15331                            &m.dev_macros,
15332                        )?;
15333                        eng.moe_down8_fma_dev(
15334                            row,
15335                            &selt,
15336                            &wt,
15337                            &act,
15338                            &mut dst,
15339                            n_ff_exp,
15340                            n_embd,
15341                            n_used,
15342                            n_expert,
15343                            m.down_exps.qtype,
15344                            m.down_exps.row_bytes,
15345                        )?;
15346                    }
15347                }
15348                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
15349                c.hits += (t * 3 * n_used) as u64;
15350                Ok(())
15351            })?;
15352        }
15353
15354        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
15355        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
15356        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
15357        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
15358        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
15359            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15360        {
15361            let n_ff_sh = gate_shexp.out_features();
15362            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
15363            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
15364            let verify_t = t > 1 && t < PRIME_MIN_T;
15365            let (sg_gate, sg_up) = if t == 1 {
15366                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
15367            } else if verify_t {
15368                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
15369                // rides one shared quantize + one fused2 batched launch instead of two
15370                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
15371                let mut fused = None;
15372                if crate::spec::spec_fused_t()
15373                    && (2..=4).contains(&t)
15374                    && e.uses_q8_1_fast(gate_shexp)
15375                    && e.uses_q8_1_fast(up_shexp)
15376                {
15377                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
15378                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
15379                }
15380                match fused {
15381                    Some(pair) => pair,
15382                    None => (
15383                        e.matmul_decode_exact(gate_shexp, z, t)?,
15384                        e.matmul_decode_exact(up_shexp, z, t)?,
15385                    ),
15386                }
15387            } else {
15388                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
15389            };
15390            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
15391            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
15392            let sh = if verify_t {
15393                e.matmul_decode_exact(down_shexp, &sa, t)?
15394            } else {
15395                e.matmul(down_shexp, &sa, t)?
15396            };
15397            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
15398            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
15399            // between the two arms; prefill keeps the batched cuBLASLt linear).
15400            let g = match &m.gate_inp_shexp {
15401                Some(gate_inp_shexp) => {
15402                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
15403                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
15404                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
15405                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15406                    } else {
15407                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
15408                        let mut g = e.uninit(t)?;
15409                        e.sigmoid(&gs, &mut g, t)?;
15410                        g
15411                    }
15412                }
15413                None => e.htod(&vec![1.0f32; t])?,
15414            };
15415            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
15416        }
15417
15418        Ok(moe_out)
15419    }
15420
15421    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
15422    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
15423    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
15424    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
15425    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
15426    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
15427    /// the collected raw pointers cannot move between collection and launch (single-threaded
15428    /// decode; the lock is held only for collection, launches are stream-ordered after any
15429    /// prior same-stream staging writes).
15430    #[allow(clippy::too_many_arguments)]
15431    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
15432    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
15433    #[allow(clippy::too_many_arguments)]
15434    fn moe_gdec_token_q8(
15435        e: &Engine,
15436        m: &MoeWeights,
15437        il: u16,
15438        max_block: usize,
15439        zq: &CudaSlice<i8>,
15440        zd: &CudaSlice<f32>,
15441        sel: &[u32],
15442        w: &[f32],
15443        moe_out: &mut CudaSlice<f32>,
15444        tok: usize,
15445        n_embd: usize,
15446        n_ff_exp: usize,
15447        n_used: usize,
15448    ) -> Result<bool, Box<dyn std::error::Error>> {
15449        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
15450        use cudarc::driver::DevicePtr;
15451        let ptrs = e.with_moe_cache(max_block, |c, eng| {
15452            let mut g = [0u64; 8];
15453            let mut u = [0u64; 8];
15454            let mut d = [0u64; 8];
15455            for (j, &ex) in sel.iter().enumerate() {
15456                let ex = ex as u16;
15457                let (Some(sg), Some(su), Some(sd)) = (
15458                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
15459                    c.resident(BlockId::new(il, PROJ_UP, ex)),
15460                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
15461                ) else {
15462                    return Ok(None);
15463                };
15464                let __s = eng.stream();
15465                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
15466                let (pu, _e1) = c.slot(su).device_ptr(&__s);
15467                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
15468                g[j] = pg;
15469                u[j] = pu;
15470                d[j] = pd;
15471            }
15472            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
15473                for &ex in sel {
15474                    let ex = ex as u16;
15475                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
15476                        c.note_profile_hit(BlockId::new(il, proj, ex));
15477                    }
15478                }
15479            }
15480            c.hits += (3 * n_used) as u64;
15481            Ok(Some((g, u, d)))
15482        })?;
15483        let Some((g, u, d)) = ptrs else {
15484            return Ok(false);
15485        };
15486        let mut wv = [0f32; 8];
15487        wv[..n_used].copy_from_slice(w);
15488        let act = e.moe_gate_up_silu8_q8(
15489            crate::WPtr8(g),
15490            crate::WPtr8(u),
15491            zq,
15492            zd,
15493            n_embd,
15494            n_ff_exp,
15495            n_used,
15496            m.gate_exps.qtype,
15497            m.up_exps.qtype,
15498            m.gate_exps.row_bytes,
15499            m.up_exps.row_bytes,
15500        )?;
15501        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
15502        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
15503        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
15504        e.moe_down8_fma_q8(
15505            crate::WPtr8(d),
15506            crate::F32x8(wv),
15507            &aq2,
15508            &ad2,
15509            &mut dst,
15510            n_ff_exp,
15511            n_embd,
15512            n_used,
15513            m.down_exps.qtype,
15514            m.down_exps.row_bytes,
15515        )?;
15516        Ok(true)
15517    }
15518
15519    /// glm5_next's fused MoE epilogue for ONE token-layer: the sigmoid router's already-selected
15520    /// `(sel, w)`, the PRE-clamped SwiGLU and the per-expert NVFP4 macro fold, in one launch
15521    /// pair. Returns `false` when the cache cannot hold this token's `3*n_used` blocks at once,
15522    /// in which case the caller must run the sequential loop (which zeroes its own row).
15523    ///
15524    /// WHY IT DOES NOT NEED A RESIDENT LAYER, unlike `moe_gdec_token_q8`. gdec collects pointers
15525    /// from blocks that are ALREADY resident and bails on the first miss, because a miss would
15526    /// mean an admission that could move a slot under the pointers it has already taken. This arm
15527    /// inverts the order: it ADMITS all `3*n_used` blocks first, through the same
15528    /// `dispatch_source` the sequential loop calls per projection (a hit copies nothing, a miss
15529    /// runs the identical `memcpy_htod` into a slot), and only then takes the addresses, in a
15530    /// second pass, with the cache lock still held. Nothing can move between the last admission
15531    /// and the pointer read, and the kernels are issued on the compute stream immediately after —
15532    /// the same in-order guarantee the sequential loop already relies on when it dispatches
15533    /// expert j+1 after launching expert j's kernels.
15534    ///
15535    /// The slot-capacity check is the fail-closed seam: `n_slots()` is a whole-cache bound, and
15536    /// with fewer than `3*n_used` slots an admission is guaranteed to evict one of this token's
15537    /// own blocks. The second pass re-reads `resident()` for every block rather than trusting the
15538    /// dispatch's return, so an eviction the capacity check did not predict falls through loudly
15539    /// to the sequential loop instead of dereferencing a reused slot.
15540    ///
15541    /// BIT-IDENTITY CLASS. Against the sequential loop this arm is a DISPATCH-class change, not a
15542    /// provenance one: the same block bytes and the same macro scales, but the gate/up dots are
15543    /// the fused kernel's warp reduction rather than `qmatvec_expert_q8`'s per-projection one, and
15544    /// the down accumulation is `moe_down8_fma_q8`'s slot-ordered `__fmaf_rn` chain rather than 8
15545    /// separate `axpy_into` calls. Those chains are the ones the gdec family documents as
15546    /// reproducing the sequential chain exactly; `tests/glm5_moe_epilogue_gpu.rs::the_two_arms_agree`
15547    /// measures the actual bit disagreement rather than asserting the claim.
15548    #[allow(clippy::too_many_arguments)]
15549    fn moe_fused_epi_token_q8(
15550        e: &Engine,
15551        m: &MoeWeights,
15552        il: u16,
15553        max_block: usize,
15554        zq: &CudaSlice<i8>,
15555        zd: &CudaSlice<f32>,
15556        sel: &[u32],
15557        w: &[f32],
15558        moe_out: &mut CudaSlice<f32>,
15559        tok: usize,
15560        n_embd: usize,
15561        n_ff_exp: usize,
15562        n_used: usize,
15563        limit: f32,
15564    ) -> Result<bool, Box<dyn std::error::Error>> {
15565        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_DOWN, PROJ_GATE, PROJ_UP};
15566        use cudarc::driver::DevicePtr;
15567        debug_assert!(
15568            limit > 1e-6,
15569            "the fused epilogue's kernel collapses every gate to silu(0) at limit 0"
15570        );
15571        debug_assert_eq!(sel.len(), n_used);
15572        debug_assert_eq!(w.len(), n_used);
15573
15574        let ptrs = e.with_moe_cache(max_block, |c, eng| {
15575            // Fail closed: below this bound an admission MUST evict one of this token's own
15576            // blocks, so there is no pointer set that stays valid.
15577            if c.n_slots() < 3 * n_used {
15578                return Ok(None);
15579            }
15580            // PASS 1 — admit. Identical dispatch to the sequential loop's `moe_cached_gemm_q8`,
15581            // projection for projection; only the GEMM is deferred.
15582            for &ex in sel.iter() {
15583                let ex_usize = ex as usize;
15584                for (proj, exps) in [
15585                    (PROJ_GATE, &m.gate_exps),
15586                    (PROJ_UP, &m.up_exps),
15587                    (PROJ_DOWN, &m.down_exps),
15588                ] {
15589                    let id = BlockId::new(il, proj, ex as u16);
15590                    let DispatchSlot::Resident(_) =
15591                        c.dispatch_source(id, exps.expert_source(ex_usize), eng)?;
15592                }
15593            }
15594            // PASS 2 — take the fixed slot addresses, with nothing left to admit.
15595            let mut g = [0u64; 8];
15596            let mut u = [0u64; 8];
15597            let mut d = [0u64; 8];
15598            for (j, &ex) in sel.iter().enumerate() {
15599                let ex = ex as u16;
15600                let (Some(sg), Some(su), Some(sd)) = (
15601                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
15602                    c.resident(BlockId::new(il, PROJ_UP, ex)),
15603                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
15604                ) else {
15605                    return Ok(None);
15606                };
15607                let __s = eng.stream();
15608                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
15609                let (pu, _e1) = c.slot(su).device_ptr(&__s);
15610                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
15611                g[j] = pg;
15612                u[j] = pu;
15613                d[j] = pd;
15614            }
15615            Ok(Some((g, u, d)))
15616        })?;
15617        let Some((g, u, d)) = ptrs else {
15618            return Ok(false);
15619        };
15620
15621        Self::moe_fused_epi_launch(
15622            e, m, zq, zd, sel, w, g, u, d, moe_out, tok, n_embd, n_ff_exp, n_used, limit,
15623            false, // SLRU slot provenance: interleaved bytes always
15624        )?;
15625        Ok(true)
15626    }
15627
15628    /// The fused epilogue's ONLY launch path, shared by both provenances (SLRU slot addresses and
15629    /// device-resident slab base+stride). Everything that could differ semantically between them
15630    /// — the per-expert macro fold, the clamp, the kernel pair, the dispatch counter — lives here
15631    /// exactly once, so the two arms cannot drift into being different programs. The callers
15632    /// differ only in how they filled `g`/`u`/`d`.
15633    ///
15634    /// Per-expert macro scales in router slot order: gate/up ride the kernel's epilogue exactly
15635    /// where `ffn_act_lim`'s gs/us ride the unfused loop, and down folds into the routing weight
15636    /// exactly where `axpy_into`'s `w[j] * macro_scale(ex)` folds it. `macro_scale` answers 1.0
15637    /// for a macro-free bank, so a k-quant GGUF takes this path with no fold and no branch.
15638    #[allow(clippy::too_many_arguments)]
15639    fn moe_fused_epi_launch(
15640        e: &Engine,
15641        m: &MoeWeights,
15642        zq: &CudaSlice<i8>,
15643        zd: &CudaSlice<f32>,
15644        sel: &[u32],
15645        w: &[f32],
15646        g: [u64; 8],
15647        u: [u64; 8],
15648        d: [u64; 8],
15649        moe_out: &mut CudaSlice<f32>,
15650        tok: usize,
15651        n_embd: usize,
15652        n_ff_exp: usize,
15653        n_used: usize,
15654        limit: f32,
15655        rp: bool,
15656    ) -> Result<(), Box<dyn std::error::Error>> {
15657        let mut gs = [0f32; 8];
15658        let mut us = [0f32; 8];
15659        let mut wv = [0f32; 8];
15660        // memra#147: split-plane slab provenance -> the kernels take the QT_NVFP4_RP arm.
15661        let (qt_g, qt_u, qt_d) = (
15662            crate::rp_qt(rp, m.gate_exps.qtype),
15663            crate::rp_qt(rp, m.up_exps.qtype),
15664            crate::rp_qt(rp, m.down_exps.qtype),
15665        );
15666        for (j, &ex) in sel.iter().enumerate() {
15667            let ex = ex as usize;
15668            gs[j] = m.gate_exps.macro_scale(ex);
15669            us[j] = m.up_exps.macro_scale(ex);
15670            wv[j] = w[j] * m.down_exps.macro_scale(ex);
15671        }
15672        // MEMRA_B200_MATVEC_ARM occupancy arm (lane/b200-matvec-occupancy-20260902, sm_100a
15673        // only, default OFF): the warp-packed `_w4` twins are bit-identical per (o,j)/(o) to
15674        // the shipped kernels below — packing only changes which block/warp computes an
15675        // output, never the per-output arithmetic order. See `b200_matvec_arm_on` and
15676        // docs/FLAGS.md.
15677        // MEMRA_B200_GEMV_V2's MoE v2 twins are MEASURED AND NOT DISPATCHED (box receipt
15678        // 2026-09-02, gate-gemv-bench on B200 dev 0, median us over N=5):
15679        //
15680        //   moe_gate_up_preclamp8_q8   shipped 55.0   _w4 53.3   v2 54.4   -> v2 1.011x, no gain
15681        //   moe_down8_fma_q8           shipped 37.6   _w4 36.0   v2 43.2   -> v2 0.870x, REGRESSION
15682        //
15683        // Both v2 arms came back bit-identical, so this is a speed verdict, not a correctness
15684        // one. The down twin's one-block-per-row / warp-per-expert form multiplied the launch
15685        // width by 8 and got SLOWER: eight warps in a block now contend for the same L2 sectors
15686        // and the shipped kernel's single warp was already covering the latency it was accused
15687        // of exposing. The gate/up twin is inside noise of `_w4` and its own static receipt
15688        // predicted that (longest LDG burst moved 18 -> 19; the shipped `sl` unroll already had
15689        // the depth). So this call site keeps the shipped/`_w4` pair unconditionally: the v2
15690        // kernels stay in the fatbin and in `b200_matvec_bench` as measured arms, and the door
15691        // no longer overrides `MEMRA_B200_MATVEC_ARM` here. `MEMRA_B200_GEMV_V2` is a bf16-row
15692        // and kda6 door now; see docs/FLAGS.md and the lane doc's section 4.
15693        let use_w4 = crate::b200_matvec_arm_on();
15694        let act = if use_w4 {
15695            e.moe_gate_up_preclamp8_q8_w4(
15696                crate::WPtr8(g),
15697                crate::WPtr8(u),
15698                zq,
15699                zd,
15700                crate::F32x8(gs),
15701                crate::F32x8(us),
15702                limit,
15703                n_embd,
15704                n_ff_exp,
15705                n_used,
15706                qt_g,
15707                qt_u,
15708                m.gate_exps.row_bytes,
15709                m.up_exps.row_bytes,
15710            )?
15711        } else {
15712            e.moe_gate_up_preclamp8_q8(
15713                crate::WPtr8(g),
15714                crate::WPtr8(u),
15715                zq,
15716                zd,
15717                crate::F32x8(gs),
15718                crate::F32x8(us),
15719                limit,
15720                n_embd,
15721                n_ff_exp,
15722                n_used,
15723                qt_g,
15724                qt_u,
15725                m.gate_exps.row_bytes,
15726                m.up_exps.row_bytes,
15727            )?
15728        };
15729        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
15730        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
15731        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
15732        if use_w4 {
15733            e.moe_down8_fma_q8_w4(
15734                crate::WPtr8(d),
15735                crate::F32x8(wv),
15736                &aq2,
15737                &ad2,
15738                &mut dst,
15739                n_ff_exp,
15740                n_embd,
15741                n_used,
15742                qt_d,
15743                m.down_exps.row_bytes,
15744            )?;
15745        } else {
15746            e.moe_down8_fma_q8(
15747                crate::WPtr8(d),
15748                crate::F32x8(wv),
15749                &aq2,
15750                &ad2,
15751                &mut dst,
15752                n_ff_exp,
15753                n_embd,
15754                n_used,
15755                qt_d,
15756                m.down_exps.row_bytes,
15757            )?;
15758        }
15759        crate::MOE_FUSED_EPI_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
15760        Ok(())
15761    }
15762
15763    /// The verify-rows batched routed-expert program (lane/glm5-vrest): one layer-call's
15764    /// WHOLE t x n_used pair union through the fused-epilogue kernels' rows twins —
15765    /// one gate/up+preclamp launch, one pair-major activation quantize, one down+FMA
15766    /// launch — with pointers computed from the resident slab base + ex*stride (the
15767    /// sequential slab arm's exact arithmetic) and the macro folds landing exactly where
15768    /// `ffn_act_lim` / `axpy_into` land them. `moe_out` rows are FULLY overwritten.
15769    #[allow(clippy::too_many_arguments)]
15770    // allow: the parameter list mirrors its dispatch-arm caller's contract
15771    fn moe_vrows_pairs_q8(
15772        e: &Engine,
15773        m: &MoeWeights,
15774        z: &CudaSlice<f32>,
15775        sel: VrowsSel<'_>,
15776        il: u16,
15777        (pg, pu, pd): (u64, u64, u64),
15778        rp: bool,
15779        t: usize,
15780        n_embd: usize,
15781        n_ff_exp: usize,
15782        n_used: usize,
15783        limit: f32,
15784        moe_out: &mut CudaSlice<f32>,
15785    ) -> Result<(), Box<dyn std::error::Error>> {
15786        let n_pairs = t * n_used;
15787        // Door E (MEMRA_MOE_VROWS_DEDUP_ORDER, default OFF): the gate/up launch walks the pair
15788        // union EXPERT-MAJOR, reading the visit order from a FOURTH plane appended to the pointer
15789        // table (planes: gate | up | down | order). Carrying it in the existing table is what
15790        // makes the door free on the host arm — the order plane rides the single `htod_u64_into`
15791        // that was already uploading the pointers, so no new transfer and no new pool appear.
15792        // Door M (`MEMRA_MOE_VROWS_PACK`) refuses it in the launcher, so do not build the plane
15793        // when the refuted pack door is armed.
15794        let order_on = crate::moe_vrows_dedup_order_on() && !crate::moe_vrows_pack_on();
15795        let n_planes = if order_on { 4 } else { 3 };
15796        // Door W (MEMRA_GLM5_VERIFY_WS): the whole staging set — tables, token quantize,
15797        // act, pair quantize — draws from the verify workspace and recycles at the end of
15798        // the call (vws_* are alloc_uninit/plain-drop with the door off, so the OFF arm is
15799        // byte-for-byte the shipped program). Every buffer is fully overwritten before any
15800        // read by the SAME kernels (the sites' standing uninit contract).
15801        let mut ptrs_d = e.vws_uninit_u64(n_planes * n_pairs)?;
15802        let mut scl_d = e.vws_uninit(3 * n_pairs)?;
15803        // ONE launch path, TWO table provenances (the fused-epilogue arm's own discipline):
15804        // only the plane-major (gate | up | down) pointer/scale tables are built differently,
15805        // and door D's kernel evaluates the SAME terms as the host loop, so nothing downstream
15806        // can tell the arms apart. See the `moe_vrows_tables_from_sel` kernel comment for the
15807        // term-by-term bit-identity argument.
15808        match sel {
15809            VrowsSel::Host(sel_all, w_all) => {
15810                debug_assert_eq!(sel_all.len(), n_pairs);
15811                debug_assert_eq!(w_all.len(), n_pairs);
15812                let mut ptrs = vec![0u64; n_planes * n_pairs];
15813                let mut scl = vec![0f32; 3 * n_pairs];
15814                for (p, (&ex, &w)) in sel_all.iter().zip(w_all).enumerate() {
15815                    let ex = ex as usize;
15816                    ptrs[p] = pg + (ex * m.gate_exps.expert_stride) as u64;
15817                    ptrs[n_pairs + p] = pu + (ex * m.up_exps.expert_stride) as u64;
15818                    ptrs[2 * n_pairs + p] = pd + (ex * m.down_exps.expert_stride) as u64;
15819                    scl[p] = m.gate_exps.macro_scale(ex);
15820                    scl[n_pairs + p] = m.up_exps.macro_scale(ex);
15821                    // down-proj macro folds into the accumulate weight (1.0 for non-macro
15822                    // banks) — the axpy_into fold, verbatim.
15823                    scl[2 * n_pairs + p] = w * m.down_exps.macro_scale(ex);
15824                }
15825                if order_on {
15826                    // The order plane rides the SAME upload — the door adds no HtoD on this arm.
15827                    ptrs[3 * n_pairs..].copy_from_slice(&crate::vrows_expert_major_order(sel_all));
15828                    // The box receipt: the slab reads whose repeat visit this schedule places
15829                    // inside the reuse window (host arm only — see MOE_VROWS_SLAB_READS_AVOIDED).
15830                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
15831                    crate::MOE_VROWS_SLAB_READS_AVOIDED
15832                        .fetch_add(visits - distinct, std::sync::atomic::Ordering::Relaxed);
15833                }
15834                e.htod_u64_into(&ptrs, &mut ptrs_d)?;
15835                e.htod_f32_into(&scl, &mut scl_d)?;
15836                // MEMRA_MOE_VROWS_DEDUP_STAT: size the ONLY remaining byte lever on this pair
15837                // (LANE.md §1 — it already runs at ~90% of theoretical DRAM peak, so the sole
15838                // way to cut it further is reading a shared expert slab once for the rows that
15839                // share it). `1 - distinct/visits` IS that lever; measuring it costs a bitset.
15840                if crate::moe_vrows_dedup_stat_on() {
15841                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
15842                    debug_assert_eq!(visits, n_pairs as u64);
15843                    crate::MOE_VROWS_PAIR_VISITS
15844                        .fetch_add(visits, std::sync::atomic::Ordering::Relaxed);
15845                    crate::MOE_VROWS_PAIR_DISTINCT
15846                        .fetch_add(distinct, std::sync::atomic::Ordering::Relaxed);
15847                    crate::moe_vrows_dedup_report();
15848                }
15849            }
15850            VrowsSel::Dev(sel_d, selw_d) => {
15851                let macros = match (
15852                    m.gate_exps.macros.as_deref(),
15853                    m.up_exps.macros.as_deref(),
15854                    m.down_exps.macros.as_deref(),
15855                ) {
15856                    (Some(g), Some(u), Some(d)) => Some((g, u, d)),
15857                    // A partially-macro bank would need per-plane 1.0 defaults the kernel does
15858                    // not carry; the serving artifact's three planes are all present or all
15859                    // absent, so refuse rather than guess.
15860                    (None, None, None) => None,
15861                    _ => {
15862                        return Err("vrows device tables: expert macro planes are not uniform \
15863                                    across gate/up/down"
15864                            .into());
15865                    }
15866                };
15867                e.moe_vrows_tables_from_sel(
15868                    sel_d,
15869                    selw_d,
15870                    il,
15871                    macros,
15872                    (pg, pu, pd),
15873                    (
15874                        m.gate_exps.expert_stride,
15875                        m.up_exps.expert_stride,
15876                        m.down_exps.expert_stride,
15877                    ),
15878                    n_pairs,
15879                    &mut ptrs_d,
15880                    &mut scl_d,
15881                )?;
15882                if order_on {
15883                    // Door E on the device arm: one extra launch (the host arm gets the plane for
15884                    // free inside its existing upload). Bit-identical to the host's stable sort by
15885                    // (expert id, pair index) — gated directly against it in glm5_dedup_sched_gpu.
15886                    e.moe_vrows_order_from_sel(sel_d, n_pairs, &mut ptrs_d)?;
15887                }
15888                if crate::MOE_VROWS_DEV_TABLES_DISPATCHES
15889                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
15890                    == 0
15891                {
15892                    eprintln!(
15893                        "[moe-vrows-dev-tables] engaged: pointer/scale tables built on device \
15894                         from the router's own sel/w; the per-layer pinned readback and its \
15895                         cuStreamSynchronize are skipped (MEMRA_MOE_VROWS_DEV_TABLES=1)"
15896                    );
15897                }
15898            }
15899        }
15900        // Token rows quantized in one launch; per-row q8_1 bytes are position-independent
15901        // (the batched-MMVQ class), bit-gated against the per-token quantize_q8_1_view.
15902        let (mut zq, mut zd) = (
15903            e.vws_uninit_i8(t * n_embd)?,
15904            e.vws_uninit(t * (n_embd / 32))?,
15905        );
15906        e.quantize_q8_1_into(z, t, n_embd, &mut zq, &mut zd)?;
15907        // LABEL FROM THE PROVENANCE, not from the call site. `moe_vrows_pairs_q8` is also the
15908        // spec-verify walk's launcher with HOST-built tables, so a hard-coded "device" here would
15909        // be a lie on that path — and a mislabelled arm is exactly the class of mistake that cost
15910        // a box window in take 7.
15911        let vrows_arm = match sel {
15912            VrowsSel::Dev(..) => "device",
15913            VrowsSel::Host(..) => "vrows-host",
15914        };
15915        Self::trace_moe_act(e, vrows_arm, il, t, z, &zq, &zd);
15916        let act = e.moe_gate_up_preclamp8_q8_rows(
15917            &ptrs_d,
15918            &scl_d,
15919            &zq,
15920            &zd,
15921            limit,
15922            n_embd,
15923            n_ff_exp,
15924            n_used,
15925            n_pairs,
15926            crate::rp_qt(rp, m.gate_exps.qtype),
15927            crate::rp_qt(rp, m.up_exps.qtype),
15928            m.gate_exps.row_bytes,
15929            m.up_exps.row_bytes,
15930        )?;
15931        // Pair-major activation quantize: [n_pairs, n_ff] rows in one launch.
15932        let (mut aq2, mut ad2) = (
15933            e.vws_uninit_i8(n_pairs * n_ff_exp)?,
15934            e.vws_uninit(n_pairs * (n_ff_exp / 32))?,
15935        );
15936        e.quantize_q8_1_into(&act, n_pairs, n_ff_exp, &mut aq2, &mut ad2)?;
15937        e.moe_down8_fma_q8_rows(
15938            &ptrs_d,
15939            &scl_d,
15940            &aq2,
15941            &ad2,
15942            moe_out,
15943            n_ff_exp,
15944            n_embd,
15945            n_used,
15946            n_pairs,
15947            crate::rp_qt(rp, m.down_exps.qtype),
15948            m.down_exps.row_bytes,
15949        )?;
15950        Self::trace_moe_out(e, vrows_arm, il, moe_out);
15951        // Everything above is dead after the down launch (stream-ordered reuse is safe on
15952        // this engine's stream, the same guarantee the async free relies on).
15953        e.vws_recycle_u64(ptrs_d);
15954        e.vws_recycle(scl_d);
15955        e.vws_recycle_i8(zq);
15956        e.vws_recycle(zd);
15957        e.vws_recycle(act);
15958        e.vws_recycle_i8(aq2);
15959        e.vws_recycle(ad2);
15960        if crate::MOE_VROWS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
15961            eprintln!(
15962                "[glm5-vrows] verify MoE batched across rows: pairs={n_pairs} (t={t} x \
15963                 {n_used}), one gate/up+preclamp launch + one down/FMA launch per layer-call \
15964                 (rides MEMRA_GLM5_VERIFY_BATCH); arm doors: pack={} dedup_order={} \
15965                 b200_matvec={} dev_tables={} (env MEMRA_MOE_VROWS_DEV_TABLES={}) — the `_rows` \
15966                 pair has several dispatch twins and only the plain one has ever run at t=1, so a \
15967                 box log has to say which it took. `dev_tables` is THIS CALL's provenance, not the \
15968                 env: the decode-graph door owns both halves and routes the tables through the \
15969                 device build whatever `MEMRA_MOE_VROWS_DEV_TABLES` says (run 7B printed the env \
15970                 and read false while the door was building on device, which cost a box window)",
15971                crate::moe_vrows_pack_on(),
15972                crate::moe_vrows_dedup_order_on(),
15973                crate::b200_matvec_arm_on(),
15974                matches!(sel, VrowsSel::Dev(..)),
15975                crate::moe_vrows_dev_tables_on(),
15976            );
15977        }
15978        Ok(())
15979    }
15980
15981    /// GROUPED MoE PREFILL for the sigmoid-router glm5_next class (`MEMRA_MOE_GROUPED_PREFILL`,
15982    /// default ON since 2026-08-29, `=0` rollback). One call covers a whole prefill chunk's
15983    /// routed-expert FFN for one layer:
15984    ///
15985    ///   1. ROUTER: the SAME m-invariant `moe_router_logits` + `moe_route_sigmoid_cfg` host
15986    ///      oracle invocation the sequential arm makes, so selected experts and routing weights
15987    ///      are BIT-identical to the sequential arm by construction. Only the GEMM accumulation
15988    ///      order may move (the grouped GEMM is measured non-bit-stable,
15989    ///      `run_tensor_parallel_routes_nvfp4_prime_grouped`'s MEMRA_MOE_DETERM note), which is
15990    ///      why the acceptance gate is reference-band + routing-exactness, not byte identity.
15991    ///   2. TOKEN-SORT BY EXPERT: host counting sort of the (token, expert) pairs into an
15992    ///      expert-major CSR (vLLM's `moe_align_block_size` shape; same O(pairs) build the
15993    ///      softmax `moe_ffn_pairs` arm and the step37 grouped prime use).
15994    ///   3. ONE GROUPED TENSOR-CORE GEMM PER PROJECTION over the resident NVFP4 slab
15995    ///      (`moe_f16_grouped`, the sk single-kernel visitor with the NVFP4 direct tile
15996    ///      loaders; the step37 grouped-prime kernel class, 170-270 TFLOP/s on its lane's
15997    ///      sizing rows, generalized off the TP runtime to the single-device `dev_exps`
15998    ///      pointer-table provenance). Each expert's weights stream through tensor cores ONCE
15999    ///      per layer per chunk instead of once per (token, expert): at t=4096 that replaces
16000    ///      the sequential loop's 49 launches x 4096 tokens (~200k launches and ~113 MB x 4096
16001    ///      of expert VRAM re-reads per layer) with a ~15-launch chunk-wide program.
16002    ///   4. EPILOGUE: glm5_next's PRE-clamped SwiGLU `silu(min(g,l)) * clamp(u,±l)` with the
16003    ///      per-expert `weight_scale_2` macro fold: gate/up macros land BEFORE the nonlinearity
16004    ///      (`scale_rows` per CSR row; silu is nonlinear, so the fold cannot commute past it),
16005    ///      down macros fold into the scatter weight, exactly where the fused epilogue and the
16006    ///      sequential loop's `ffn_act_lim`/`axpy_into` put them.
16007    ///   5. SCATTER: permute CSR rows back to pair order, then the slot-ordered weighted
16008    ///      per-token accumulation (`moe_pairs_scatter`, the sequential-axpy accumulation
16009    ///      class). Shared expert rides the canonical clamp-aware grouped add.
16010    ///
16011    /// Returns `Ok(None)` (fail closed to the sequential arm) for every unqualified shape:
16012    /// no local resident slab, f16g door off, a projection the grouped kernel cannot walk, or
16013    /// an expert count past the sk visitor's group cap. Numeric class: f16-mirror activations
16014    /// (`moe_f16g_act` row-normalized f16), the class the softmax pairs f16g arm and the
16015    /// step37 grouped prime already serve prefill with; gated by
16016    /// `tests/glm5_moe_grouped_prefill_gpu.rs` against `memra_reference` at the fused-epilogue
16017    /// gate's tolerance class, plus the run-gen first-token argmax gate on real prompts.
16018    fn moe_ffn_grouped_prefill_sigmoid(
16019        e: &Engine,
16020        m: &MoeWeights,
16021        z: &CudaSlice<f32>,
16022        t: usize,
16023        cfg: &ModelConfig,
16024        il: u16,
16025    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16026        // Placement: the LOCAL resident slab only. The SLRU cannot serve a 4096-token chunk's
16027        // expert working set (a glm5 layer is 288 x 3 blocks against ~285 slots/layer on the
16028        // serving recipe), and a remote slab must never be dereferenced (m=1 peer reads are the
16029        // measured 34-150x class). Fail closed: the sequential arm stages as before.
16030        // Named decline, once per process (the loud-refusal law): on the 2x B200 pair the
16031        // resident decision never ran under MEMRA_ST_PINNED=1 (host-pinned store), this arm
16032        // returned None silently on every layer, and a 24k-token prompt primed through the
16033        // per-token dispatch at ~220 tok/s while the flag announce said "on". The 4-card 1M
16034        // window recorded the same "executes 0 times" shape. A boot must say why.
16035        fn decline_once(reason: &str, t: usize, il: u16) {
16036            static DECLINED: std::sync::atomic::AtomicBool =
16037                std::sync::atomic::AtomicBool::new(false);
16038            if !DECLINED.swap(true, std::sync::atomic::Ordering::Relaxed) {
16039                eprintln!(
16040                    "[moe-grouped-prefill] DECLINED t={t} il={il}: {reason} -> the sequential \
16041                     per-token dispatch serves this prime (logged once per process)"
16042                );
16043            }
16044        }
16045        let Some(dev) = m
16046            .dev_exps
16047            .as_ref()
16048            .filter(|d| moe_slab_enabled() && d.dev == e.ctx().ordinal())
16049        else {
16050            decline_once(
16051                "no LOCAL resident expert slab (dev_exps is None: the resident-experts decision \
16052                 did not select RESIDENT on this device, e.g. MEMRA_ST_PINNED=1 / \
16053                 MEMRA_MOE_RESIDENT=0 / a budget below the bank; or MEMRA_MOE_SLAB=0)",
16054                t,
16055                il,
16056            );
16057            return Ok(None);
16058        };
16059        if crate::moe_f16g_mode() == 0 {
16060            decline_once(
16061                "MEMRA_MOE_F16G=0 (the grouped f16 GEMM class is off)",
16062                t,
16063                il,
16064            );
16065            return Ok(None);
16066        }
16067        // MODE 1 IS REFUSED ON THIS WALK (2026-09-02, lane/b200-prefill-roofline). Mode 1 is
16068        // dequant-into-a-workspace + `cublasGemmGroupedBatchedEx`, and that API issues through
16069        // cuBLAS-INTERNAL streams that are NOT ordered with ours — the round-46 NaN race this
16070        // file's own header records, whose only mitigation here is a full stream sync placed
16071        // AFTER the `h2f_scaled` pass, i.e. after the unordered read has already happened.
16072        //
16073        // MEASURED, and the reason this is a refusal rather than a note. Box boot D on the 2x
16074        // B200 pair (`MEMRA_PRIME_CHUNK=4096 MEMRA_MOE_F16G=1`, everything else as the passing
16075        // boot C): the trunk went silently wrong somewhere in the 42 MoE layers — by layer 43 a
16076        // 4096-token chunk reported `n_active=14` where boots A-C showed 209-283, which is a
16077        // router seeing a destroyed residual — then the corrupt logits produced a token id near
16078        // `i32::MAX`, the embed gather panicked, the GPU worker died, the respawn died on the
16079        // same prime, and every request after it was connection refused. One door, one numeric
16080        // race, a fleet outage. glm5_next reaches this arm with 283 concurrent groups where the
16081        // families that qualified mode 1 reach it with a fraction of that, which is the most
16082        // likely reason the race is deterministic here and was not there.
16083        //
16084        // Scope, stated so this is not read as wider than it is: this refuses mode 1 for the
16085        // glm5 sigmoid grouped PREFILL only. Every other mode-1 caller is untouched and keeps
16086        // whatever evidence it already had. The replacement for this walk is the door's own
16087        // dequant-once arm, which runs on OUR stream and cannot take this race.
16088        if crate::moe_f16g_mode() == 1 {
16089            decline_once(
16090                "MEMRA_MOE_F16G=1 is REFUSED on the glm5 sigmoid grouped prefill:                  cublasGemmGroupedBatchedEx issues on cuBLAS-internal streams unordered with                  ours, and on this walk's 283-group shape that race silently destroyed the                  trunk and killed the worker (2026-09-02 boot D, research/glm5-b200-20260902/                 box/prefill/). Use MEMRA_MOE_F16G=2 (the default) or the MEMRA_B200_PRIME_V2                  dequant-once arm",
16091                t,
16092                il,
16093            );
16094            return Ok(None);
16095        }
16096        // MEMRA_MOE_GATE is the BYTE-identity oracle between sequential-class dispatches; this
16097        // arm is a different numeric class with its own reference-band gate, so it must not
16098        // shadow that comparison.
16099        if std::env::var("MEMRA_MOE_GATE").is_ok() {
16100            return Ok(None);
16101        }
16102        let moe = cfg
16103            .moe
16104            .as_ref()
16105            .ok_or("grouped sigmoid prefill requires MoE model metadata")?;
16106        let n_embd = cfg.n_embd as usize;
16107        let n_expert = moe.expert_count as usize;
16108        let n_used = moe.expert_used_count as usize;
16109        let n_ff_exp = moe.expert_ff_length as usize;
16110        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
16111            && f16g_proj_ok(m.up_exps.qtype, n_embd)
16112            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
16113        {
16114            decline_once(
16115                "an expert projection qtype/shape is outside the grouped f16 GEMM class",
16116                t,
16117                il,
16118            );
16119            return Ok(None);
16120        }
16121        // The sk visitor's direct-lane group cap (mirrors the grouped prime's guard). glm5's
16122        // 288 experts fit; a wider bank falls closed rather than erring mid-forward.
16123        if n_expert > 512 || n_used == 0 || n_used > 8 {
16124            return Ok(None);
16125        }
16126        let sigmoid = cfg
16127            .sigmoid_router()
16128            .ok_or("grouped sigmoid prefill requires the sigmoid router")?;
16129        // glm5_next carries the PRE form on every clamped layer (`clamp_exp_at`); a POST-form
16130        // arch reaching this arm is an unqualified semantic program, and the clamp-form law
16131        // says no dispatch site may pick a form by default. Err, not assert.
16132        let lim_exp = cfg.clamp_exp_at(il as u32);
16133        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
16134            return Err(
16135                "grouped sigmoid prefill is qualified for the PRE-clamped SwiGLU form only; \
16136                 a POST-clamp layer must ride the sequential arm"
16137                    .into(),
16138            );
16139        }
16140
16141        // 1. ROUTER. One selector shared with the sequential arm so changing dispatch cannot
16142        // change logits, selected expert ids, or routing weights (the routing-exactness half of
16143        // the acceptance gate holds by construction). The host readback here is the same one
16144        // the sequential arm performs; killing it is the L4 host-sync diet, not this arm.
16145        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
16146        Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
16147        let (sel_all, w_all) =
16148            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
16149        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
16150        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
16151        Self::trace_moe_input(e, il, t, n_embd, z)?;
16152
16153        let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
16154        let mut mt = std::time::Instant::now();
16155        let mut phase = |on: bool| -> f64 {
16156            if on {
16157                let _ = e.stream().synchronize();
16158                let v = mt.elapsed().as_secs_f64() * 1e3;
16159                mt = std::time::Instant::now();
16160                v
16161            } else {
16162                0.0
16163            }
16164        };
16165        let d_router = phase(mprof);
16166
16167        // 2. TOKEN-SORT BY EXPERT: expert-major CSR over the (token, expert) pairs.
16168        let n_pairs = t * n_used;
16169        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
16170            return Err("grouped sigmoid prefill geometry".into());
16171        }
16172        let mut buckets: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
16173        for (p, &s_id) in sel_all.iter().take(n_pairs).enumerate() {
16174            let s_id = s_id as usize;
16175            if s_id >= n_expert {
16176                return Err(format!("grouped prefill selection {s_id} >= {n_expert}").into());
16177            }
16178            buckets[s_id].push(p as i32);
16179        }
16180        let mut ex_ids: Vec<i32> = Vec::new();
16181        let mut ex_off: Vec<i32> = vec![0];
16182        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
16183        for (e_id, b) in buckets.iter().enumerate() {
16184            if !b.is_empty() {
16185                ex_ids.push(e_id as i32);
16186                ex_pairs.extend_from_slice(b);
16187                ex_off.push(ex_pairs.len() as i32);
16188            }
16189        }
16190        let n_active = ex_ids.len();
16191        if n_active == 0 || n_active > 512 {
16192            return Err(format!("grouped prefill n_active {n_active} outside 1..=512").into());
16193        }
16194        let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
16195
16196        // 3. MACRO PLANES. Gate/up per-CSR-row scales (before silu); down folds into the
16197        // scatter weight; `macro_scale` answers 1.0 for a macro-free bank, so the fold is
16198        // skipped rather than launched as a no-op.
16199        let wd: Vec<f32> = (0..n_pairs)
16200            .map(|p| w_all[p] * m.down_exps.macro_scale(sel_all[p] as usize))
16201            .collect();
16202
16203        // Interleaved gate/up slab strides (see moe_ffn_pairs / moe_ffn_dev).
16204        let (rbg_d, rbu_d) = if dev.gu_il {
16205            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
16206            (sxx, sxx)
16207        } else {
16208            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
16209        };
16210
16211        let exi = e.htod_i32(&ex_ids)?;
16212        let exo = e.htod_i32(&ex_off)?;
16213        let exp_d = e.htod_i32(&ex_pairs)?;
16214        let csr_tok_d = e.htod_i32(&csr_tok)?;
16215        let pw = e.htod(&wd)?;
16216
16217        // 4. GATE/UP grouped GEMMs over the bank, CSR order end-to-end.
16218        let (z16, zs) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
16219        let mut g = e.moe_f16_grouped(
16220            &dev.ptr_row,
16221            0,
16222            n_expert,
16223            &exi,
16224            &ex_off,
16225            &exo,
16226            &z16,
16227            &zs,
16228            n_embd,
16229            n_ff_exp,
16230            n_active,
16231            n_pairs,
16232            crate::rp_qt(dev.rp, m.gate_exps.qtype), // memra#147: slot-major resident slab
16233            rbg_d,
16234        )?;
16235        if m.gate_exps.macros.is_some() {
16236            let mg: Vec<f32> = ex_pairs
16237                .iter()
16238                .map(|&p| m.gate_exps.macro_scale(sel_all[p as usize] as usize))
16239                .collect();
16240            let mg_d = e.htod(&mg)?;
16241            e.scale_rows(&mut g, &mg_d, n_ff_exp, n_pairs)?;
16242        }
16243        let mut u = e.moe_f16_grouped(
16244            &dev.ptr_row,
16245            1,
16246            n_expert,
16247            &exi,
16248            &ex_off,
16249            &exo,
16250            &z16,
16251            &zs,
16252            n_embd,
16253            n_ff_exp,
16254            n_active,
16255            n_pairs,
16256            crate::rp_qt(dev.rp, m.up_exps.qtype),
16257            rbu_d,
16258        )?;
16259        if m.up_exps.macros.is_some() {
16260            let mu: Vec<f32> = ex_pairs
16261                .iter()
16262                .map(|&p| m.up_exps.macro_scale(sel_all[p as usize] as usize))
16263                .collect();
16264            let mu_d = e.htod(&mu)?;
16265            e.scale_rows(&mut u, &mu_d, n_ff_exp, n_pairs)?;
16266        }
16267
16268        // 5. EPILOGUE: glm5_next's PRE-clamped SwiGLU (the POST form was refused above); a
16269        // config with no live limit takes the plain-silu pair kernel.
16270        let act = match lim_exp {
16271            Some(SwigluClamp::Pre(limit)) => {
16272                let mut a = e.uninit(n_pairs * n_ff_exp)?;
16273                // Scales are 1.0: the per-expert macros already landed via scale_rows (an
16274                // exact *1.0 inside the kernel keeps the value chain unchanged).
16275                e.swiglu_preclamped_mul_scaled(
16276                    &g,
16277                    &u,
16278                    1.0,
16279                    1.0,
16280                    limit,
16281                    &mut a,
16282                    n_pairs * n_ff_exp,
16283                )?;
16284                a
16285            }
16286            None => e.moe_pairs_silu_mul(&g, &u, n_pairs * n_ff_exp)?,
16287            Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
16288        };
16289        let d_gemm_gu = phase(mprof);
16290
16291        // 6. DOWN grouped GEMM (CSR order), permute back to pair order, weighted scatter.
16292        let (a16, a_s) = e.moe_f16g_act(&act, None, n_ff_exp, n_pairs)?;
16293        let d_csr = e.moe_f16_grouped(
16294            &dev.ptr_row,
16295            2,
16296            n_expert,
16297            &exi,
16298            &ex_off,
16299            &exo,
16300            &a16,
16301            &a_s,
16302            n_ff_exp,
16303            n_embd,
16304            n_active,
16305            n_pairs,
16306            crate::rp_qt(dev.rp, m.down_exps.qtype),
16307            m.down_exps.row_bytes,
16308        )?;
16309        let y_pair = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
16310        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
16311        let tids: Vec<i32> = (0..n_pairs as i32).collect();
16312        let toff_d = e.htod_i32(&toff)?;
16313        let tids_d = e.htod_i32(&tids)?;
16314        // The scatter fully overwrites every (token, col): slot-ordered accumulation over the
16315        // token's n_used pairs, the sequential-axpy class.
16316        let mut moe_out = e.uninit(t * n_embd)?;
16317        e.moe_pairs_scatter(&y_pair, &pw, &toff_d, &tids_d, &mut moe_out, t, n_embd)?;
16318        let d_down = phase(mprof);
16319
16320        // 7. SHARED EXPERT: the canonical clamp-aware grouped add (reads clamp_shexp_at and
16321        // the optional shexp gate; glm5_next has a live shared expert on every MoE layer).
16322        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
16323        if mprof {
16324            let d_shared = phase(true);
16325            eprintln!(
16326                "[moe-grouped-prefill-prof] il={il} t={t} router={d_router:.1}ms \
16327                 gemm_gu={d_gemm_gu:.1}ms down_scatter={d_down:.1}ms shared={d_shared:.1}ms"
16328            );
16329        }
16330
16331        crate::MOE_GROUPED_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
16332        // Once per layer per process: the engagement receipt line (the A/B greps for it; the
16333        // both-arms flag announce lives at the dispatch site).
16334        static GPF_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16335        let layer_bit = 1u64 << (il as u64 % 64);
16336        if GPF_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit == 0 {
16337            eprintln!(
16338                "[moe-grouped-prefill] execute layer={il} tokens={t} n_active={n_active} \
16339                 provenance=resident-slab router=sigmoid-host-oracle epilogue=pre-clamped \
16340                 macro_fold=gate-up-rows+down-weight performance_claim=false \
16341                 (logged once per layer)"
16342            );
16343        }
16344        Ok(Some(moe_out))
16345    }
16346
16347    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16348    fn moe_gdec_token(
16349        e: &Engine,
16350        m: &MoeWeights,
16351        il: u16,
16352        max_block: usize,
16353        zt: &cudarc::driver::CudaView<f32>,
16354        sel: &[u32],
16355        w: &[f32],
16356        moe_out: &mut CudaSlice<f32>,
16357        tok: usize,
16358        n_embd: usize,
16359        n_ff_exp: usize,
16360        n_used: usize,
16361    ) -> Result<bool, Box<dyn std::error::Error>> {
16362        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16363        use cudarc::driver::DevicePtr;
16364        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
16365        let ptrs = e.with_moe_cache(max_block, |c, eng| {
16366            let mut g = [0u64; 8];
16367            let mut u = [0u64; 8];
16368            let mut d = [0u64; 8];
16369            for (j, &ex) in sel.iter().enumerate() {
16370                let ex = ex as u16;
16371                let (Some(sg), Some(su), Some(sd)) = (
16372                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
16373                    c.resident(BlockId::new(il, PROJ_UP, ex)),
16374                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
16375                ) else {
16376                    return Ok(None);
16377                };
16378                let __s = eng.stream();
16379                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
16380                let (pu, _e1) = c.slot(su).device_ptr(&__s);
16381                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
16382                g[j] = pg;
16383                u[j] = pu;
16384                d[j] = pd;
16385            }
16386            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
16387                for &ex in sel {
16388                    let ex = ex as u16;
16389                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
16390                        c.note_profile_hit(BlockId::new(il, proj, ex));
16391                    }
16392                }
16393            }
16394            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
16395            Ok(Some((g, u, d)))
16396        })?;
16397        let Some((g, u, d)) = ptrs else {
16398            return Ok(false);
16399        };
16400        let mut wv = [0f32; 8];
16401        wv[..n_used].copy_from_slice(w);
16402        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
16403        let act = e.moe_gate_up_silu8(
16404            crate::WPtr8(g),
16405            crate::WPtr8(u),
16406            zt,
16407            n_embd,
16408            n_ff_exp,
16409            n_used,
16410            m.gate_exps.qtype,
16411            m.up_exps.qtype,
16412            m.gate_exps.row_bytes,
16413            m.up_exps.row_bytes,
16414        )?;
16415        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
16416        e.moe_down8_fma_into(
16417            crate::WPtr8(d),
16418            crate::F32x8(wv),
16419            &act,
16420            &mut dst,
16421            n_ff_exp,
16422            n_embd,
16423            n_used,
16424            m.down_exps.qtype,
16425            m.down_exps.row_bytes,
16426        )?;
16427        Ok(true)
16428    }
16429
16430    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
16431    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
16432    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
16433    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
16434    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
16435    fn moe_cached_gemm_q8(
16436        e: &Engine,
16437        il: u16,
16438        proj: u8,
16439        ex: usize,
16440        m: &MoeWeights,
16441        max_block: usize,
16442        aq: &CudaSlice<i8>,
16443        ad: &CudaSlice<f32>,
16444    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16445        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
16446        let exps = match proj {
16447            PROJ_GATE => &m.gate_exps,
16448            PROJ_UP => &m.up_exps,
16449            _ => &m.down_exps,
16450        };
16451        let layout = exps.expert_layout(ex);
16452        let id = BlockId::new(il, proj, ex as u16);
16453        let source = exps.expert_source(ex);
16454        e.with_moe_cache(max_block, |c, eng| {
16455            let slot = c.dispatch_source(id, source, eng)?;
16456            let DispatchSlot::Resident(sl) = slot;
16457            let buf = c.slot(sl);
16458            eng.qmatvec_expert_q8(
16459                buf,
16460                0..layout.len,
16461                aq,
16462                ad,
16463                1,
16464                exps.in_f,
16465                exps.out_f,
16466                layout.qtype,
16467                layout.row_bytes,
16468            )
16469        })
16470    }
16471
16472    fn moe_cached_gemm(
16473        e: &Engine,
16474        il: u16,
16475        proj: u8,
16476        ex: usize,
16477        m: &MoeWeights,
16478        max_block: usize,
16479        x: &cudarc::driver::CudaView<f32>,
16480    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16481        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
16482        let exps = match proj {
16483            PROJ_GATE => &m.gate_exps,
16484            PROJ_UP => &m.up_exps,
16485            _ => &m.down_exps,
16486        };
16487        let layout = exps.expert_layout(ex);
16488        let id = BlockId::new(il, proj, ex as u16);
16489        let source = exps.expert_source(ex);
16490        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
16491        e.with_moe_cache(max_block, |c, eng| {
16492            let slot = c.dispatch_source(id, source, eng)?;
16493            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
16494            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
16495            let DispatchSlot::Resident(sl) = slot;
16496            let buf = c.slot(sl);
16497            m.qmatvec_view(
16498                eng,
16499                buf,
16500                0..layout.len,
16501                x,
16502                1,
16503                exps.in_f,
16504                exps.out_f,
16505                layout.qtype,
16506                layout.row_bytes,
16507            )
16508        })
16509    }
16510
16511    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
16512    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
16513    /// so the current forward's backend assignment and output remain unchanged.
16514    fn moe_profile_admit_expert(
16515        e: &Engine,
16516        il: u16,
16517        ex: usize,
16518        m: &MoeWeights,
16519        max_block: usize,
16520    ) -> Result<(), Box<dyn std::error::Error>> {
16521        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16522        e.with_moe_cache(max_block, |cache, eng| {
16523            for (proj, exps) in [
16524                (PROJ_GATE, &m.gate_exps),
16525                (PROJ_UP, &m.up_exps),
16526                (PROJ_DOWN, &m.down_exps),
16527            ] {
16528                let id = BlockId::new(il, proj, ex as u16);
16529                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
16530            }
16531            Ok(())
16532        })
16533    }
16534
16535    /// Read a projection from the immutable residency set when present; otherwise use one
16536    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
16537    #[allow(clippy::too_many_arguments)]
16538    fn moe_frozen_gemm(
16539        e: &Engine,
16540        il: u16,
16541        proj: u8,
16542        ex: usize,
16543        m: &MoeWeights,
16544        max_block: usize,
16545        x: &cudarc::driver::CudaView<f32>,
16546        scratch: &mut Option<CudaSlice<u8>>,
16547        scratch_len: usize,
16548    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16549        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
16550        let exps = match proj {
16551            PROJ_GATE => &m.gate_exps,
16552            PROJ_UP => &m.up_exps,
16553            _ => &m.down_exps,
16554        };
16555        let layout = exps.expert_layout(ex);
16556        let id = BlockId::new(il, proj, ex as u16);
16557        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
16558            let Some(slot) = cache.resident(id) else {
16559                return Ok(None);
16560            };
16561            let buf = cache.slot(slot);
16562            Ok(Some(m.qmatvec_view(
16563                eng,
16564                buf,
16565                0..layout.len,
16566                x,
16567                1,
16568                exps.in_f,
16569                exps.out_f,
16570                layout.qtype,
16571                layout.row_bytes,
16572            )?))
16573        })? {
16574            return Ok(output);
16575        }
16576        if scratch.is_none() {
16577            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
16578        }
16579        let scratch = scratch.as_mut().unwrap();
16580        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
16581        m.qmatvec_view(
16582            e,
16583            scratch,
16584            0..layout.len,
16585            x,
16586            1,
16587            exps.in_f,
16588            exps.out_f,
16589            layout.qtype,
16590            layout.row_bytes,
16591        )
16592    }
16593
16594    fn moe_prefetch_expert(
16595        e: &Engine,
16596        il: u16,
16597        ex: usize,
16598        m: &MoeWeights,
16599        max_block: usize,
16600        keep: &[crate::moe_cache::BlockId],
16601    ) -> Result<(), Box<dyn std::error::Error>> {
16602        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16603        e.with_moe_cache(max_block, |c, eng| {
16604            for (proj, exps) in [
16605                (PROJ_GATE, &m.gate_exps),
16606                (PROJ_UP, &m.up_exps),
16607                (PROJ_DOWN, &m.down_exps),
16608            ] {
16609                let id = BlockId::new(il, proj, ex as u16);
16610                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
16611            }
16612            Ok(())
16613        })
16614    }
16615
16616    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
16617    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
16618    fn moe_prefetch_disk_expert(
16619        e: &Engine,
16620        il: u16,
16621        ex: usize,
16622        m: &MoeWeights,
16623        max_block: usize,
16624        keep: &[crate::moe_cache::BlockId],
16625    ) -> Result<(), Box<dyn std::error::Error>> {
16626        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16627        e.with_moe_cache(max_block, |c, eng| {
16628            for (proj, exps) in [
16629                (PROJ_GATE, &m.gate_exps),
16630                (PROJ_UP, &m.up_exps),
16631                (PROJ_DOWN, &m.down_exps),
16632            ] {
16633                let source = exps.expert_source(ex);
16634                if let crate::model::ExpertSource::Disk { .. } = &source {
16635                    let id = BlockId::new(il, proj, ex as u16);
16636                    let _ = c.prefetch_source(id, source, keep, eng)?;
16637                }
16638            }
16639            Ok(())
16640        })
16641    }
16642
16643    #[inline]
16644    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
16645        let _ = m.gate_exps.prefetch_expert_pages(ex);
16646        let _ = m.up_exps.prefetch_expert_pages(ex);
16647        let _ = m.down_exps.prefetch_expert_pages(ex);
16648    }
16649}
16650
16651// ================================================================================================
16652// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
16653//
16654// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
16655// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
16656// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
16657//
16658// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
16659// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
16660// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
16661// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
16662// identical to the per-token loop regardless of expert processing order.
16663//
16664// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
16665// ================================================================================================
16666
16667impl HybridModel {
16668    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
16669    /// sequential fused q8 program over the token axis; clamped layers use the separate
16670    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
16671    #[allow(clippy::too_many_arguments)]
16672    fn moe_ffn_grouped_resident_q8(
16673        e: &Engine,
16674        m: &MoeWeights,
16675        z: &CudaSlice<f32>,
16676        t: usize,
16677        cfg: &ModelConfig,
16678        il: u16,
16679        sel_all: &[u32],
16680        w_all: &[f32],
16681        table: &CudaSlice<u64>,
16682        gu_il: bool,
16683    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16684        crate::moe_rp_refuse(
16685            m.dev_exps.as_ref().is_some_and(|d| d.rp),
16686            "moe_ffn_grouped_resident_q8",
16687        )?; // memra#147: no split-plane arm here
16688        let moe = cfg.moe.as_ref().unwrap();
16689        let n_embd = cfg.n_embd as usize;
16690        let n_expert = moe.expert_count as usize;
16691        let n_used = moe.expert_used_count as usize;
16692        let n_ff_exp = moe.expert_ff_length as usize;
16693        let n_pairs = t * n_used;
16694        debug_assert_eq!(sel_all.len(), n_pairs);
16695        debug_assert_eq!(w_all.len(), n_pairs);
16696        debug_assert!(
16697            m.gate_exps.macros.is_none()
16698                && m.up_exps.macros.is_none()
16699                && m.down_exps.macros.is_none(),
16700            "resident grouped q8 does not fold per-expert macro scales",
16701        );
16702
16703        // The rows twins run the resident sequential program verbatim on grid.z = token:
16704        // fused gate/up/SiLU per slot, batched activation quantization, then the original
16705        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
16706        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
16707        // never enter the softmax router.
16708        if !cfg.swiglu_clamped_at(il as u32) {
16709            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
16710            let sel_d = e.htod_i32(&sel)?;
16711            let w_d = e.htod(w_all)?;
16712            let (gate_row_bytes, up_row_bytes) = if gu_il {
16713                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
16714                (combined, combined)
16715            } else {
16716                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
16717            };
16718            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
16719            let act = e.moe_gate_up_silu8_dev_q8_rows(
16720                table,
16721                &sel_d,
16722                &zq,
16723                &zd,
16724                t,
16725                n_embd,
16726                n_ff_exp,
16727                n_used,
16728                n_expert,
16729                m.gate_exps.qtype,
16730                m.up_exps.qtype,
16731                gate_row_bytes,
16732                up_row_bytes,
16733                &m.dev_macros,
16734            )?;
16735            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
16736            let mut moe_out = e.uninit(t * n_embd)?;
16737            e.moe_down8_fma_dev_q8_rows_g(
16738                table,
16739                &sel_d,
16740                &w_d,
16741                &aq2,
16742                &ad2,
16743                &mut moe_out,
16744                t,
16745                n_ff_exp,
16746                n_embd,
16747                n_used,
16748                n_expert,
16749                m.down_exps.qtype,
16750                m.down_exps.row_bytes,
16751            )?;
16752
16753            if std::env::var("MEMRA_MOE_STATS").is_ok() {
16754                let mut counts = vec![0usize; n_expert];
16755                for &expert in sel_all {
16756                    counts[expert as usize] += 1;
16757                }
16758                let mut sizes: Vec<usize> =
16759                    counts.into_iter().filter(|&count| count != 0).collect();
16760                sizes.sort_unstable();
16761                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
16762                println!(
16763                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
16764                     m_e: min={} median={} mean={mean:.1} max={}",
16765                    sizes.len(),
16766                    n_expert,
16767                    sizes.first().copied().unwrap_or(0),
16768                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
16769                    sizes.last().copied().unwrap_or(0),
16770                );
16771            }
16772            return Ok(moe_out);
16773        }
16774
16775        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
16776        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
16777        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
16778        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
16779        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
16780        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
16781        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
16782
16783        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
16784        for (pair, &expert) in pair_ex.iter().enumerate() {
16785            by_expert[expert as usize].push(pair as i32);
16786        }
16787
16788        let pair_tok_d = e.htod_i32(&pair_tok)?;
16789        let pair_ex_d = e.htod_i32(&pair_ex)?;
16790        let pair_w_d = e.htod(w_all)?;
16791        let tok_off_d = e.htod_i32(&tok_off)?;
16792        let tok_ids_d = e.htod_i32(&tok_ids)?;
16793
16794        let matvec = |proj: i32,
16795                      pair_rows: &CudaSlice<i32>,
16796                      aq: &CudaSlice<i8>,
16797                      ad: &CudaSlice<f32>,
16798                      in_f: usize,
16799                      out_f: usize,
16800                      qtype: i32,
16801                      row_bytes: usize|
16802         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16803            e.moe_pairs_matvec_q8(
16804                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
16805                row_bytes,
16806            )
16807        };
16808
16809        let (gate_row_bytes, up_row_bytes) = if gu_il {
16810            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
16811            (combined, combined)
16812        } else {
16813            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
16814        };
16815        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
16816        let gate = matvec(
16817            0,
16818            &pair_tok_d,
16819            &zq,
16820            &zd,
16821            n_embd,
16822            n_ff_exp,
16823            m.gate_exps.qtype,
16824            gate_row_bytes,
16825        )?;
16826        let up = matvec(
16827            1,
16828            &pair_tok_d,
16829            &zq,
16830            &zd,
16831            n_embd,
16832            n_ff_exp,
16833            m.up_exps.qtype,
16834            up_row_bytes,
16835        )?;
16836        let mut act = e.uninit(n_pairs * n_ff_exp)?;
16837        Self::ffn_act_lim(
16838            e,
16839            cfg,
16840            &gate,
16841            &up,
16842            1.0,
16843            1.0,
16844            cfg.clamp_exp_at(il as u32),
16845            &mut act,
16846            n_pairs * n_ff_exp,
16847        )?;
16848        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
16849        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
16850        let pair_self_d = e.htod_i32(&pair_self)?;
16851        let down = matvec(
16852            2,
16853            &pair_self_d,
16854            &aq2,
16855            &ad2,
16856            n_ff_exp,
16857            n_embd,
16858            m.down_exps.qtype,
16859            m.down_exps.row_bytes,
16860        )?;
16861        let mut moe_out = e.uninit(t * n_embd)?;
16862        e.moe_pairs_scatter(
16863            &down,
16864            &pair_w_d,
16865            &tok_off_d,
16866            &tok_ids_d,
16867            &mut moe_out,
16868            t,
16869            n_embd,
16870        )?;
16871
16872        if std::env::var("MEMRA_MOE_STATS").is_ok() {
16873            let mut sizes: Vec<usize> = by_expert
16874                .iter()
16875                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
16876                .collect();
16877            sizes.sort_unstable();
16878            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
16879            println!(
16880                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
16881                 m_e: min={} median={} mean={mean:.1} max={}",
16882                sizes.len(),
16883                n_expert,
16884                sizes.first().copied().unwrap_or(0),
16885                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
16886                sizes.last().copied().unwrap_or(0),
16887            );
16888        }
16889        Ok(moe_out)
16890    }
16891
16892    /// MEMRA_SHEXP_SPLIT worker: the shared expert's gate/up/down rows split across both
16893    /// devices (dev1 idles during E3), act halves exchanged both ways, downs row-split —
16894    /// per-element/per-row programs identical, so `sh` is BIT-IDENTICAL to the single-device
16895    /// arm. Process-static workspace pinned by the gate tensor pointer; rank1 holds one-time
16896    /// row-half replicas (~10MB/layer x 42). Returns None when ineligible.
16897    #[allow(clippy::too_many_arguments)]
16898    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
16899    fn shexp_split_matvec(
16900        e: &Engine,
16901        rank1: &Engine,
16902        wg: &CudaSlice<u8>,
16903        wu: &CudaSlice<u8>,
16904        wd: &CudaSlice<u8>,
16905        z: &CudaSlice<f32>,
16906        lim: Option<SwigluClamp>,
16907        cfg: &ModelConfig,
16908        il: u16,
16909        n_embd: usize,
16910        n_ff_sh: usize,
16911    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16912        use cudarc::driver::DevicePtr;
16913        if !n_ff_sh.is_multiple_of(2) || !n_embd.is_multiple_of(2) {
16914            return Ok(None);
16915        }
16916        let hf = n_ff_sh / 2;
16917        let nd = n_embd / 2;
16918        struct Rep {
16919            wg1: CudaSlice<u8>,
16920            wu1: CudaSlice<u8>,
16921            wd1: CudaSlice<u8>,
16922        }
16923        struct SplitWs {
16924            pin_dev: usize,
16925            // e side
16926            gate0: CudaSlice<f32>,
16927            up0: CudaSlice<f32>,
16928            act: CudaSlice<f32>,
16929            sh_buf: CudaSlice<f32>,
16930            ev_z: cudarc::driver::CudaEvent,
16931            ev_act0: cudarc::driver::CudaEvent,
16932            // rank1 side
16933            z1: CudaSlice<f32>,
16934            g1: CudaSlice<f32>,
16935            u1: CudaSlice<f32>,
16936            a1h: CudaSlice<f32>,
16937            act1: CudaSlice<f32>,
16938            y1: CudaSlice<f32>,
16939            ev_act1: cudarc::driver::CudaEvent,
16940            ev_y1: cudarc::driver::CudaEvent,
16941            raw_act_e: u64,
16942            raw_sh_e: u64,
16943            raw_z1: u64,
16944            raw_a1h: u64,
16945            raw_act1: u64,
16946            raw_y1: u64,
16947        }
16948        static WS: std::sync::Mutex<Option<SplitWs>> = std::sync::Mutex::new(None);
16949        static REPS: std::sync::Mutex<Option<std::collections::HashMap<u64, Rep>>> =
16950            std::sync::Mutex::new(None);
16951        let mut guard = WS.lock().map_err(|_| "shexp split lock is poisoned")?;
16952        let mut reps_guard = REPS.lock().map_err(|_| "shexp reps lock is poisoned")?;
16953        let reps = reps_guard.get_or_insert_with(std::collections::HashMap::new);
16954        let pins = e.ctx().ordinal();
16955        if guard.as_ref().is_none_or(|w| w.pin_dev != pins) {
16956            let (gate0, up0, act, sh_buf, ev_z, ev_act0) = {
16957                let _m = e.gpu.enter_main()?;
16958                (
16959                    e.htod(&vec![0.0f32; hf])?,
16960                    e.htod(&vec![0.0f32; hf])?,
16961                    e.htod(&vec![0.0f32; n_ff_sh])?,
16962                    e.htod(&vec![0.0f32; n_embd])?,
16963                    e.ctx().new_event(None)?,
16964                    e.ctx().new_event(None)?,
16965                )
16966            };
16967            let (z1, g1, u1, a1h, act1, y1, ev_act1, ev_y1) = {
16968                let _r = rank1.gpu.enter_main()?;
16969                (
16970                    rank1.htod(&vec![0.0f32; n_embd])?,
16971                    rank1.htod(&vec![0.0f32; hf])?,
16972                    rank1.htod(&vec![0.0f32; hf])?,
16973                    rank1.htod(&vec![0.0f32; hf])?,
16974                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
16975                    rank1.htod(&vec![0.0f32; nd])?,
16976                    rank1.ctx().new_event(None)?,
16977                    rank1.ctx().new_event(None)?,
16978                )
16979            };
16980            let (raw_act_e, raw_sh_e) = {
16981                let _m = e.gpu.enter_main()?;
16982                let stream = e.stream();
16983                let (a, _g0) = act.device_ptr(&stream);
16984                let (b, _g1) = sh_buf.device_ptr(&stream);
16985                (a, b)
16986            };
16987            let (raw_z1, raw_a1h, raw_act1, raw_y1) = {
16988                let _r = rank1.gpu.enter_main()?;
16989                let rs = rank1.stream();
16990                let (a, _g0) = z1.device_ptr(&rs);
16991                let (b, _g1) = a1h.device_ptr(&rs);
16992                let (c, _g2) = act1.device_ptr(&rs);
16993                let (d, _g3) = y1.device_ptr(&rs);
16994                (a, b, c, d)
16995            };
16996            *guard = Some(SplitWs {
16997                pin_dev: pins,
16998                gate0,
16999                up0,
17000                act,
17001                sh_buf,
17002                ev_z,
17003                ev_act0,
17004                z1,
17005                g1,
17006                u1,
17007                a1h,
17008                act1,
17009                y1,
17010                ev_act1,
17011                ev_y1,
17012                raw_act_e,
17013                raw_sh_e,
17014                raw_z1,
17015                raw_a1h,
17016                raw_act1,
17017                raw_y1,
17018            });
17019        }
17020        let ws = guard.as_mut().expect("armed above");
17021        let wg_pin = {
17022            let _m = e.gpu.enter_main()?;
17023            let stream = e.stream();
17024            let (p, _g) = wg.device_ptr(&stream);
17025            p
17026        };
17027        if !reps.contains_key(&wg_pin) {
17028            // One-time per layer: upload rank1's row halves (gate/up rows [hf..], down rows [nd..]).
17029            let up = |src: &CudaSlice<u8>,
17030                      off_bytes: usize,
17031                      len: usize|
17032             -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
17033                use cudarc::driver::sys;
17034                let sptr = {
17035                    let _m = e.gpu.enter_main()?;
17036                    let stream = e.stream();
17037                    let (p, _g) = src.device_ptr(&stream);
17038                    p + off_bytes as u64
17039                };
17040                let dst = {
17041                    let _r = rank1.gpu.enter_main()?;
17042                    rank1.alloc_u8_uninit(len)?
17043                };
17044                let dptr = {
17045                    let _r = rank1.gpu.enter_main()?;
17046                    let rs = rank1.stream();
17047                    let (p, _g) = dst.device_ptr(&rs);
17048                    p
17049                };
17050                let _r = rank1.gpu.enter_main()?;
17051                let r = unsafe {
17052                    sys::cuMemcpyAsync(
17053                        dptr as sys::CUdeviceptr,
17054                        sptr as sys::CUdeviceptr,
17055                        len,
17056                        rank1.stream().cu_stream() as sys::CUstream,
17057                    )
17058                };
17059                if r != sys::CUresult::CUDA_SUCCESS {
17060                    return Err(format!("shexp split replica upload: {r:?}").into());
17061                }
17062                rank1.stream().synchronize()?;
17063                Ok(dst)
17064            };
17065            let wg1 = up(wg, hf * n_embd * 2, hf * n_embd * 2)?;
17066            let wu1 = up(wu, hf * n_embd * 2, hf * n_embd * 2)?;
17067            let wd1 = up(wd, nd * n_ff_sh * 2, nd * n_ff_sh * 2)?;
17068            reps.insert(wg_pin, Rep { wg1, wu1, wd1 });
17069        }
17070        let _ = il;
17071        // Per token, evented split flow.
17072        let raw_z = {
17073            let _m = e.gpu.enter_main()?;
17074            let stream = e.stream();
17075            let (p, _g) = z.device_ptr(&stream);
17076            ws.ev_z.record(&stream)?;
17077            p
17078        };
17079        // rank1: pull z, its dual half, its act half; push act half to e; pull e's act half.
17080        {
17081            let rep = reps.get(&wg_pin).expect("uploaded above");
17082            let _r = rank1.gpu.enter_main()?;
17083            rank1.stream().wait(&ws.ev_z)?;
17084            crate::tp::raw_copy_bytes(ws.raw_z1, raw_z, n_embd * 4, rank1)?;
17085            let SplitWs {
17086                z1, g1, u1, a1h, ..
17087            } = &mut *ws;
17088            rank1.matvec_bf16_dual_into(&rep.wg1, &rep.wu1, z1, g1, u1, n_embd, hf)?;
17089            Self::ffn_act_lim(rank1, cfg, g1, u1, 1.0, 1.0, lim, a1h, hf)?;
17090            // local place into act1[hf..] + P2P push into e's act[hf..]
17091            crate::tp::raw_copy_bytes(ws.raw_act1 + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
17092            crate::tp::raw_copy_bytes(ws.raw_act_e + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
17093            ws.ev_act1.record(&rank1.stream())?;
17094        }
17095        // e: its dual half + act half; publish it; then wait rank1's half and run down lo.
17096        {
17097            let _m = e.gpu.enter_main()?;
17098            let SplitWs {
17099                gate0, up0, act, ..
17100            } = &mut *ws;
17101            let wg_lo = wg.slice(0..hf * n_embd * 2);
17102            let wu_lo = wu.slice(0..hf * n_embd * 2);
17103            e.matvec_bf16_dual_view_into(&wg_lo, &wu_lo, z, gate0, up0, n_embd, hf)?;
17104            Self::ffn_act_lim(e, cfg, gate0, up0, 1.0, 1.0, lim, act, hf)?;
17105            ws.ev_act0.record(&e.stream())?;
17106        }
17107        // rank1: pull e's act half into act1[0..hf], run down hi, push y half.
17108        {
17109            let rep = reps.get(&wg_pin).expect("uploaded above");
17110            let _r = rank1.gpu.enter_main()?;
17111            rank1.stream().wait(&ws.ev_act0)?;
17112            crate::tp::raw_copy_bytes(ws.raw_act1, ws.raw_act_e, hf * 4, rank1)?;
17113            let SplitWs { act1, y1, .. } = &mut *ws;
17114            rank1.matvec_bf16_into(&rep.wd1, act1, y1, n_ff_sh, nd)?;
17115            crate::tp::raw_copy_bytes(ws.raw_sh_e + (nd * 4) as u64, ws.raw_y1, nd * 4, rank1)?;
17116            ws.ev_y1.record(&rank1.stream())?;
17117        }
17118        // e: down lo into sh_buf[0..nd]; join rank1's half; hand back an owned sh.
17119        {
17120            let _m = e.gpu.enter_main()?;
17121            e.stream().wait(&ws.ev_act1)?;
17122            let SplitWs { act, sh_buf, .. } = &mut *ws;
17123            let wd_lo = wd.slice(0..nd * n_ff_sh * 2);
17124            e.matvec_bf16_view_into(&wd_lo, act, sh_buf, n_ff_sh, nd)?;
17125            e.stream().wait(&ws.ev_y1)?;
17126            let mut sh = e.uninit(n_embd)?;
17127            {
17128                let mut dst = sh.slice_mut(0..n_embd);
17129                e.stream()
17130                    .memcpy_dtod(&ws.sh_buf.slice(0..n_embd), &mut dst)?;
17131            }
17132            Ok(Some(sh))
17133        }
17134    }
17135
17136    /// SHEXP OVERLAP issue (MEMRA_SHEXP_OVERLAP=1): the shared expert reads only `z`, so
17137    /// its kernels (dual matvec+SwiGLU, down) are issued on e's stream from the routes
17138    /// PREJOIN hook — they execute while the peer rank drains its sweep, filling dev0's
17139    /// join wait. The down lands in ITS OWN row (the pre-#2e split program — receipted
17140    /// bit-identical to the fused down+addscale) and `shexp_overlap_apply` adds it after
17141    /// the join with the exact add_scaled_rows expression: values unchanged.
17142    fn shexp_overlap_issue(
17143        e: &Engine,
17144        m: &MoeWeights,
17145        z: &CudaSlice<f32>,
17146        cfg: &ModelConfig,
17147        il: u16,
17148        n_embd: usize,
17149    ) -> Result<bool, Box<dyn std::error::Error>> {
17150        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
17151            return Ok(false);
17152        }
17153        let (
17154            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
17155            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
17156            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
17157        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
17158        else {
17159            return Ok(false);
17160        };
17161        let n_ff_sh = m
17162            .gate_shexp
17163            .as_ref()
17164            .expect("matched Some above")
17165            .out_features();
17166        // The dual-silu epilogue is step35's POST form only; a PRE-clamped layer declines here
17167        // and takes the unfused seam.
17168        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
17169            return Ok(false);
17170        };
17171        let mut guard = SHEXP_OV_WS
17172            .lock()
17173            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
17174        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
17175        if guard
17176            .as_ref()
17177            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
17178        {
17179            *guard = Some((
17180                pins.0,
17181                pins.1,
17182                pins.2,
17183                e.uninit(n_ff_sh)?,
17184                e.uninit(n_embd)?,
17185            ));
17186        }
17187        let (_, _, _, act, sh) = guard.as_mut().expect("armed above");
17188        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim)?;
17189        e.matvec_bf16_into(wd, act, sh, n_ff_sh, n_embd)?;
17190        drop(guard);
17191        Ok(true)
17192    }
17193
17194    /// SHEXP ON DEV1 issue (MEMRA_SHEXP_DEV1=1): the shared expert runs on rank1 — the
17195    /// idle device — with replica weights (one-time P2P upload), the SAME kernels and the
17196    /// SAME split program as the dev0 overlap (dual matvec+SwiGLU, f32acc down), so the
17197    /// values are bit-identical. z rides one 16KB P2P pull behind an e-stream event; the
17198    /// down row lands root-resident (single P2P store pass); apply waits ev_done on e.
17199    #[allow(clippy::too_many_arguments)]
17200    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
17201    fn shexp_dev1_issue(
17202        e: &Engine,
17203        rank1: &Engine,
17204        m: &MoeWeights,
17205        z: &CudaSlice<f32>,
17206        cfg: &ModelConfig,
17207        il: u16,
17208        n_embd: usize,
17209    ) -> Result<bool, Box<dyn std::error::Error>> {
17210        use cudarc::driver::DevicePtr;
17211        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
17212            return Ok(false);
17213        }
17214        let (
17215            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
17216            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
17217            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
17218        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
17219        else {
17220            return Ok(false);
17221        };
17222        let n_ff_sh = m
17223            .gate_shexp
17224            .as_ref()
17225            .expect("matched Some above")
17226            .out_features();
17227        // POST-form epilogue only (see `fused_post_limit`): a PRE-clamped layer declines
17228        // (Ok(false) = nothing issued, caller falls back per column).
17229        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
17230            return Ok(false);
17231        };
17232        // Shared scratch, geometry-keyed.
17233        let mut ws_guard = SHEXP_D1_WS
17234            .lock()
17235            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
17236        if ws_guard
17237            .as_ref()
17238            .is_none_or(|(k, ..)| *k != (n_embd, n_ff_sh))
17239        {
17240            let (act1, z1, ev_done) = {
17241                let _r1 = rank1.gpu.enter_main()?;
17242                (
17243                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
17244                    rank1.htod(&vec![0.0f32; n_embd])?,
17245                    rank1.ctx().new_event(None)?,
17246                )
17247            };
17248            let (sh_root, ev_z) = {
17249                let _main = e.gpu.enter_main()?;
17250                (e.htod(&vec![0.0f32; n_embd])?, e.ctx().new_event(None)?)
17251            };
17252            *ws_guard = Some(((n_embd, n_ff_sh), act1, z1, sh_root, ev_z, ev_done));
17253        }
17254        // Per-LAYER weight replicas (gate/up/down differ per layer): one-time P2P upload.
17255        let mut reps_guard = SHEXP_D1_REPS
17256            .lock()
17257            .map_err(|_| "shexp dev1 replica lock is poisoned")?;
17258        let reps = reps_guard.get_or_insert_with(Default::default);
17259        if !reps.contains_key(&il) {
17260            let (wg1, wu1, wd1) = {
17261                let _r1 = rank1.gpu.enter_main()?;
17262                (
17263                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
17264                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
17265                    rank1.alloc_u8_uninit(n_embd * n_ff_sh * 2)?,
17266                )
17267            };
17268            for (src, dst) in [(wg, &wg1), (wu, &wu1), (wd, &wd1)] {
17269                let s_ptr = {
17270                    let _main = e.gpu.enter_main()?;
17271                    let stream = e.stream();
17272                    let (p, _g) = src.device_ptr(&stream);
17273                    p
17274                };
17275                let d_ptr = {
17276                    let _r1 = rank1.gpu.enter_main()?;
17277                    let stream = rank1.stream();
17278                    let (p, _g) = dst.device_ptr(&stream);
17279                    p
17280                };
17281                let _r1 = rank1.gpu.enter_main()?;
17282                crate::tp::raw_copy_bytes(d_ptr, s_ptr, src.len(), rank1)?;
17283            }
17284            {
17285                let _r1 = rank1.gpu.enter_main()?;
17286                rank1.stream().synchronize()?;
17287            }
17288            reps.insert(il, (wg1, wu1, wd1));
17289        }
17290        let (wg1, wu1, wd1) = reps.get(&il).expect("armed above");
17291        let (_, act1, z1, sh_root, ev_z, ev_done) = ws_guard.as_mut().expect("armed above");
17292        // z ready on e's stream -> rank1 pulls it, runs the split shexp, pushes the down
17293        // row root-side (single store pass), rings ev_done.
17294        let (raw_z, raw_sh) = {
17295            let _main = e.gpu.enter_main()?;
17296            let stream = e.stream();
17297            let (a, _g0) = z.device_ptr(&stream);
17298            let (b, _g1) = sh_root.device_ptr(&stream);
17299            ev_z.record(&stream)?;
17300            (a, b)
17301        };
17302        {
17303            let _r1 = rank1.gpu.enter_main()?;
17304            rank1.stream().wait(ev_z)?;
17305            let raw_z1 = {
17306                let stream = rank1.stream();
17307                let (p, _g) = z1.device_ptr(&stream);
17308                p
17309            };
17310            crate::tp::raw_copy_bytes(raw_z1, raw_z, n_embd * 4, rank1)?;
17311            rank1.matvec_bf16_dual_silu_into(wg1, wu1, z1, act1, n_embd, n_ff_sh, lim)?;
17312            // down writes the ROOT-resident row over P2P via the raw-output twin of
17313            // matvec_bf16_into: reuse the view launcher with a slice view is not possible
17314            // cross-device, so launch on the raw pointer.
17315            rank1.matvec_bf16_raw_out(wd1, act1, raw_sh, n_ff_sh, n_embd)?;
17316            ev_done.record(&rank1.stream())?;
17317        }
17318        Ok(true)
17319    }
17320
17321    /// Apply the dev1 shared expert: wait ev_done on e, then the exact add_scaled_rows.
17322    fn shexp_dev1_apply(
17323        e: &Engine,
17324        output: &mut CudaSlice<f32>,
17325        n_embd: usize,
17326    ) -> Result<(), Box<dyn std::error::Error>> {
17327        let guard = SHEXP_D1_WS
17328            .lock()
17329            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
17330        let (pin, _, _, sh_root, _, ev_done) =
17331            guard.as_ref().ok_or("shexp dev1 apply without issue")?;
17332        if pin.0 != n_embd {
17333            return Err("shexp dev1 width drifted".into());
17334        }
17335        let _main = e.gpu.enter_main()?;
17336        e.stream().wait(ev_done)?;
17337        static ONES_D1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
17338            std::sync::Mutex::new(None);
17339        let mut og = ONES_D1.lock().map_err(|_| "ones lock is poisoned")?;
17340        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
17341            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
17342        }
17343        let ones = &og.as_ref().expect("armed above").1;
17344        e.add_scaled_rows(sh_root, ones, output, n_embd, 1)?;
17345        Ok(())
17346    }
17347
17348    /// MOE TAIL FUSION M1 helper: pre-arm the overlap ws + persistent ones row and
17349    /// return their RAW pointers (None when the overlap is ineligible — the caller then
17350    /// takes the split path). Mirrors shexp_overlap_issue's eligibility exactly.
17351    fn shexp_overlap_tail_ptrs(
17352        e: &Engine,
17353        m: &MoeWeights,
17354        cfg: &ModelConfig,
17355        n_embd: usize,
17356    ) -> Result<Option<(u64, u64)>, Box<dyn std::error::Error>> {
17357        use cudarc::driver::DevicePtr;
17358        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
17359            return Ok(None);
17360        }
17361        let (
17362            Some(crate::model::GpuTensor::FloatBf16 { .. }),
17363            Some(crate::model::GpuTensor::FloatBf16 { .. }),
17364            Some(crate::model::GpuTensor::FloatBf16 { .. }),
17365        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
17366        else {
17367            return Ok(None);
17368        };
17369        let n_ff_sh = m
17370            .gate_shexp
17371            .as_ref()
17372            .expect("matched Some above")
17373            .out_features();
17374        let mut guard = SHEXP_OV_WS
17375            .lock()
17376            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
17377        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
17378        if guard
17379            .as_ref()
17380            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
17381        {
17382            *guard = Some((
17383                pins.0,
17384                pins.1,
17385                pins.2,
17386                e.uninit(n_ff_sh)?,
17387                e.uninit(n_embd)?,
17388            ));
17389        }
17390        let sh_raw = {
17391            let (_, _, _, _, sh) = guard.as_ref().expect("armed above");
17392            let stream = e.stream();
17393            let (p, _g) = sh.device_ptr(&stream);
17394            p
17395        };
17396        static ONES_T3: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
17397            std::sync::Mutex::new(None);
17398        let mut og = ONES_T3.lock().map_err(|_| "ones lock is poisoned")?;
17399        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
17400            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
17401        }
17402        let ones_raw = {
17403            let stream = e.stream();
17404            let (p, _g) = og.as_ref().expect("armed above").1.device_ptr(&stream);
17405            p
17406        };
17407        Ok(Some((sh_raw, ones_raw)))
17408    }
17409
17410    /// Apply the overlapped shared expert: output[r] += sh[r] * 1.0 — the exact
17411    /// add_scaled_rows program the split path used (persistent ones row, no htod).
17412    fn shexp_overlap_apply(
17413        e: &Engine,
17414        output: &mut CudaSlice<f32>,
17415        n_embd: usize,
17416    ) -> Result<(), Box<dyn std::error::Error>> {
17417        let guard = SHEXP_OV_WS
17418            .lock()
17419            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
17420        let (_, ne, _, _, sh) = guard.as_ref().ok_or("shexp overlap apply without issue")?;
17421        if *ne != n_embd {
17422            return Err("shexp overlap width drifted".into());
17423        }
17424        static ONES_OV: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
17425            std::sync::Mutex::new(None);
17426        let mut og = ONES_OV.lock().map_err(|_| "ones lock is poisoned")?;
17427        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
17428            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
17429        }
17430        let ones = &og.as_ref().expect("armed above").1;
17431        e.add_scaled_rows(sh, ones, output, n_embd, 1)?;
17432        Ok(())
17433    }
17434
17435    fn moe_ffn_grouped_add_shared(
17436        e: &Engine,
17437        m: &MoeWeights,
17438        z: &CudaSlice<f32>,
17439        t: usize,
17440        cfg: &ModelConfig,
17441        il: u16,
17442        moe_out: &mut CudaSlice<f32>,
17443    ) -> Result<(), Box<dyn std::error::Error>> {
17444        // MEMRA_STEP_TP_TIMING=1: shared-expert wall (syncs e's stream at exit to bill the
17445        // queued matmuls here rather than at the next host readback).
17446        static SHEXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17447        static SHEXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17448        let shexp_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17449        let shexp_started = shexp_timing.then(std::time::Instant::now);
17450        let result = Self::moe_ffn_grouped_add_shared_inner(e, m, z, t, cfg, il, moe_out);
17451        if let Some(started) = shexp_started {
17452            use std::sync::atomic::Ordering;
17453            e.stream().synchronize()?;
17454            let ns = SHEXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17455                + started.elapsed().as_nanos() as u64;
17456            let calls = SHEXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17457            if calls.is_multiple_of(430) {
17458                eprintln!(
17459                    "[moe-shexp-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
17460                    ns as f64 / 1.0e6,
17461                    ns as f64 / calls as f64 / 1.0e3,
17462                );
17463            }
17464        }
17465        result
17466    }
17467
17468    #[allow(clippy::too_many_arguments)]
17469    fn moe_ffn_grouped_add_shared_inner(
17470        e: &Engine,
17471        m: &MoeWeights,
17472        z: &CudaSlice<f32>,
17473        t: usize,
17474        cfg: &ModelConfig,
17475        il: u16,
17476        moe_out: &mut CudaSlice<f32>,
17477    ) -> Result<(), Box<dyn std::error::Error>> {
17478        let n_embd = cfg.n_embd as usize;
17479        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
17480            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
17481        {
17482            let n_ff_sh = gate_shexp.out_features();
17483            let lim = cfg.clamp_shexp_at(il as u32);
17484            // T=1 DECODE FUSION (2026-08-20): the ffn_swiglu_decode fast-path program, ported
17485            // — one shared quantize feeds gate+up (matmul at m=1 quantizes internally per
17486            // call with identical bytes, so sharing it is bit-identical), the dual NVFP4/Q8
17487            // launch covers both when available, and silu_mul_scaled_q8_1 emits down's
17488            // operand pre-quantized (kernel_check-proven identities). This path measured
17489            // 167us/layer as separate matmuls + 5 allocs at decode.
17490            let fused = t == 1
17491                && lim.is_none()
17492                && cfg.m3.is_none()
17493                && e.uses_q8_1_fast(gate_shexp)
17494                && e.uses_q8_1_fast(up_shexp);
17495            let canonical_w4a16_rows =
17496                t <= 32 && m.step_ep.as_ref().is_some_and(|ep| ep.nvfp4_device_routes);
17497            // MEMRA_BF16_MMV class: both projections in ONE launch (bit-identical per row to
17498            // the two matvec_bf16 launches matmul would issue). W4A16 distributed execution
17499            // uses the same row program for decode and verify; a t=1-only fusion accumulated
17500            // sub-ULP residual drift from the first MoE layer onward.
17501            let bf16_dual = if (t == 1 || canonical_w4a16_rows)
17502                && crate::Engine::bf16_mmv_on()
17503                && n_embd.is_multiple_of(8)
17504            {
17505                match (gate_shexp, up_shexp) {
17506                    (
17507                        crate::model::GpuTensor::FloatBf16 { data: wg, .. },
17508                        crate::model::GpuTensor::FloatBf16 { data: wu, .. },
17509                    ) => Some((wg, wu)),
17510                    _ => None,
17511                }
17512            } else {
17513                None
17514            };
17515            let sh = if let Some((wg, wu)) = bf16_dual {
17516                // Persistent shared-expert workspace: sizes are constant across every MoE
17517                // layer, so one process-level set pinned by (device, n_embd, n_ff_sh) removes
17518                // the four per-layer allocations. Buffers are fully overwritten each call.
17519                type SharedExpertWorkspace = (
17520                    usize,
17521                    usize,
17522                    usize,
17523                    CudaSlice<f32>,
17524                    CudaSlice<f32>,
17525                    CudaSlice<f32>,
17526                    CudaSlice<f32>,
17527                );
17528                static SHEXP_WS: std::sync::Mutex<
17529                    Option<std::collections::HashMap<usize, SharedExpertWorkspace>>,
17530                > = std::sync::Mutex::new(None);
17531                let down_bf16 = match down_shexp {
17532                    crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
17533                    _ => None,
17534                };
17535                let mut guard = SHEXP_WS
17536                    .lock()
17537                    .map_err(|_| "shexp workspace lock is poisoned")?;
17538                let capacity = if canonical_w4a16_rows { 32 } else { 1 };
17539                let device = e.ctx().ordinal();
17540                let workspaces = guard.get_or_insert_with(Default::default);
17541                if workspaces
17542                    .get(&device)
17543                    .is_none_or(|(ne, nf, cap, ..)| (*ne, *nf, *cap) != (n_embd, n_ff_sh, capacity))
17544                {
17545                    workspaces.insert(
17546                        device,
17547                        (
17548                            n_embd,
17549                            n_ff_sh,
17550                            capacity,
17551                            e.uninit(capacity * n_ff_sh)?,
17552                            e.uninit(capacity * n_ff_sh)?,
17553                            e.uninit(capacity * n_ff_sh)?,
17554                            e.uninit(capacity * n_embd)?,
17555                        ),
17556                    );
17557                }
17558                // MEMRA_SHEXP_SPLIT=1: bit-identical row split across both devices; falls
17559                // through to the single-device arm when ineligible.
17560                {
17561                    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17562                    let split_on = *ON
17563                        .get_or_init(|| std::env::var("MEMRA_SHEXP_SPLIT").as_deref() == Ok("1"));
17564                    if split_on
17565                        && t == 1
17566                        && let (Some(wd), Some(rank1)) = (
17567                            match down_shexp {
17568                                crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
17569                                _ => None,
17570                            },
17571                            m.step_tp.as_ref().and_then(|st| st.runtime.rank_engine(1)),
17572                        )
17573                        && let Some(sh) = Self::shexp_split_matvec(
17574                            e, rank1, wg, wu, wd, z, lim, cfg, il, n_embd, n_ff_sh,
17575                        )?
17576                    {
17577                        drop(guard);
17578                        let gate = match &m.gate_inp_shexp {
17579                            Some(gate_inp_shexp) => {
17580                                e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
17581                            }
17582                            None => e.htod(&vec![1.0f32; t])?,
17583                        };
17584                        e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
17585                        return Ok(());
17586                    }
17587                }
17588                let (_, _, _, gate, up, act, sh_buf) = workspaces
17589                    .get_mut(&device)
17590                    .expect("shexp workspace initialized above");
17591                // FUSION #2b needs the POST form its epilogue hardcodes; m3's swigluoai and
17592                // glm5_next's PRE clamp both take the unfused dual-matmul + ffn_act_lim arm.
17593                if let (true, Ok(lim_post)) = (cfg.m3.is_none(), Self::fused_post_limit(lim)) {
17594                    // dual matvec + SwiGLU act in one launch — exact dual per-row program +
17595                    // exact silu/clamped expression, bit-identical.
17596                    if canonical_w4a16_rows {
17597                        e.matvec_bf16_dual_silu_rows_into(
17598                            wg, wu, z, act, n_embd, n_ff_sh, lim_post, t,
17599                        )?;
17600                    } else {
17601                        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim_post)?;
17602                    }
17603                    let _ = (&gate, &up);
17604                } else {
17605                    e.matvec_bf16_dual_into(wg, wu, z, gate, up, n_embd, n_ff_sh)?;
17606                    Self::ffn_act_lim(e, cfg, gate, up, 1.0, 1.0, lim, act, n_ff_sh)?;
17607                }
17608                if let Some(down) = down_bf16 {
17609                    if canonical_w4a16_rows {
17610                        e.matvec_bf16_rows_into(down, act, sh_buf, n_ff_sh, n_embd, t)?;
17611                        let mut sh = e.uninit(t * n_embd)?;
17612                        {
17613                            let mut dst = sh.slice_mut(0..t * n_embd);
17614                            e.stream()
17615                                .memcpy_dtod(&sh_buf.slice(0..t * n_embd), &mut dst)?;
17616                        }
17617                        sh
17618                    } else {
17619                        // FUSION #2e (gate-less shexp only, MEMRA_FUSE_DOWN_ADDSCALE=0 reverts):
17620                        // down matvec + scaled accumulate straight into moe_out in ONE launch —
17621                        // exact f32acc per-row program + the exact add_scaled_rows expression
17622                        // (dst[r] += y_r * 1.0). Replaces down + ownership alloc + 16KB copy +
17623                        // add_scaled (3 launches + alloc -> 1 launch); bit-identical because the
17624                        // accumulate consumes the same f32 the split path stored and reloaded.
17625                        static FUSE_DA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17626                        let fuse_da = *FUSE_DA.get_or_init(|| {
17627                            std::env::var("MEMRA_FUSE_DOWN_ADDSCALE").as_deref() != Ok("0")
17628                        });
17629                        if fuse_da && m.gate_inp_shexp.is_none() {
17630                            static ONES1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
17631                                std::sync::Mutex::new(None);
17632                            let mut og = ONES1.lock().map_err(|_| "shexp ones lock is poisoned")?;
17633                            if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
17634                                *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
17635                            }
17636                            let ones = &og.as_ref().expect("armed above").1;
17637                            e.matvec_bf16_down_addscale_into(
17638                                down, act, ones, moe_out, n_ff_sh, n_embd,
17639                            )?;
17640                            return Ok(());
17641                        }
17642                        e.matvec_bf16_into(down, act, sh_buf, n_ff_sh, n_embd)?;
17643                        let sh = e.uninit(n_embd)?;
17644                        // One alloc keeps the ownership contract; the copy is 16KB on-stream.
17645                        let mut sh = sh;
17646                        {
17647                            let mut dst = sh.slice_mut(0..n_embd);
17648                            e.stream().memcpy_dtod(&sh_buf.slice(0..n_embd), &mut dst)?;
17649                        }
17650                        sh
17651                    }
17652                } else {
17653                    e.matmul(down_shexp, act, t)?
17654                }
17655            } else if fused {
17656                let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
17657                let pair = match e.matmul_pre_dual_noscale(gate_shexp, up_shexp, &zq, &zd, 1)? {
17658                    Some((gate, up)) => Some((gate, up)),
17659                    None => {
17660                        match (
17661                            e.matmul_pre_noscale(gate_shexp, &zq, &zd, 1)?,
17662                            e.matmul_pre_noscale(up_shexp, &zq, &zd, 1)?,
17663                        ) {
17664                            (Some(gate), Some(up)) => Some((gate, up)),
17665                            _ => None,
17666                        }
17667                    }
17668                };
17669                match pair {
17670                    Some(((gate, gs), (up, us))) => {
17671                        if e.uses_q8_1_fast(down_shexp) {
17672                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff_sh)?;
17673                            e.matmul_pre(down_shexp, &aq, &ad, &gate, 1)?
17674                        } else {
17675                            let mut act = e.uninit(n_ff_sh)?;
17676                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff_sh)?;
17677                            e.matmul(down_shexp, &act, 1)?
17678                        }
17679                    }
17680                    None => {
17681                        let gate = e.matmul_pre(gate_shexp, &zq, &zd, z, 1)?;
17682                        let up = e.matmul_pre(up_shexp, &zq, &zd, z, 1)?;
17683                        let mut act = e.uninit(n_ff_sh)?;
17684                        Self::ffn_act(e, cfg, &gate, &up, &mut act, n_ff_sh)?;
17685                        e.matmul(down_shexp, &act, 1)?
17686                    }
17687                }
17688            } else {
17689                let sg_gate = e.matmul(gate_shexp, z, t)?;
17690                let sg_up = e.matmul(up_shexp, z, t)?;
17691                let mut sa = e.uninit(t * n_ff_sh)?;
17692                Self::ffn_act_lim(
17693                    e,
17694                    cfg,
17695                    &sg_gate,
17696                    &sg_up,
17697                    1.0,
17698                    1.0,
17699                    lim,
17700                    &mut sa,
17701                    t * n_ff_sh,
17702                )?;
17703                e.matmul(down_shexp, &sa, t)?
17704            };
17705            let gate = match &m.gate_inp_shexp {
17706                Some(gate_inp_shexp) => {
17707                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
17708                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
17709                    } else {
17710                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
17711                        let mut gate = e.uninit(t)?;
17712                        e.sigmoid(&raw, &mut gate, t)?;
17713                        gate
17714                    }
17715                }
17716                // t=1 hot path: the per-layer htod of a ones row is a PAGEABLE H2D that
17717                // synchronizes the stream — measured as the biggest per-layer host gap
17718                // (44.6us x 42, eager gap table 2026-08-21). One persistent ones row per
17719                // device serves every layer; larger t (prefill) keeps the plain htod.
17720                None if t == 1 => {
17721                    static ONES: std::sync::Mutex<
17722                        Option<std::collections::HashMap<usize, CudaSlice<f32>>>,
17723                    > = std::sync::Mutex::new(None);
17724                    let mut guard = ONES.lock().map_err(|_| "shexp ones lock is poisoned")?;
17725                    let device = e.ctx().ordinal();
17726                    let rows = guard.get_or_insert_with(Default::default);
17727                    // One entry lookup, not three (contains_key + insert + get). The vacant arm
17728                    // stays fallible, which is why this is `match` and not `or_insert_with`.
17729                    use std::collections::hash_map::Entry;
17730                    let ones = match rows.entry(device) {
17731                        Entry::Occupied(occupied) => occupied.into_mut(),
17732                        Entry::Vacant(vacant) => vacant.insert(e.htod(&[1.0f32])?),
17733                    };
17734                    e.add_scaled_rows(&sh, ones, moe_out, n_embd, t)?;
17735                    return Ok(());
17736                }
17737                None => e.htod(&vec![1.0f32; t])?,
17738            };
17739            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
17740        }
17741        Ok(())
17742    }
17743
17744    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
17745    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
17746    pub(crate) fn moe_ffn_grouped(
17747        e: &Engine,
17748        m: &MoeWeights,
17749        z: &CudaSlice<f32>,
17750        t: usize,
17751        cfg: &ModelConfig,
17752        il: u16,
17753        max_block: usize,
17754    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17755        let moe = cfg.moe.as_ref().unwrap();
17756        let n_embd = cfg.n_embd as usize;
17757        let n_expert = moe.expert_count as usize;
17758        let n_used = moe.expert_used_count as usize;
17759        let n_ff_exp = moe.expert_ff_length as usize;
17760        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
17761        let lim_exp = cfg.clamp_exp_at(il as u32);
17762
17763        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
17764        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
17765        // enters the softmax-only pairs/dev router.
17766        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
17767        if let Some(sig) = cfg.sigmoid_router() {
17768            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
17769        }
17770        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
17771            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
17772        } else {
17773            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
17774        };
17775        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
17776        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
17777        Self::trace_moe_input(e, il, t, n_embd, z)?;
17778
17779        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
17780        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
17781        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
17782        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
17783        let no_exp_macros = m.gate_exps.macros.is_none()
17784            && m.up_exps.macros.is_none()
17785            && m.down_exps.macros.is_none();
17786        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
17787            m.has_uniform_expert_layout()
17788                && no_exp_macros
17789                && moe_q8_enabled_for_model(cfg, m)
17790                && moe_slab_enabled()
17791                && dev.dev == e.ctx().ordinal()
17792        });
17793        if let Some(dev) = resident_q8 {
17794            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
17795                e,
17796                m,
17797                z,
17798                t,
17799                cfg,
17800                il,
17801                &sel_all,
17802                &w_all,
17803                &dev.ptr_row,
17804                dev.gu_il,
17805            )?;
17806            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
17807            return Ok(moe_out);
17808        }
17809
17810        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
17811        // For each expert e, we need: which tokens use it, their positions in z, their top-k
17812        // slot index (for bit-identical accumulation), and their weights.
17813        struct ExpertGroup {
17814            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
17815            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
17816            weights: Vec<f32>,      // renormalized weight for that token-expert pair
17817        }
17818        let mut groups: Vec<ExpertGroup> = (0..n_expert)
17819            .map(|_| ExpertGroup {
17820                tok_indices: Vec::new(),
17821                slot_indices: Vec::new(),
17822                weights: Vec::new(),
17823            })
17824            .collect();
17825
17826        for tok in 0..t {
17827            for j in 0..n_used {
17828                let ex = sel_all[tok * n_used + j] as usize;
17829                let w = w_all[tok * n_used + j];
17830                groups[ex].tok_indices.push(tok as i32);
17831                groups[ex].slot_indices.push(j as i32);
17832                groups[ex].weights.push(w);
17833            }
17834        }
17835
17836        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
17837        // Each token's 8 expert contributions land in their respective slots.
17838        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
17839        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
17840
17841        // Expert weight dimensions (used in both cache and staging paths).
17842        let g_len = m.gate_exps.max_expert_bytes();
17843        let u_len = m.up_exps.max_expert_bytes();
17844        let d_len = m.down_exps.max_expert_bytes();
17845        let moe_q8 = moe_q8_enabled_for_model(cfg, m);
17846        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
17847        // Interleaved GU slabs require the pointer-table fast path above.
17848        let slab_local = m
17849            .dev_exps
17850            .as_ref()
17851            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
17852        let use_cache =
17853            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
17854        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
17855        // also does: a local resident slab or a live SLRU dispatch.
17856        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
17857
17858        // GPU scratch for staging (only allocated without a local slab or cache).
17859        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
17860            (
17861                Some(e.alloc_u8(g_len)?),
17862                Some(e.alloc_u8(u_len)?),
17863                Some(e.alloc_u8(d_len)?),
17864            )
17865        } else {
17866            (None, None, None)
17867        };
17868
17869        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
17870        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
17871        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
17872        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
17873        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
17874        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
17875        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
17876        // at long prompts where every expert stages regardless. Order is FREE to change without
17877        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
17878        // regardless of expert processing order (the whole point of the slots).
17879        let mut order: Vec<usize> = (0..n_expert)
17880            .filter(|&ex| !groups[ex].tok_indices.is_empty())
17881            .collect();
17882        order.sort_by(|&a, &b| {
17883            groups[b]
17884                .tok_indices
17885                .len()
17886                .cmp(&groups[a].tok_indices.len())
17887                .then(a.cmp(&b))
17888        });
17889        let mut m_dist: Vec<usize> = Vec::new(); // for stats
17890        let page_window = moe_page_prefetch_window();
17891        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
17892        if worker_disk_prefetch
17893            && let Some(first) = grouped_worker_prefetch_position(order.len(), None)
17894        {
17895            Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
17896        }
17897        for (order_pos, &ex) in order.iter().enumerate() {
17898            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
17899                Self::moe_prefetch_host_expert(order[next], m);
17900            }
17901            if worker_disk_prefetch
17902                && let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos))
17903            {
17904                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
17905                let keep = [
17906                    BlockId::new(il, PROJ_GATE, ex as u16),
17907                    BlockId::new(il, PROJ_UP, ex as u16),
17908                    BlockId::new(il, PROJ_DOWN, ex as u16),
17909                ];
17910                Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
17911            }
17912            let grp = &groups[ex];
17913            let m_e = grp.tok_indices.len();
17914            m_dist.push(m_e);
17915            let gl = m.gate_exps.expert_layout(ex);
17916            let ul = m.up_exps.expert_layout(ex);
17917            let dl = m.down_exps.expert_layout(ex);
17918
17919            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
17920            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
17921            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
17922            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
17923            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
17924            let dmac = m.down_exps.macro_scale(ex);
17925            let weight_d = if dmac == 1.0 {
17926                e.htod(&grp.weights)?
17927            } else {
17928                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
17929                e.htod(&scaled)?
17930            };
17931
17932            // GATHER: collect m_e activation rows from z into a contiguous buffer.
17933            let mut gathered = e.zeros(m_e * n_embd)?;
17934            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
17935            let gv = gathered.slice(0..m_e * n_embd);
17936
17937            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
17938            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
17939            let y = if let Some(dev) = slab_local {
17940                let gate_start = ex * m.gate_exps.expert_stride;
17941                let up_start = ex * m.up_exps.expert_stride;
17942                let down_start = ex * m.down_exps.expert_stride;
17943                if grouped_q8 {
17944                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
17945                    let gate = e.qmatvec_expert_q8(
17946                        &dev.gate,
17947                        gate_start..gate_start + gl.len,
17948                        &zq,
17949                        &zd,
17950                        m_e,
17951                        m.gate_exps.in_f,
17952                        m.gate_exps.out_f,
17953                        gl.qtype,
17954                        gl.row_bytes,
17955                    )?;
17956                    let up = e.qmatvec_expert_q8(
17957                        &dev.up,
17958                        up_start..up_start + ul.len,
17959                        &zq,
17960                        &zd,
17961                        m_e,
17962                        m.up_exps.in_f,
17963                        m.up_exps.out_f,
17964                        ul.qtype,
17965                        ul.row_bytes,
17966                    )?;
17967                    let mut act = e.uninit(m_e * n_ff_exp)?;
17968                    Self::ffn_act_lim(
17969                        e,
17970                        cfg,
17971                        &gate,
17972                        &up,
17973                        m.gate_exps.macro_scale(ex),
17974                        m.up_exps.macro_scale(ex),
17975                        lim_exp,
17976                        &mut act,
17977                        m_e * n_ff_exp,
17978                    )?;
17979                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
17980                    e.qmatvec_expert_q8(
17981                        &dev.down,
17982                        down_start..down_start + dl.len,
17983                        &aq2,
17984                        &ad2,
17985                        m_e,
17986                        m.down_exps.in_f,
17987                        m.down_exps.out_f,
17988                        dl.qtype,
17989                        dl.row_bytes,
17990                    )?
17991                } else {
17992                    let gate = m.qmatvec_view(
17993                        e,
17994                        &dev.gate,
17995                        gate_start..gate_start + gl.len,
17996                        &gv,
17997                        m_e,
17998                        m.gate_exps.in_f,
17999                        m.gate_exps.out_f,
18000                        gl.qtype,
18001                        gl.row_bytes,
18002                    )?;
18003                    let up = m.qmatvec_view(
18004                        e,
18005                        &dev.up,
18006                        up_start..up_start + ul.len,
18007                        &gv,
18008                        m_e,
18009                        m.up_exps.in_f,
18010                        m.up_exps.out_f,
18011                        ul.qtype,
18012                        ul.row_bytes,
18013                    )?;
18014                    let mut act = e.uninit(m_e * n_ff_exp)?;
18015                    Self::ffn_act_lim(
18016                        e,
18017                        cfg,
18018                        &gate,
18019                        &up,
18020                        m.gate_exps.macro_scale(ex),
18021                        m.up_exps.macro_scale(ex),
18022                        lim_exp,
18023                        &mut act,
18024                        m_e * n_ff_exp,
18025                    )?;
18026                    let actv = act.slice(0..m_e * n_ff_exp);
18027                    m.qmatvec_view(
18028                        e,
18029                        &dev.down,
18030                        down_start..down_start + dl.len,
18031                        &actv,
18032                        m_e,
18033                        m.down_exps.in_f,
18034                        m.down_exps.out_f,
18035                        dl.qtype,
18036                        dl.row_bytes,
18037                    )?
18038                }
18039            } else if use_cache {
18040                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
18041                if grouped_q8 {
18042                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
18043                    let gate = e.with_moe_cache(max_block, |cache, eng| {
18044                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
18045                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
18046                        eng.qmatvec_expert_q8(
18047                            cache.buf(slot),
18048                            0..gl.len,
18049                            &zq,
18050                            &zd,
18051                            m_e,
18052                            m.gate_exps.in_f,
18053                            m.gate_exps.out_f,
18054                            gl.qtype,
18055                            gl.row_bytes,
18056                        )
18057                    })?;
18058                    let up = e.with_moe_cache(max_block, |cache, eng| {
18059                        let id = BlockId::new(il, PROJ_UP, ex as u16);
18060                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
18061                        eng.qmatvec_expert_q8(
18062                            cache.buf(slot),
18063                            0..ul.len,
18064                            &zq,
18065                            &zd,
18066                            m_e,
18067                            m.up_exps.in_f,
18068                            m.up_exps.out_f,
18069                            ul.qtype,
18070                            ul.row_bytes,
18071                        )
18072                    })?;
18073                    let mut act = e.uninit(m_e * n_ff_exp)?;
18074                    Self::ffn_act_lim(
18075                        e,
18076                        cfg,
18077                        &gate,
18078                        &up,
18079                        m.gate_exps.macro_scale(ex),
18080                        m.up_exps.macro_scale(ex),
18081                        lim_exp,
18082                        &mut act,
18083                        m_e * n_ff_exp,
18084                    )?;
18085                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
18086                    e.with_moe_cache(max_block, |cache, eng| {
18087                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
18088                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
18089                        eng.qmatvec_expert_q8(
18090                            cache.buf(slot),
18091                            0..dl.len,
18092                            &aq2,
18093                            &ad2,
18094                            m_e,
18095                            m.down_exps.in_f,
18096                            m.down_exps.out_f,
18097                            dl.qtype,
18098                            dl.row_bytes,
18099                        )
18100                    })?
18101                } else {
18102                    let gate = e.with_moe_cache(max_block, |cache, eng| {
18103                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
18104                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
18105                        m.qmatvec_view(
18106                            eng,
18107                            cache.buf(slot),
18108                            0..gl.len,
18109                            &gv,
18110                            m_e,
18111                            m.gate_exps.in_f,
18112                            m.gate_exps.out_f,
18113                            gl.qtype,
18114                            gl.row_bytes,
18115                        )
18116                    })?;
18117                    let up = e.with_moe_cache(max_block, |cache, eng| {
18118                        let id = BlockId::new(il, PROJ_UP, ex as u16);
18119                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
18120                        m.qmatvec_view(
18121                            eng,
18122                            cache.buf(slot),
18123                            0..ul.len,
18124                            &gv,
18125                            m_e,
18126                            m.up_exps.in_f,
18127                            m.up_exps.out_f,
18128                            ul.qtype,
18129                            ul.row_bytes,
18130                        )
18131                    })?;
18132                    let mut act = e.uninit(m_e * n_ff_exp)?;
18133                    Self::ffn_act_lim(
18134                        e,
18135                        cfg,
18136                        &gate,
18137                        &up,
18138                        m.gate_exps.macro_scale(ex),
18139                        m.up_exps.macro_scale(ex),
18140                        lim_exp,
18141                        &mut act,
18142                        m_e * n_ff_exp,
18143                    )?;
18144                    let actv = act.slice(0..m_e * n_ff_exp);
18145                    e.with_moe_cache(max_block, |cache, eng| {
18146                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
18147                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
18148                        m.qmatvec_view(
18149                            eng,
18150                            cache.buf(slot),
18151                            0..dl.len,
18152                            &actv,
18153                            m_e,
18154                            m.down_exps.in_f,
18155                            m.down_exps.out_f,
18156                            dl.qtype,
18157                            dl.row_bytes,
18158                        )
18159                    })?
18160                }
18161            } else {
18162                let sg = scratch_g.as_mut().unwrap();
18163                let su = scratch_u.as_mut().unwrap();
18164                let sd = scratch_d.as_mut().unwrap();
18165                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
18166                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
18167                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
18168                if grouped_q8 {
18169                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
18170                    let gate = e.qmatvec_expert_q8(
18171                        sg,
18172                        0..gl.len,
18173                        &zq,
18174                        &zd,
18175                        m_e,
18176                        m.gate_exps.in_f,
18177                        m.gate_exps.out_f,
18178                        gl.qtype,
18179                        gl.row_bytes,
18180                    )?;
18181                    let up = e.qmatvec_expert_q8(
18182                        su,
18183                        0..ul.len,
18184                        &zq,
18185                        &zd,
18186                        m_e,
18187                        m.up_exps.in_f,
18188                        m.up_exps.out_f,
18189                        ul.qtype,
18190                        ul.row_bytes,
18191                    )?;
18192                    let mut act = e.uninit(m_e * n_ff_exp)?;
18193                    Self::ffn_act_lim(
18194                        e,
18195                        cfg,
18196                        &gate,
18197                        &up,
18198                        m.gate_exps.macro_scale(ex),
18199                        m.up_exps.macro_scale(ex),
18200                        lim_exp,
18201                        &mut act,
18202                        m_e * n_ff_exp,
18203                    )?;
18204                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
18205                    e.qmatvec_expert_q8(
18206                        sd,
18207                        0..dl.len,
18208                        &aq2,
18209                        &ad2,
18210                        m_e,
18211                        m.down_exps.in_f,
18212                        m.down_exps.out_f,
18213                        dl.qtype,
18214                        dl.row_bytes,
18215                    )?
18216                } else {
18217                    let gate = m.qmatvec_view(
18218                        e,
18219                        sg,
18220                        0..gl.len,
18221                        &gv,
18222                        m_e,
18223                        m.gate_exps.in_f,
18224                        m.gate_exps.out_f,
18225                        gl.qtype,
18226                        gl.row_bytes,
18227                    )?;
18228                    let up = m.qmatvec_view(
18229                        e,
18230                        su,
18231                        0..ul.len,
18232                        &gv,
18233                        m_e,
18234                        m.up_exps.in_f,
18235                        m.up_exps.out_f,
18236                        ul.qtype,
18237                        ul.row_bytes,
18238                    )?;
18239                    let mut act = e.uninit(m_e * n_ff_exp)?;
18240                    Self::ffn_act_lim(
18241                        e,
18242                        cfg,
18243                        &gate,
18244                        &up,
18245                        m.gate_exps.macro_scale(ex),
18246                        m.up_exps.macro_scale(ex),
18247                        lim_exp,
18248                        &mut act,
18249                        m_e * n_ff_exp,
18250                    )?;
18251                    let actv = act.slice(0..m_e * n_ff_exp);
18252                    m.qmatvec_view(
18253                        e,
18254                        sd,
18255                        0..dl.len,
18256                        &actv,
18257                        m_e,
18258                        m.down_exps.in_f,
18259                        m.down_exps.out_f,
18260                        dl.qtype,
18261                        dl.row_bytes,
18262                    )?
18263                }
18264            };
18265
18266            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
18267            e.scatter_slot(
18268                &y,
18269                &tok_idx_d,
18270                &slot_idx_d,
18271                &weight_d,
18272                &mut slot_buf,
18273                &mut wbuf,
18274                n_embd,
18275                n_used,
18276                m_e,
18277            )?;
18278        }
18279
18280        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
18281        let mut moe_out = e.zeros(t * n_embd)?;
18282        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
18283
18284        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
18285        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
18286            m_dist.sort_unstable();
18287            let active = m_dist.len();
18288            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
18289            let median = m_dist[active / 2];
18290            let max_m = *m_dist.last().unwrap();
18291            let min_m = m_dist[0];
18292            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
18293            println!(
18294                "moe-grouped il={il} t={t} active={active}/{n_expert} \
18295                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
18296                      above_gemm_threshold(>=16)={above16}/{active}"
18297            );
18298        }
18299
18300        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
18301        Ok(moe_out)
18302    }
18303
18304    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
18305    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
18306    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
18307    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
18308    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
18309    /// expert-sum order identical to the sequential path.
18310    pub(crate) fn moe_ffn_lockstep(
18311        &self,
18312        e: &Engine,
18313        m: &MoeWeights,
18314        zbatch: &CudaSlice<f32>,
18315        mrows: usize,
18316        il: u16,
18317        max_block: usize,
18318    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18319        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
18320        let cfg = &self.cfg;
18321        let moe = cfg.moe.as_ref().unwrap();
18322        let n_embd = cfg.n_embd as usize;
18323        let n_expert = moe.expert_count as usize;
18324        let n_used = moe.expert_used_count as usize;
18325        let n_ff_exp = moe.expert_ff_length as usize;
18326        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
18327        let lim_exp = cfg.clamp_exp_at(il as u32);
18328        let lim_shexp = cfg.clamp_shexp_at(il as u32);
18329
18330        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
18331        if let Some(sig) = cfg.sigmoid_router() {
18332            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
18333        }
18334        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
18335            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
18336        } else {
18337            Self::moe_route_cfg(
18338                e,
18339                &logits,
18340                mrows,
18341                n_expert,
18342                n_used,
18343                m.active_experts.as_deref(),
18344            )?
18345        };
18346        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
18347
18348        // Residency split at whole-expert granularity against the (frozen) cache.
18349        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
18350            Ok((0..n_expert)
18351                .map(|ex| {
18352                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
18353                        .into_iter()
18354                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
18355                })
18356                .collect())
18357        })?;
18358
18359        struct Group {
18360            rows: Vec<i32>,
18361            slots: Vec<i32>,
18362            weights: Vec<f32>,
18363        }
18364        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
18365        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
18366        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
18367            Default::default();
18368        for row in 0..mrows {
18369            for j in 0..n_used {
18370                let ex = sel_all[row * n_used + j] as usize;
18371                let w = w_all[row * n_used + j];
18372                if resident_expert[ex] {
18373                    let group = groups.entry(ex).or_insert_with(|| Group {
18374                        rows: Vec::new(),
18375                        slots: Vec::new(),
18376                        weights: Vec::new(),
18377                    });
18378                    group.rows.push(row as i32);
18379                    group.slots.push(j as i32);
18380                    group.weights.push(w);
18381                } else {
18382                    crate::cpu_experts::record_incomplete_gpu_residency(0);
18383                    cpu_rows[row].push((ex, w));
18384                    cpu_by_expert.entry(ex).or_default().push((row, w));
18385                }
18386            }
18387        }
18388
18389        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
18390        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
18391        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
18392        // order per row differs from the sequential single-call chunk — part of the
18393        // documented lockstep numeric class.
18394        let host_rows = e.dtoh(zbatch)?;
18395        let rows_ok = crate::cpu_experts::rows_supported();
18396        enum CpuPart {
18397            Single { row: usize },
18398            Rows { rows: Vec<usize> },
18399        }
18400        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
18401        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
18402        if rows_ok {
18403            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
18404                .into_iter()
18405                .filter(|(_, rows)| rows.len() >= 2)
18406                .collect();
18407            shared.sort_by_key(|(ex, _)| *ex);
18408            for (ex, mut row_weights) in shared {
18409                row_weights.sort_by_key(|(row, _)| *row);
18410                let inputs: Vec<(&[f32], f32)> = row_weights
18411                    .iter()
18412                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
18413                    .collect();
18414                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
18415                    .map_err(std::io::Error::other)?;
18416                for &(row, _) in &row_weights {
18417                    rows_served.insert((row, ex));
18418                }
18419                tickets.push((
18420                    CpuPart::Rows {
18421                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
18422                    },
18423                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
18424                ));
18425            }
18426        }
18427        for (row, selected) in cpu_rows.iter().enumerate() {
18428            let leftover: Vec<(usize, f32)> = selected
18429                .iter()
18430                .copied()
18431                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
18432                .collect();
18433            if leftover.is_empty() {
18434                continue;
18435            }
18436            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
18437            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
18438                .map_err(std::io::Error::other)?;
18439            tickets.push((
18440                CpuPart::Single { row },
18441                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
18442            ));
18443        }
18444
18445        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
18446        let mut wbuf = e.zeros(mrows * n_used)?;
18447        let mut order: Vec<usize> = groups.keys().copied().collect();
18448        order.sort_by(|&a, &b| {
18449            groups[&b]
18450                .rows
18451                .len()
18452                .cmp(&groups[&a].rows.len())
18453                .then(a.cmp(&b))
18454        });
18455        for &ex in &order {
18456            let group = &groups[&ex];
18457            let m_e = group.rows.len();
18458            let gl = m.gate_exps.expert_layout(ex);
18459            let ul = m.up_exps.expert_layout(ex);
18460            let dl = m.down_exps.expert_layout(ex);
18461            let row_idx_d = e.htod_i32(&group.rows)?;
18462            let slot_idx_d = e.htod_i32(&group.slots)?;
18463            let dmac = m.down_exps.macro_scale(ex);
18464            let weight_d = if dmac == 1.0 {
18465                e.htod(&group.weights)?
18466            } else {
18467                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
18468                e.htod(&scaled)?
18469            };
18470            let mut gathered = e.zeros(m_e * n_embd)?;
18471            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
18472            let gv = gathered.slice(0..m_e * n_embd);
18473            let gate = e.with_moe_cache(max_block, |c, eng| {
18474                let slot = c
18475                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
18476                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
18477                m.qmatvec_view(
18478                    eng,
18479                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
18480                    0..gl.len,
18481                    &gv,
18482                    m_e,
18483                    m.gate_exps.in_f,
18484                    m.gate_exps.out_f,
18485                    gl.qtype,
18486                    gl.row_bytes,
18487                )
18488            })?;
18489            let up = e.with_moe_cache(max_block, |c, eng| {
18490                let slot = c
18491                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
18492                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
18493                m.qmatvec_view(
18494                    eng,
18495                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
18496                    0..ul.len,
18497                    &gv,
18498                    m_e,
18499                    m.up_exps.in_f,
18500                    m.up_exps.out_f,
18501                    ul.qtype,
18502                    ul.row_bytes,
18503                )
18504            })?;
18505            let mut act = e.zeros(m_e * n_ff_exp)?;
18506            Self::ffn_act_lim(
18507                e,
18508                cfg,
18509                &gate,
18510                &up,
18511                m.gate_exps.macro_scale(ex),
18512                m.up_exps.macro_scale(ex),
18513                lim_exp,
18514                &mut act,
18515                m_e * n_ff_exp,
18516            )?;
18517            let actv = act.slice(0..m_e * n_ff_exp);
18518            let y = e.with_moe_cache(max_block, |c, eng| {
18519                let slot = c
18520                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
18521                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
18522                m.qmatvec_view(
18523                    eng,
18524                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
18525                    0..dl.len,
18526                    &actv,
18527                    m_e,
18528                    m.down_exps.in_f,
18529                    m.down_exps.out_f,
18530                    dl.qtype,
18531                    dl.row_bytes,
18532                )
18533            })?;
18534            e.scatter_slot(
18535                &y,
18536                &row_idx_d,
18537                &slot_idx_d,
18538                &weight_d,
18539                &mut slot_buf,
18540                &mut wbuf,
18541                n_embd,
18542                n_used,
18543                m_e,
18544            )?;
18545        }
18546        let mut moe_out = e.zeros(mrows * n_embd)?;
18547        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
18548
18549        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
18550        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
18551        for (part, ticket) in tickets {
18552            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
18553            let mut add_row = |row: usize, chunk: &[f32]| {
18554                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
18555                for (accumulator, value) in sum.iter_mut().zip(chunk) {
18556                    *accumulator += value;
18557                }
18558            };
18559            match part {
18560                CpuPart::Single { row } => add_row(row, &cpu_output),
18561                CpuPart::Rows { rows } => {
18562                    for (slot, row) in rows.into_iter().enumerate() {
18563                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
18564                    }
18565                }
18566            }
18567        }
18568        for (row, sum) in row_sums.into_iter().enumerate() {
18569            let Some(sum) = sum else { continue };
18570            let cpu_output = e.htod(&sum)?;
18571            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
18572            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
18573        }
18574
18575        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
18576            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
18577        {
18578            let n_ff_sh = gate_shexp.out_features();
18579            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
18580            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
18581            let mut sa = e.zeros(mrows * n_ff_sh)?;
18582            Self::ffn_act_lim(
18583                e,
18584                cfg,
18585                &sg_gate,
18586                &sg_up,
18587                1.0,
18588                1.0,
18589                lim_shexp,
18590                &mut sa,
18591                mrows * n_ff_sh,
18592            )?;
18593            let sh = e.matmul(down_shexp, &sa, mrows)?;
18594            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
18595            // decode matches the single-sequence decode chain bit-for-bit.
18596            let g = match &m.gate_inp_shexp {
18597                Some(gate_inp_shexp) => {
18598                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
18599                }
18600                None => e.htod(&vec![1.0f32; mrows])?,
18601            };
18602            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
18603        }
18604
18605        Ok(moe_out)
18606    }
18607}
18608
18609// ============================ gemma4 (R8 verified wiring) ==================================
18610// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
18611// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
18612// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
18613// gemma variants after the correctness gate).
18614impl HybridModel {
18615    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
18616    /// Rotary width the gemma-4 config DECLARES for layer `il`, per SWA/global class — the
18617    /// `n_rot` argument of the fused `rms_norm_qkv_rope*` kernels
18618    /// (`Engine::full_width_rope_only`, lane/graph-s-key-exactness-20260819).
18619    ///
18620    /// It exists so the fusion's full-width assumption is checked against the CONFIG instead of
18621    /// against itself: passing `hd` twice would guard nothing. On every gemma-4 artifact both
18622    /// routes set `rope_dims_{global,swa} == key_length_{global,swa}` (GGUF reads
18623    /// `rope.dimension_count(_swa)`, the safetensors route derives them from
18624    /// `global_head_dim`/`head_dim`), which is exactly WHY the fusion is legal here — gemma's
18625    /// "partial" rotary is per-dimension frequency scaling through the `rope_freqs` factors
18626    /// tensor, not an n_rot truncation. An arch that truncates instead (qwen3.5: 64 of 256;
18627    /// step35 full-attn: 64 of 128) must not reach these kernels, and now cannot silently.
18628    pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
18629        let g = self
18630            .cfg
18631            .gemma4
18632            .as_ref()
18633            .expect("gemma4_rope_dims on a non-gemma4 config");
18634        if g.swa_pattern[il] {
18635            g.rope_dims_swa as usize
18636        } else {
18637            g.rope_dims_global as usize
18638        }
18639    }
18640
18641    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
18642        let g = self.cfg.gemma4.as_ref().unwrap();
18643        let swa = g.swa_pattern[il];
18644        let hd = if swa {
18645            g.key_length_swa
18646        } else {
18647            g.key_length_global
18648        } as usize;
18649        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
18650        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
18651        // rows exact (softmax over one element) while every later position drifted).
18652        (
18653            hd,
18654            g.head_count_kv[il] as usize,
18655            self.cfg.n_head as usize,
18656            if swa {
18657                g.rope_base_swa
18658            } else {
18659                g.rope_base_global
18660            },
18661            1.0,
18662            swa,
18663        )
18664    }
18665
18666    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
18667    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
18668    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
18669    pub(crate) fn gemma4_suppress(
18670        &self,
18671        e: &Engine,
18672        ld: &mut CudaSlice<f32>,
18673        t: usize,
18674    ) -> Result<(), Box<dyn std::error::Error>> {
18675        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
18676            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
18677            // stage as primary, and this tail runs only after the last stage). The assert turns
18678            // that argued invariant into a checked one: any topology violating primary==head
18679            // trips here in debug instead of silently peer-reading a device-0 buffer.
18680            #[cfg(debug_assertions)]
18681            crate::debug_assert_tensor_stream_device(
18682                ids,
18683                &e.stream(),
18684                "gemma4_suppress.suppress_d",
18685            );
18686            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
18687        }
18688        Ok(())
18689    }
18690
18691    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
18692    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
18693    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
18694    /// only (v0): attends within `tokens` via the f32 sdpa.
18695    #[allow(clippy::too_many_arguments)]
18696    /// ONE-PROGRAM seam (lane/gemma-restore-exactness-20260819, DEFAULT OFF): route dense
18697    /// gemma SWA prefill attention through the WINDOWED arm at every prompt length instead of
18698    /// switching program at `t > sliding_window`. The door is the measured cause of the
18699    /// prefix-restore exactness violation: below it the SWA layers attend unwindowed on f32
18700    /// operands, above it windowed on bf16 operands, so rows[0..932) of a 932-row prime and of
18701    /// a 1048-row prime differ in 59 of 60 layers. `=1` collapses it to one program so the
18702    /// published prefix KV stops depending on the total prompt length. Off by default because
18703    /// collapsing changes the bytes short prompts get today; the flip is an owner call.
18704    fn gemma_fa_one_program() -> bool {
18705        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18706        *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
18707    }
18708
18709    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
18710    fn gemma4_attn_prime(
18711        &self,
18712        e: &Engine,
18713        fa: &crate::hybrid::FullAttnLayer,
18714        il: usize,
18715        h: &CudaSlice<f32>,
18716        pos_d: &CudaSlice<i32>,
18717        t: usize,
18718        cache: Option<&mut Cache>,
18719        island: Option<&CudaSlice<i32>>,
18720    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18721        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18722        let eps = self.cfg.rms_eps;
18723        let aux = self.gemma4_aux.as_ref().unwrap();
18724        let ones = aux.ones(e);
18725        #[cfg(debug_assertions)]
18726        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
18727
18728        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
18729        // (h stays borrowed across the triple, so the cache key can't go stale).
18730        e.mmq_act_begin();
18731        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
18732        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
18733            let v = e.dtoh(&q0)?;
18734            let nan = v.iter().filter(|x| x.is_nan()).count();
18735            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
18736            eprintln!(
18737                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
18738                v.len()
18739            );
18740        }
18741        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
18742        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
18743        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
18744        let v0 = if swa {
18745            e.matmul(&fa.wv, h, t)?
18746        } else {
18747            e.clone_dtod(&k0)?
18748        };
18749        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
18750            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
18751                let v = e.dtoh(buf)?;
18752                let nan = v.iter().filter(|x| x.is_nan()).count();
18753                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
18754                eprintln!(
18755                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
18756                    v.len()
18757                );
18758            }
18759        }
18760
18761        let mut q = e.uninit(t * nh * hd)?;
18762        let mut k = e.uninit(t * nkv * hd)?;
18763        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
18764        let mut v = e.uninit(t * nkv * hd)?;
18765        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
18766        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
18767        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
18768        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18769        // Island primes take the mask-capable naive kernel below; keep the operands f32
18770        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
18771        let emit = island.is_none()
18772            && t >= 16
18773            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
18774            && *EMIT.get_or_init(|| {
18775                std::env::var("MEMRA_FA_EMIT")
18776                    .map(|s| s != "0")
18777                    .unwrap_or(true)
18778            });
18779        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
18780        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
18781        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
18782        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
18783        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
18784        let v_f16 = emit
18785            && crate::fa_f16pv_on()
18786            && match hd {
18787                512 => true,
18788                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
18789                _ => false,
18790            };
18791        if emit {
18792            e.rms_norm_qkv_w4b(
18793                &q0,
18794                &k0,
18795                &v0,
18796                fa.q_norm.float_data(),
18797                fa.k_norm.float_data(),
18798                ones,
18799                &mut q,
18800                &mut k,
18801                &mut v,
18802                &mut vb,
18803                hd,
18804                nh * t,
18805                nkv * t,
18806                eps,
18807                v_f16,
18808            )?;
18809        } else {
18810            e.rms_norm_qkv(
18811                &q0,
18812                &k0,
18813                &v0,
18814                fa.q_norm.float_data(),
18815                fa.k_norm.float_data(),
18816                ones,
18817                &mut q,
18818                &mut k,
18819                &mut v,
18820                hd,
18821                nh * t,
18822                nkv * t,
18823                eps,
18824            )?;
18825        }
18826
18827        let ff = if swa {
18828            None
18829        } else {
18830            Some(
18831                aux.rope_freqs(e)
18832                    .expect("gemma4 global rope needs rope_freqs.weight"),
18833            )
18834        };
18835        #[cfg(debug_assertions)]
18836        if let Some(ff) = ff {
18837            crate::debug_assert_tensor_stream_device(
18838                ff,
18839                &e.stream(),
18840                "gemma4_attn_prime.rope_freqs",
18841            );
18842        }
18843        if emit {
18844            e.rope_neox2_bf16e(
18845                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
18846            )?;
18847        } else {
18848            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
18849        }
18850
18851        if let Some(cache) = cache {
18852            let kvl = cache.kv[il].as_mut().unwrap();
18853            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
18854            e.append_kv_quantized_rows(
18855                &k,
18856                &v,
18857                &mut kvl.k,
18858                &mut kvl.v,
18859                kvl.len,
18860                t,
18861                kvl.kv_dim_k,
18862                kvl.kv_dim_v,
18863                kvl.k_tok_bytes,
18864                kvl.v_tok_bytes,
18865                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
18866            )?;
18867            kvl.len += t;
18868        }
18869        let mut attn = e.zeros(t * nh * hd)?;
18870        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
18871        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
18872        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
18873        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
18874        if let Some(span) = island {
18875            // Masked-prefill arm: every layer routes through the island-aware naive
18876            // kernel (correctness-first, same posture as the vision tower v1). The
18877            // window argument keeps the R6 shortcut: 0 while the prompt fits the
18878            // window, the real window beyond it.
18879            let w = if swa && t > win { win } else { 0 };
18880            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
18881        } else if swa && (t > win || Self::gemma_fa_one_program()) {
18882            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
18883                if emit {
18884                    e.fa_prefill_w_pre(
18885                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
18886                    )?;
18887                } else {
18888                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
18889                }
18890            } else {
18891                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
18892            }
18893        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
18894            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
18895        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
18896            if emit {
18897                e.fa_prefill_hd512_pre(
18898                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
18899                )?;
18900            } else {
18901                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
18902            }
18903        } else {
18904            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
18905        }
18906        e.matmul(&fa.wo, &attn, t)
18907    }
18908
18909    /// Back-compat wrapper (pure prefill, no cache).
18910    fn gemma4_attn(
18911        &self,
18912        e: &Engine,
18913        fa: &crate::hybrid::FullAttnLayer,
18914        il: usize,
18915        h: &CudaSlice<f32>,
18916        pos_d: &CudaSlice<i32>,
18917        t: usize,
18918    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18919        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
18920    }
18921
18922    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
18923    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
18924    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
18925    /// the q8z epilogue is quantize_q8_1 verbatim).
18926    fn gemma4_moe_q8(
18927        &self,
18928        e: &Engine,
18929        m: &crate::hybrid::MoeWeights,
18930        bits: &crate::hybrid::Gemma4MoeBits,
18931        mq: &(CudaSlice<i8>, CudaSlice<f32>),
18932        router_in: &CudaSlice<f32>,
18933        t: usize,
18934    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18935        let cfg = &self.cfg;
18936        let moe = cfg.moe.as_ref().unwrap();
18937        let n_embd = cfg.n_embd as usize;
18938        let n_expert = moe.expert_count as usize;
18939        let n_used = moe.expert_used_count as usize;
18940        let n_ff_exp = moe.expert_ff_length as usize;
18941        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
18942        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
18943        // the pair's 12us is kernel time, not launch gaps.
18944        let logits = if crate::router_kernel_on() {
18945            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
18946        } else {
18947            e.matmul(&m.gate_inp, router_in, t)?
18948        };
18949        let dev = m.dev_exps.as_ref().unwrap();
18950        let (sel_d, w_d) =
18951            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
18952        let (zq, zd) = mq;
18953        if t == 1 {
18954            let selv = sel_d.slice(0..n_used);
18955            let wv = w_d.slice(0..n_used);
18956            let act = e.moe_gate_up_gelu8_dev_q8(
18957                &dev.ptr_row,
18958                &selv,
18959                zq,
18960                zd,
18961                n_embd,
18962                n_ff_exp,
18963                n_used,
18964                n_expert,
18965                m.gate_exps.qtype,
18966                m.up_exps.qtype,
18967                m.gate_exps.row_bytes,
18968                m.up_exps.row_bytes,
18969            )?;
18970            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
18971            let mut moe_out = e.uninit(n_embd)?;
18972            e.moe_down8_fma_dev_q8(
18973                &dev.ptr_row,
18974                &selv,
18975                &wv,
18976                &aq2,
18977                &ad2,
18978                &mut moe_out.slice_mut(0..n_embd),
18979                n_ff_exp,
18980                n_embd,
18981                n_used,
18982                n_expert,
18983                m.down_exps.qtype,
18984                m.down_exps.row_bytes,
18985            )?;
18986            return Ok(moe_out);
18987        }
18988        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
18989        let act = if csr {
18990            e.moe_gate_up_gelu8_dev_q8_csr(
18991                &dev.ptr_row,
18992                &sel_d,
18993                zq,
18994                zd,
18995                t * n_used,
18996                n_embd,
18997                n_ff_exp,
18998                n_used,
18999                n_expert,
19000                m.gate_exps.qtype,
19001                m.up_exps.qtype,
19002                m.gate_exps.row_bytes,
19003                m.up_exps.row_bytes,
19004            )?
19005        } else {
19006            e.moe_gate_up_gelu8_dev_q8_rows(
19007                &dev.ptr_row,
19008                &sel_d,
19009                zq,
19010                zd,
19011                t,
19012                n_embd,
19013                n_ff_exp,
19014                n_used,
19015                n_expert,
19016                m.gate_exps.qtype,
19017                m.up_exps.qtype,
19018                m.gate_exps.row_bytes,
19019                m.up_exps.row_bytes,
19020            )?
19021        };
19022        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
19023        let mut moe_out = e.uninit(t * n_embd)?;
19024        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
19025        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
19026        e.moe_down8_fma_dev_q8_rows_g(
19027            &dev.ptr_row,
19028            &sel_d,
19029            &w_d,
19030            &aq2,
19031            &ad2,
19032            &mut moe_out,
19033            t,
19034            n_ff_exp,
19035            n_embd,
19036            n_used,
19037            n_expert,
19038            m.down_exps.qtype,
19039            m.down_exps.row_bytes,
19040        )?;
19041        Ok(moe_out)
19042    }
19043
19044    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
19045    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
19046    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
19047    fn gemma4_moe(
19048        &self,
19049        e: &Engine,
19050        m: &crate::hybrid::MoeWeights,
19051        bits: &crate::hybrid::Gemma4MoeBits,
19052        moe_in: &CudaSlice<f32>,
19053        router_in: &CudaSlice<f32>,
19054        t: usize,
19055    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19056        crate::moe_rp_refuse(m.dev_exps.as_ref().is_some_and(|d| d.rp), "gemma4_moe")?; // memra#147: no split-plane arm here
19057        let cfg = &self.cfg;
19058        let moe = cfg.moe.as_ref().unwrap();
19059        let n_embd = cfg.n_embd as usize;
19060        let n_expert = moe.expert_count as usize;
19061        let n_used = moe.expert_used_count as usize;
19062        let n_ff_exp = moe.expert_ff_length as usize;
19063
19064        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
19065        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
19066        // batched matmul only at real prefill.
19067        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
19068            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
19069        } else {
19070            e.matmul(&m.gate_inp, router_in, t)?
19071        };
19072
19073        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
19074        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
19075        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
19076        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
19077        if t < PRIME_MIN_T
19078            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
19079            && expert_dp4a_supported(m.gate_exps.qtype)
19080            && expert_dp4a_supported(m.up_exps.qtype)
19081            && expert_dp4a_supported(m.down_exps.qtype)
19082            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
19083        {
19084            let dev = m.dev_exps.as_ref().unwrap();
19085            let (sel_d, w_d) =
19086                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
19087            if t == 1 {
19088                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
19089                let selv = sel_d.slice(0..n_used);
19090                let wv = w_d.slice(0..n_used);
19091                let act = e.moe_gate_up_gelu8_dev_q8(
19092                    &dev.ptr_row,
19093                    &selv,
19094                    &zq,
19095                    &zd,
19096                    n_embd,
19097                    n_ff_exp,
19098                    n_used,
19099                    n_expert,
19100                    m.gate_exps.qtype,
19101                    m.up_exps.qtype,
19102                    m.gate_exps.row_bytes,
19103                    m.up_exps.row_bytes,
19104                )?;
19105                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
19106                let mut moe_out = e.uninit(n_embd)?;
19107                e.moe_down8_fma_dev_q8(
19108                    &dev.ptr_row,
19109                    &selv,
19110                    &wv,
19111                    &aq2,
19112                    &ad2,
19113                    &mut moe_out.slice_mut(0..n_embd),
19114                    n_ff_exp,
19115                    n_embd,
19116                    n_used,
19117                    n_expert,
19118                    m.down_exps.qtype,
19119                    m.down_exps.row_bytes,
19120                )?;
19121                return Ok(moe_out);
19122            }
19123            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
19124            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
19125            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
19126            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
19127            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
19128            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
19129            let act = if csr {
19130                e.moe_gate_up_gelu8_dev_q8_csr(
19131                    &dev.ptr_row,
19132                    &sel_d,
19133                    &zq,
19134                    &zd,
19135                    t * n_used,
19136                    n_embd,
19137                    n_ff_exp,
19138                    n_used,
19139                    n_expert,
19140                    m.gate_exps.qtype,
19141                    m.up_exps.qtype,
19142                    m.gate_exps.row_bytes,
19143                    m.up_exps.row_bytes,
19144                )?
19145            } else {
19146                e.moe_gate_up_gelu8_dev_q8_rows(
19147                    &dev.ptr_row,
19148                    &sel_d,
19149                    &zq,
19150                    &zd,
19151                    t,
19152                    n_embd,
19153                    n_ff_exp,
19154                    n_used,
19155                    n_expert,
19156                    m.gate_exps.qtype,
19157                    m.up_exps.qtype,
19158                    m.gate_exps.row_bytes,
19159                    m.up_exps.row_bytes,
19160                )?
19161            };
19162            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
19163            let mut moe_out = e.uninit(t * n_embd)?;
19164            e.moe_down8_fma_dev_q8_rows_g(
19165                &dev.ptr_row,
19166                &sel_d,
19167                &w_d,
19168                &aq2,
19169                &ad2,
19170                &mut moe_out,
19171                t,
19172                n_ff_exp,
19173                n_embd,
19174                n_used,
19175                n_expert,
19176                m.down_exps.qtype,
19177                m.down_exps.row_bytes,
19178            )?;
19179            return Ok(moe_out);
19180        }
19181
19182        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
19183        for (i, &sx) in sel_all.iter().enumerate() {
19184            w_all[i] *= bits.per_expert_scale[sx as usize];
19185        }
19186
19187        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
19188        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
19189        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
19190        if t >= PRIME_MIN_T
19191            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
19192            && expert_dp4a_supported(m.gate_exps.qtype)
19193            && expert_dp4a_supported(m.up_exps.qtype)
19194            && expert_dp4a_supported(m.down_exps.qtype)
19195            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
19196        {
19197            let dev = m.dev_exps.as_ref().unwrap();
19198            let n_pairs = t * n_used;
19199            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
19200            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
19201            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
19202            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
19203            let pt = e.htod_i32(&pair_tok)?;
19204            let pw = e.htod(&w_all)?;
19205            let toff = e.htod_i32(&tok_off)?;
19206            let tids = e.htod_i32(&tok_ids)?;
19207            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
19208            for p in 0..n_pairs {
19209                by_ex[pair_ex[p] as usize].push(p as i32);
19210            }
19211            let mut ex_ids: Vec<i32> = Vec::new();
19212            let mut ex_off: Vec<i32> = vec![0];
19213            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
19214            for (ex, list) in by_ex.iter().enumerate() {
19215                if list.is_empty() {
19216                    continue;
19217                }
19218                ex_ids.push(ex as i32);
19219                ex_pairs.extend_from_slice(list);
19220                ex_off.push(ex_pairs.len() as i32);
19221            }
19222            let n_active = ex_ids.len();
19223            let exi = e.htod_i32(&ex_ids)?;
19224            let exo = e.htod_i32(&ex_off)?;
19225            let exp_d = e.htod_i32(&ex_pairs)?;
19226            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
19227            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
19228            // end-to-end (gelu is elementwise), one row permute before the scatter. The
19229            // ragged down k (704) needs no padding here — cublas takes any k.
19230            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
19231            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
19232            // Hopper default — see moe_f16g_gemma_on.
19233            if crate::moe_f16g_gemma_on()
19234                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
19235                && f16g_proj_ok(m.up_exps.qtype, n_embd)
19236                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
19237            {
19238                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
19239                let csr_tok_d = e.htod_i32(&csr_tok)?;
19240                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
19241                let g_csr = e.moe_f16_grouped(
19242                    &dev.ptr_row,
19243                    0,
19244                    n_expert,
19245                    &exi,
19246                    &ex_off,
19247                    &exo,
19248                    &z_f16,
19249                    &z_s,
19250                    n_embd,
19251                    n_ff_exp,
19252                    n_active,
19253                    n_pairs,
19254                    m.gate_exps.qtype,
19255                    m.gate_exps.row_bytes,
19256                )?;
19257                let u_csr = e.moe_f16_grouped(
19258                    &dev.ptr_row,
19259                    1,
19260                    n_expert,
19261                    &exi,
19262                    &ex_off,
19263                    &exo,
19264                    &z_f16,
19265                    &z_s,
19266                    n_embd,
19267                    n_ff_exp,
19268                    n_active,
19269                    n_pairs,
19270                    m.up_exps.qtype,
19271                    m.up_exps.row_bytes,
19272                )?;
19273                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
19274                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
19275                let d_csr = e.moe_f16_grouped(
19276                    &dev.ptr_row,
19277                    2,
19278                    n_expert,
19279                    &exi,
19280                    &ex_off,
19281                    &exo,
19282                    &a_f16,
19283                    &a_s,
19284                    n_ff_exp,
19285                    n_embd,
19286                    n_active,
19287                    n_pairs,
19288                    m.down_exps.qtype,
19289                    m.down_exps.row_bytes,
19290                )?;
19291                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
19292                let mut moe_out = e.uninit(t * n_embd)?;
19293                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
19294                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
19295                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
19296                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
19297                    eprintln!(
19298                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
19299                        scan(&yd),
19300                        scan(&mo)
19301                    );
19302                }
19303                return Ok(moe_out);
19304            }
19305            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
19306            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
19307            let mma = n_embd.is_multiple_of(256)
19308                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
19309            let (gate, up) = if mma {
19310                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
19311                (
19312                    e.mmq_iq_experts(
19313                        &dev.ptr_row,
19314                        0,
19315                        n_expert,
19316                        &exi,
19317                        &exo,
19318                        &exp_d,
19319                        &pt,
19320                        &z_scr,
19321                        n_embd,
19322                        n_ff_exp,
19323                        n_active,
19324                        n_pairs,
19325                        t,
19326                        m.gate_exps.qtype,
19327                        m.gate_exps.row_bytes,
19328                    )?,
19329                    e.mmq_iq_experts(
19330                        &dev.ptr_row,
19331                        1,
19332                        n_expert,
19333                        &exi,
19334                        &exo,
19335                        &exp_d,
19336                        &pt,
19337                        &z_scr,
19338                        n_embd,
19339                        n_ff_exp,
19340                        n_active,
19341                        n_pairs,
19342                        t,
19343                        m.up_exps.qtype,
19344                        m.up_exps.row_bytes,
19345                    )?,
19346                )
19347            } else {
19348                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
19349                (
19350                    e.moe_pairs_matvec_q8_dec(
19351                        &dev.ptr_row,
19352                        0,
19353                        &exi,
19354                        &exo,
19355                        &exp_d,
19356                        &pt,
19357                        &zq,
19358                        &zd,
19359                        n_embd,
19360                        n_ff_exp,
19361                        n_expert,
19362                        n_active,
19363                        n_pairs,
19364                        m.gate_exps.qtype,
19365                        m.gate_exps.row_bytes,
19366                    )?,
19367                    e.moe_pairs_matvec_q8_dec(
19368                        &dev.ptr_row,
19369                        1,
19370                        &exi,
19371                        &exo,
19372                        &exp_d,
19373                        &pt,
19374                        &zq,
19375                        &zd,
19376                        n_embd,
19377                        n_ff_exp,
19378                        n_expert,
19379                        n_active,
19380                        n_pairs,
19381                        m.up_exps.qtype,
19382                        m.up_exps.row_bytes,
19383                    )?,
19384                )
19385            };
19386            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
19387            let pself = e.htod_i32(&pair_self)?;
19388            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
19389            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
19390            // to the 256-val superblock (768) while the act quantizer's zero padding
19391            // makes every padded-k product exactly zero (weight overread bytes multiply
19392            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
19393            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
19394            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
19395            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
19396            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
19397            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
19398            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
19399            let y_down = if mma {
19400                let in_pad = n_ff_exp.div_ceil(256) * 256;
19401                let a_scr = if crate::moe_fuse_actq_on() {
19402                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
19403                } else {
19404                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
19405                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
19406                };
19407                e.mmq_iq_experts(
19408                    &dev.ptr_row,
19409                    2,
19410                    n_expert,
19411                    &exi,
19412                    &exo,
19413                    &exp_d,
19414                    &pself,
19415                    &a_scr,
19416                    in_pad,
19417                    n_embd,
19418                    n_active,
19419                    n_pairs,
19420                    n_pairs,
19421                    m.down_exps.qtype,
19422                    m.down_exps.row_bytes,
19423                )?
19424            } else {
19425                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
19426                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
19427                e.moe_pairs_matvec_q8_dec(
19428                    &dev.ptr_row,
19429                    2,
19430                    &exi,
19431                    &exo,
19432                    &exp_d,
19433                    &pself,
19434                    &aq2,
19435                    &ad2,
19436                    n_ff_exp,
19437                    n_embd,
19438                    n_expert,
19439                    n_active,
19440                    n_pairs,
19441                    m.down_exps.qtype,
19442                    m.down_exps.row_bytes,
19443                )?
19444            };
19445            let mut moe_out = e.uninit(t * n_embd)?;
19446            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
19447            return Ok(moe_out);
19448        }
19449
19450        let g_len = m.gate_exps.expert_stride;
19451        let u_len = m.up_exps.expert_stride;
19452        let d_len = m.down_exps.expert_stride;
19453        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
19454        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
19455        // the spill fallback.
19456        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
19457        let (mut sg, mut su, mut sd) = if dev.is_some() {
19458            (None, None, None)
19459        } else {
19460            (
19461                Some(e.alloc_u8_uninit(g_len)?),
19462                Some(e.alloc_u8_uninit(u_len)?),
19463                Some(e.alloc_u8_uninit(d_len)?),
19464            )
19465        };
19466        let mut moe_out = e.zeros(t * n_embd)?;
19467        for tok in 0..t {
19468            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
19469            let w = &w_all[tok * n_used..(tok + 1) * n_used];
19470            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
19471            for (j, &ex) in sel.iter().enumerate() {
19472                let ex = ex as usize;
19473                let gate = match dev {
19474                    Some(d) => m.qmatvec_view(
19475                        e,
19476                        &d.gate,
19477                        ex * g_len..(ex + 1) * g_len,
19478                        &zt,
19479                        1,
19480                        m.gate_exps.in_f,
19481                        m.gate_exps.out_f,
19482                        m.gate_exps.qtype,
19483                        m.gate_exps.row_bytes,
19484                    )?,
19485                    None => {
19486                        let sg = sg.as_mut().unwrap();
19487                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
19488                        m.qmatvec_view(
19489                            e,
19490                            sg,
19491                            0..g_len,
19492                            &zt,
19493                            1,
19494                            m.gate_exps.in_f,
19495                            m.gate_exps.out_f,
19496                            m.gate_exps.qtype,
19497                            m.gate_exps.row_bytes,
19498                        )?
19499                    }
19500                };
19501                let up = match dev {
19502                    Some(d) => m.qmatvec_view(
19503                        e,
19504                        &d.up,
19505                        ex * u_len..(ex + 1) * u_len,
19506                        &zt,
19507                        1,
19508                        m.up_exps.in_f,
19509                        m.up_exps.out_f,
19510                        m.up_exps.qtype,
19511                        m.up_exps.row_bytes,
19512                    )?,
19513                    None => {
19514                        let su = su.as_mut().unwrap();
19515                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
19516                        m.qmatvec_view(
19517                            e,
19518                            su,
19519                            0..u_len,
19520                            &zt,
19521                            1,
19522                            m.up_exps.in_f,
19523                            m.up_exps.out_f,
19524                            m.up_exps.qtype,
19525                            m.up_exps.row_bytes,
19526                        )?
19527                    }
19528                };
19529                let mut act = e.uninit(n_ff_exp)?;
19530                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
19531                let actv = act.slice(0..n_ff_exp);
19532                let y = match dev {
19533                    Some(d) => m.qmatvec_view(
19534                        e,
19535                        &d.down,
19536                        ex * d_len..(ex + 1) * d_len,
19537                        &actv,
19538                        1,
19539                        m.down_exps.in_f,
19540                        m.down_exps.out_f,
19541                        m.down_exps.qtype,
19542                        m.down_exps.row_bytes,
19543                    )?,
19544                    None => {
19545                        let sd = sd.as_mut().unwrap();
19546                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
19547                        m.qmatvec_view(
19548                            e,
19549                            sd,
19550                            0..d_len,
19551                            &actv,
19552                            1,
19553                            m.down_exps.in_f,
19554                            m.down_exps.out_f,
19555                            m.down_exps.qtype,
19556                            m.down_exps.row_bytes,
19557                        )?
19558                    }
19559                };
19560                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
19561                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
19562            }
19563        }
19564        Ok(moe_out)
19565    }
19566
19567    /// One gemma4 trunk layer (R8): x -> x_next.
19568    fn gemma4_layer(
19569        &self,
19570        e: &Engine,
19571        il: usize,
19572        layer: &crate::hybrid::HybridLayer,
19573        x: &CudaSlice<f32>,
19574        pos_d: &CudaSlice<i32>,
19575        t: usize,
19576    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19577        let n_embd = self.cfg.n_embd as usize;
19578        let eps = self.cfg.rms_eps;
19579
19580        let mut h = e.zeros(t * n_embd)?;
19581        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
19582        let Mixer::Full(fa) = &layer.mixer else {
19583            panic!("gemma4 layer {il} not full-attn")
19584        };
19585        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
19586        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
19587        let mut cur = e.zeros(t * n_embd)?;
19588        e.rms_norm(
19589            &o,
19590            layer.post_attn_norm.float_data(),
19591            &mut cur,
19592            n_embd,
19593            t,
19594            eps,
19595        )?;
19596        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
19597    }
19598
19599    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
19600    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
19601    /// layer scale — shared verbatim by the prefill, decode and verify paths.
19602    fn gemma4_layer_tail_add(
19603        &self,
19604        e: &Engine,
19605        layer: &crate::hybrid::HybridLayer,
19606        cur: &CudaSlice<f32>,
19607        x: &CudaSlice<f32>,
19608        t: usize,
19609    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19610        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
19611    }
19612
19613    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
19614    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
19615    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19616    fn gemma4_layer_tail_add_n(
19617        &self,
19618        e: &Engine,
19619        layer: &crate::hybrid::HybridLayer,
19620        cur: &CudaSlice<f32>,
19621        x: &CudaSlice<f32>,
19622        t: usize,
19623        next_norm: Option<&CudaSlice<f32>>,
19624    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
19625        let n_embd = self.cfg.n_embd as usize;
19626        let bits = layer.gemma4.as_ref().unwrap();
19627        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
19628        let mut xn = e.uninit(t * n_embd)?;
19629        match next_norm {
19630            Some(w) => {
19631                let mut hn = e.uninit(t * n_embd)?;
19632                e.add_scale_rms_norm(
19633                    &sn,
19634                    &attn_out,
19635                    bits.layer_scale,
19636                    w,
19637                    &mut xn,
19638                    &mut hn,
19639                    n_embd,
19640                    t,
19641                    self.cfg.rms_eps,
19642                )?;
19643                Ok((xn, Some(hn)))
19644            }
19645            None => {
19646                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
19647                Ok((xn, None))
19648            }
19649        }
19650    }
19651
19652    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
19653    /// norm — returns (sn, attn_out) for the closing add+scale variants.
19654    fn gemma4_layer_tail_core(
19655        &self,
19656        e: &Engine,
19657        layer: &crate::hybrid::HybridLayer,
19658        cur: &CudaSlice<f32>,
19659        x: &CudaSlice<f32>,
19660        t: usize,
19661    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19662        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
19663    }
19664
19665    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
19666    /// means `cur` is the RAW attention output and the dense entry runs
19667    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
19668    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
19669    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
19670    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
19671    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19672    fn gemma4_layer_tail_core_pn(
19673        &self,
19674        e: &Engine,
19675        layer: &crate::hybrid::HybridLayer,
19676        cur: &CudaSlice<f32>,
19677        x: &CudaSlice<f32>,
19678        t: usize,
19679        pre_norm: Option<&CudaSlice<f32>>,
19680        defer_post_norm: bool,
19681    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19682        let n_embd = self.cfg.n_embd as usize;
19683        let eps = self.cfg.rms_eps;
19684        let bits = layer.gemma4.as_ref().unwrap();
19685
19686        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
19687        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
19688        let Some(mbits) = bits.moe_bits.as_ref() else {
19689            let crate::hybrid::Ffn::Dense {
19690                ffn_gate,
19691                ffn_up,
19692                ffn_down,
19693            } = &layer.ffn
19694            else {
19695                panic!("gemma4 dense layer without Dense ffn")
19696            };
19697            let mut attn_out = e.uninit(t * n_embd)?;
19698            let mut zsh = e.uninit(t * n_embd)?;
19699            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
19700            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
19701            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19702            match pre_norm {
19703                Some(wa) if t == 1 => {
19704                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
19705                        cur,
19706                        wa,
19707                        x,
19708                        bits.ffn_norm.float_data(),
19709                        &mut attn_out,
19710                        &mut zsh,
19711                        n_embd,
19712                        t,
19713                        eps,
19714                    )?);
19715                }
19716                Some(wa) => e.rms_pre_add_rms_norm(
19717                    cur,
19718                    wa,
19719                    x,
19720                    bits.ffn_norm.float_data(),
19721                    &mut attn_out,
19722                    &mut zsh,
19723                    n_embd,
19724                    t,
19725                    eps,
19726                )?,
19727                None => e.add_rms_norm(
19728                    cur,
19729                    x,
19730                    bits.ffn_norm.float_data(),
19731                    &mut attn_out,
19732                    &mut zsh,
19733                    n_embd,
19734                    t,
19735                    eps,
19736                )?,
19737            }
19738            let n_ff = ffn_gate.out_features();
19739            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
19740            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
19741            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
19742            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
19743            // rescue segment C — the megakernel front is closed for the dense tail.
19744            let (gate, up) = if t == 1 {
19745                let (zq, zd) = match zpair {
19746                    Some(p) => p,
19747                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
19748                };
19749                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
19750                    Some(p) => p,
19751                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
19752                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
19753                        Some(p) => p,
19754                        None => (
19755                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
19756                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
19757                        ),
19758                    },
19759                }
19760            } else {
19761                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
19762                // launch for the verify's gate+up — the up segment's blocks fill SMs as
19763                // the gate segment drains (the launch-tail mechanism behind the b-tier
19764                // plateau; first positive after six falsified in-kernel variants).
19765                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19766                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
19767                let fused = if f2b {
19768                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
19769                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
19770                } else {
19771                    None
19772                };
19773                match fused {
19774                    Some(p) => p,
19775                    None => {
19776                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
19777                        e.mmq_act_begin();
19778                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
19779                    }
19780                }
19781            };
19782            let mut act = e.uninit(t * n_ff)?;
19783            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
19784            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
19785            let f0 = if e.uses_q8_1_fast(ffn_down) {
19786                let upv = e.view(&up, t * n_ff);
19787                let up_all = upv.slice(0..t * n_ff);
19788                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
19789                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
19790            } else {
19791                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
19792                e.matmul(ffn_down, &act, t)?
19793            };
19794            if defer_post_norm {
19795                return Ok((f0, attn_out));
19796            }
19797            let mut sn = e.uninit(t * n_embd)?;
19798            e.rms_norm(
19799                &f0,
19800                bits.post_ffw_norm.float_data(),
19801                &mut sn,
19802                n_embd,
19803                t,
19804                eps,
19805            )?;
19806            return Ok((sn, attn_out));
19807        };
19808
19809        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
19810        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
19811        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
19812        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
19813        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
19814        let mut attn_out = e.uninit(t * n_embd)?;
19815        let mut router_in = e.uninit(t * n_embd)?;
19816        let fast_moe = match &layer.ffn {
19817            crate::hybrid::Ffn::Moe(m) => {
19818                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
19819                    && expert_dp4a_supported(m.gate_exps.qtype)
19820                    && expert_dp4a_supported(m.up_exps.qtype)
19821                    && expert_dp4a_supported(m.down_exps.qtype)
19822                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
19823            }
19824            _ => false,
19825        };
19826        let q8z = t < PRIME_MIN_T && fast_moe;
19827        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
19828            let (z0, m2) = e.add_rms_norm3_q8z(
19829                cur,
19830                x,
19831                bits.ffn_norm.float_data(),
19832                &mbits.router_scale_pre,
19833                mbits.pre_ffw_norm_2.float_data(),
19834                &mut attn_out,
19835                &mut router_in,
19836                n_embd,
19837                t,
19838                eps,
19839            )?;
19840            (None, Some(z0), Some(m2))
19841        } else {
19842            let mut zsh = e.uninit(t * n_embd)?;
19843            let mut moe_in = e.uninit(t * n_embd)?;
19844            e.add_rms_norm3(
19845                cur,
19846                x,
19847                bits.ffn_norm.float_data(),
19848                &mbits.router_scale_pre,
19849                mbits.pre_ffw_norm_2.float_data(),
19850                &mut attn_out,
19851                &mut zsh,
19852                &mut router_in,
19853                &mut moe_in,
19854                n_embd,
19855                t,
19856                eps,
19857            )?;
19858            (Some((zsh, moe_in)), None, None)
19859        };
19860        let attn_out2 = attn_out;
19861        #[allow(unused_variables)]
19862        let attn_out = &attn_out2;
19863        let n_ff = mbits.shared_gate.out_features();
19864        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
19865            if t == 1 {
19866                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
19867                    Some(p) => p,
19868                    None => match e.matmul_nvfp4_fused2(
19869                        &mbits.shared_gate,
19870                        &mbits.shared_up,
19871                        zq,
19872                        zd,
19873                        1,
19874                    )? {
19875                        Some(p) => p,
19876                        None => {
19877                            let h0 = e.zeros(0)?;
19878                            (
19879                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
19880                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
19881                            )
19882                        }
19883                    },
19884                }
19885            } else {
19886                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
19887                let h0 = e.zeros(0)?;
19888                (
19889                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
19890                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
19891                )
19892            }
19893        } else {
19894            let (zsh, _) = zsh_f32.as_ref().unwrap();
19895            (
19896                e.matmul(&mbits.shared_gate, zsh, t)?,
19897                e.matmul(&mbits.shared_up, zsh, t)?,
19898            )
19899        };
19900        let mut act = e.uninit(t * n_ff)?;
19901        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
19902        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
19903        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
19904            panic!("gemma4 layer not MoE")
19905        };
19906        let moe0 = match (&moe_q8, &zsh_f32) {
19907            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
19908            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
19909            _ => unreachable!(),
19910        };
19911        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
19912        let mut mlp = e.uninit(t * n_embd)?;
19913        let mut moe = e.uninit(t * n_embd)?;
19914        e.rms_norm2x(
19915            &mlp0,
19916            &moe0,
19917            mbits.post_ffw_norm_1.float_data(),
19918            mbits.post_ffw_norm_2.float_data(),
19919            &mut mlp,
19920            &mut moe,
19921            n_embd,
19922            t,
19923            eps,
19924        )?;
19925
19926        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
19927        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
19928        let mut sum = e.uninit(t * n_embd)?;
19929        let mut sn = e.uninit(t * n_embd)?;
19930        e.add_rms_norm(
19931            &mlp,
19932            &moe,
19933            bits.post_ffw_norm.float_data(),
19934            &mut sum,
19935            &mut sn,
19936            n_embd,
19937            t,
19938            eps,
19939        )?;
19940        Ok((sn, attn_out2))
19941    }
19942
19943    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
19944    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
19945    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
19946    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
19947    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
19948    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
19949    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
19950    /// decode == verify == graph parity holds by construction at either seam value.
19951    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
19952    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19953    pub(crate) fn gemma4_layer_tail_add_nq_pn(
19954        &self,
19955        e: &Engine,
19956        layer: &crate::hybrid::HybridLayer,
19957        o: &CudaSlice<f32>,
19958        x: &CudaSlice<f32>,
19959        t: usize,
19960        next_norm: Option<&CudaSlice<f32>>,
19961    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
19962    {
19963        let n_embd = self.cfg.n_embd as usize;
19964        let eps = self.cfg.rms_eps;
19965        let bits = layer.gemma4.as_ref().unwrap();
19966        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
19967            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
19968                e,
19969                layer,
19970                o,
19971                x,
19972                t,
19973                Some(layer.post_attn_norm.float_data()),
19974                true,
19975            )?;
19976            let mut xn = e.uninit(t * n_embd)?;
19977            return match next_norm {
19978                Some(w) => {
19979                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
19980                        &f0,
19981                        bits.post_ffw_norm.float_data(),
19982                        &attn_out,
19983                        bits.layer_scale,
19984                        w,
19985                        &mut xn,
19986                        n_embd,
19987                        t,
19988                        eps,
19989                    )?;
19990                    Ok((xn, Some(pair)))
19991                }
19992                None => {
19993                    let mut sn = e.uninit(t * n_embd)?;
19994                    e.rms_norm(
19995                        &f0,
19996                        bits.post_ffw_norm.float_data(),
19997                        &mut sn,
19998                        n_embd,
19999                        t,
20000                        eps,
20001                    )?;
20002                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
20003                    Ok((xn, None))
20004                }
20005            };
20006        }
20007        let mut cur = e.uninit(t * n_embd)?;
20008        e.rms_norm(
20009            o,
20010            layer.post_attn_norm.float_data(),
20011            &mut cur,
20012            n_embd,
20013            t,
20014            eps,
20015        )?;
20016        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
20017    }
20018
20019    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20020    pub(crate) fn gemma4_layer_tail_add_nq(
20021        &self,
20022        e: &Engine,
20023        layer: &crate::hybrid::HybridLayer,
20024        cur: &CudaSlice<f32>,
20025        x: &CudaSlice<f32>,
20026        t: usize,
20027        next_norm: Option<&CudaSlice<f32>>,
20028    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
20029    {
20030        let n_embd = self.cfg.n_embd as usize;
20031        let bits = layer.gemma4.as_ref().unwrap();
20032        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
20033        let mut xn = e.uninit(t * n_embd)?;
20034        match next_norm {
20035            Some(w) => {
20036                let pair = e.add_scale_rms_norm_q8_1(
20037                    &sn,
20038                    &attn_out,
20039                    bits.layer_scale,
20040                    w,
20041                    &mut xn,
20042                    n_embd,
20043                    t,
20044                    self.cfg.rms_eps,
20045                )?;
20046                Ok((xn, Some(pair)))
20047            }
20048            None => {
20049                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
20050                Ok((xn, None))
20051            }
20052        }
20053    }
20054
20055    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
20056    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
20057    fn gemma4_forward(
20058        &self,
20059        e: &Engine,
20060        tokens: &[u32],
20061        last_only: bool,
20062    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
20063        // E4B routes to its own forward regardless of the caller's entry point (forward /
20064        // forward_last / prime paths all funnel here for gemma4).
20065        if self.is_gemma4_e4b() {
20066            return self.gemma4_e4b_forward(e, tokens, last_only);
20067        }
20068        let n_embd = self.cfg.n_embd as usize;
20069        let t = tokens.len();
20070        let pos: Vec<i32> = (0..t as i32).collect();
20071        let pos_d = e.htod_i32(&pos)?;
20072
20073        let mut x = self.embed(e, tokens)?;
20074        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20075        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
20076        // the bring-up bisect vs llama-eval-callback node stats.
20077        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
20078        let stat =
20079            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
20080                let h = e.dtoh(x)?;
20081                let bad = h.iter().filter(|v| !v.is_finite()).count();
20082                let mx = h
20083                    .iter()
20084                    .filter(|v| v.is_finite())
20085                    .fold(0.0f32, |m, v| m.max(v.abs()));
20086                eprintln!(
20087                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
20088                    &h[..3]
20089                );
20090                Ok(())
20091            };
20092        if probe {
20093            stat(e, &x, "embed")?;
20094        }
20095        for (il, layer) in self.layers.iter().enumerate() {
20096            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
20097            if probe {
20098                stat(e, &x, &format!("L{il}"))?;
20099            }
20100        }
20101        let mut hn = e.zeros(t * n_embd)?;
20102        e.rms_norm(
20103            &x,
20104            self.output_norm.float_data(),
20105            &mut hn,
20106            n_embd,
20107            t,
20108            self.cfg.rms_eps,
20109        )?;
20110        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20111        let n_vocab = self.output.out_features();
20112        let logits = if last_only {
20113            let hv = e.view(&hn, t * n_embd);
20114            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
20115            let mut hlast = e.zeros(n_embd)?;
20116            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
20117            let mut ld = e.matmul(&self.output, &hlast, 1)?;
20118            e.softcap(&mut ld, cap, n_vocab)?;
20119            self.gemma4_suppress(e, &mut ld, 1)?;
20120            e.dtoh(&ld)?
20121        } else {
20122            let mut ld = e.matmul(&self.output, &hn, t)?;
20123            e.softcap(&mut ld, cap, t * n_vocab)?;
20124            self.gemma4_suppress(e, &mut ld, t)?;
20125            e.dtoh(&ld)?
20126        };
20127        Ok(logits)
20128    }
20129
20130    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
20131    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
20132    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
20133    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
20134    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
20135    pub(crate) fn gemma4_prime(
20136        &self,
20137        e: &Engine,
20138        tokens: &[u32],
20139        cache: &mut Cache,
20140        overlay: Option<&crate::vision::EmbedOverlay>,
20141    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20142        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
20143        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
20144        // whole worker process on this line. The worker now primes gemma4 monolithically and
20145        // routes continuation suffixes tokenwise; this is the per-request backstop.
20146        if cache.pos != 0 {
20147            return Err(
20148                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
20149                        — prime the full prompt in one call or decode tokenwise"
20150                    .into(),
20151            );
20152        }
20153        let n_embd = self.cfg.n_embd as usize;
20154        let eps = self.cfg.rms_eps;
20155        let t = tokens.len();
20156        let pos: Vec<i32> = (0..t as i32).collect();
20157        let pos_d = e.htod_i32(&pos)?;
20158        let mut x = self.embed(e, tokens)?;
20159        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20160        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
20161        // sqrt(n_embd) text scale — the reference scales token batches only
20162        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
20163        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
20164        // bidirectional within itself, causal+SWA everywhere else, matching the
20165        // reference's llama_set_causal_attn(false) image batch exactly.
20166        let island: Option<CudaSlice<i32>> = match overlay {
20167            Some(ov) => {
20168                // The residency law reaches this arm too (lane/glm53-vision-ppn): the splice
20169                // arithmetic below differs from the shared helper on purpose (post-scale
20170                // placement + island ids), but the pointer it reads obeys the same rule.
20171                ov.require_resident(e)?;
20172                let mut span_id = vec![-1i32; t];
20173                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
20174                    if pos + n_rows > t {
20175                        return Err(format!(
20176                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
20177                            pos + n_rows
20178                        )
20179                        .into());
20180                    }
20181                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
20182                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
20183                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
20184                        *s = i as i32;
20185                    }
20186                }
20187                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
20188                // keep the plain causal mask. Exists only so the decisive probe can show
20189                // the island mask itself changes the answer; never on in serving.
20190                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
20191                    None
20192                } else {
20193                    Some(e.htod_i32(&span_id)?)
20194                }
20195            }
20196            None => None,
20197        };
20198        for (il, layer) in self.layers.iter().enumerate() {
20199            let mut h = e.zeros(t * n_embd)?;
20200            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
20201            let Mixer::Full(fa) = &layer.mixer else {
20202                panic!("gemma4 layer not full-attn")
20203            };
20204            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
20205            if trace {
20206                let v = e.dtoh(&h)?;
20207                let nan = v.iter().filter(|x| x.is_nan()).count();
20208                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
20209            }
20210            let o =
20211                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
20212            if trace {
20213                let v = e.dtoh(&o)?;
20214                let nan = v.iter().filter(|x| x.is_nan()).count();
20215                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
20216            }
20217            let mut cur = e.zeros(t * n_embd)?;
20218            e.rms_norm(
20219                &o,
20220                layer.post_attn_norm.float_data(),
20221                &mut cur,
20222                n_embd,
20223                t,
20224                eps,
20225            )?;
20226            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
20227            self.dflash_tap(e, cache, il, &x, t)?;
20228            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
20229            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
20230                let h = e.dtoh(&x)?;
20231                let nan = h.iter().filter(|v| v.is_nan()).count();
20232                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
20233                eprintln!(
20234                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
20235                    h.len()
20236                );
20237                if nan > 0 {
20238                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
20239                }
20240            }
20241        }
20242        cache.pos += t;
20243        let hiddens = e.clone_dtod(&x)?;
20244        let xv = e.view(&x, t * n_embd);
20245        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
20246        let mut h_seed = e.zeros(n_embd)?;
20247        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
20248        let mut hn = e.uninit(n_embd)?;
20249        e.rms_norm(
20250            &h_seed,
20251            self.output_norm.float_data(),
20252            &mut hn,
20253            n_embd,
20254            1,
20255            eps,
20256        )?;
20257        let mut ld = e.matmul(&self.output, &hn, 1)?;
20258        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20259        e.softcap(&mut ld, cap, self.output.out_features())?;
20260        self.gemma4_suppress(e, &mut ld, 1)?;
20261        let logits = e.dtoh(&ld)?;
20262        Ok((logits, h_seed, hiddens))
20263    }
20264
20265    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
20266    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
20267    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
20268    /// fused norm emits q8 directly — the f32 h never materializes).
20269    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
20270    fn gemma4_decode_attn(
20271        &self,
20272        e: &Engine,
20273        fa: &crate::hybrid::FullAttnLayer,
20274        il: usize,
20275        hq: &CudaSlice<i8>,
20276        hdq: &CudaSlice<f32>,
20277        pos_d: &CudaSlice<i32>,
20278        cache: &mut Cache,
20279    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20280        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20281        let eps = self.cfg.rms_eps;
20282        let aux = self.gemma4_aux.as_ref().unwrap();
20283        let ones = aux.ones(e);
20284        #[cfg(debug_assertions)]
20285        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
20286        let (hq, hdq) = (hq, hdq);
20287        let h0 = e.zeros(0)?;
20288        let h = &h0;
20289        let (q0, k0, v0) = if swa {
20290            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
20291                Some(t3) => t3,
20292                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
20293                // match — fuse the uniform (q,k) pair and take v as its own single.
20294                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
20295                    Some((q0, k0)) => {
20296                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, h, 1)?;
20297                        (q0, k0, v0)
20298                    }
20299                    None => (
20300                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
20301                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
20302                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
20303                    ),
20304                },
20305            }
20306        } else {
20307            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
20308                Some(p) => p,
20309                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
20310                    Some(p) => p,
20311                    None => (
20312                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
20313                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
20314                    ),
20315                },
20316            };
20317            let v0 = e.clone_dtod(&k0)?;
20318            (q0, k0, v0)
20319        };
20320        let mut q = e.uninit(nh * hd)?;
20321        let mut k = e.uninit(nkv * hd)?;
20322        let mut v = e.uninit(nkv * hd)?;
20323        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
20324        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
20325        let ff = if swa {
20326            None
20327        } else {
20328            Some(
20329                aux.rope_freqs(e)
20330                    .expect("gemma4 global rope needs rope_freqs.weight"),
20331            )
20332        };
20333        #[cfg(debug_assertions)]
20334        if let Some(ff) = ff {
20335            crate::debug_assert_tensor_stream_device(
20336                ff,
20337                &e.stream(),
20338                "gemma4_decode_attn.rope_freqs",
20339            );
20340        }
20341        let kvl = cache.kv[il].as_mut().unwrap();
20342        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
20343        if crate::Engine::qkv_append_on() {
20344            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
20345            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
20346            // twin of the dc fold — bit-identical bodies, one launch per layer.
20347            e.rms_norm_qkv_rope_append(
20348                &q0,
20349                &k0,
20350                &v0,
20351                fa.q_norm.float_data(),
20352                fa.k_norm.float_data(),
20353                ones,
20354                &mut q,
20355                &mut k,
20356                &mut v,
20357                hd,
20358                self.gemma4_rope_dims(il),
20359                nh,
20360                nkv,
20361                pos_d,
20362                nh,
20363                nkv,
20364                base,
20365                1.0,
20366                ff,
20367                eps,
20368                &mut kvl.k,
20369                &mut kvl.v,
20370                kvl.len,
20371                kvl.k_tok_bytes,
20372                kvl.v_tok_bytes,
20373                kv_fp8,
20374            )?;
20375        } else {
20376            e.rms_norm_qkv_rope(
20377                &q0,
20378                &k0,
20379                &v0,
20380                fa.q_norm.float_data(),
20381                fa.k_norm.float_data(),
20382                ones,
20383                &mut q,
20384                &mut k,
20385                &mut v,
20386                hd,
20387                self.gemma4_rope_dims(il),
20388                nh,
20389                nkv,
20390                pos_d,
20391                nh,
20392                nkv,
20393                base,
20394                1.0,
20395                ff,
20396                eps,
20397            )?;
20398            e.append_kv_quantized(
20399                &k,
20400                &v,
20401                &mut kvl.k,
20402                &mut kvl.v,
20403                kvl.len,
20404                kvl.kv_dim_k,
20405                kvl.kv_dim_v,
20406                kvl.k_tok_bytes,
20407                kvl.v_tok_bytes,
20408                kv_fp8,
20409            )?;
20410        }
20411        kvl.len += 1;
20412        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
20413        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
20414        // positional). Globals attend the full history.
20415        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20416        let mut attn = e.uninit(nh * hd)?;
20417        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
20418        if !swa
20419            && hd == 512
20420            && kvl.len >= crate::fa512_min_tkv()
20421            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20422        {
20423            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
20424            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
20425            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
20426            let base = kvl.len as i32;
20427            e.i32_set_k(&mut kvl.len_d, base)?;
20428            e.fa_decode_rows(
20429                &q,
20430                &kp,
20431                &vp,
20432                &mut attn,
20433                hd,
20434                nh,
20435                nkv,
20436                kvl.len - 1,
20437                1,
20438                scale,
20439                kvl.k_tok_bytes,
20440                kvl.v_tok_bytes,
20441                Some((&kvl.len_d, -1)),
20442                false,
20443                false,
20444                None,
20445            )?;
20446            return e.matmul(&fa.wo, &attn, 1);
20447        }
20448        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
20449        if swa
20450            && kvl.len > win
20451            && hd == 256
20452            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20453        {
20454            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
20455            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
20456            let base = kvl.len as i32;
20457            e.i32_set_k(&mut kvl.len_d, base)?;
20458            e.fa_decode_rows_w(
20459                &q,
20460                &kp,
20461                &vp,
20462                &mut attn,
20463                hd,
20464                nh,
20465                nkv,
20466                &kvl.len_d,
20467                -1,
20468                1,
20469                scale,
20470                win,
20471                kvl.k_tok_bytes,
20472                kvl.v_tok_bytes,
20473                None,
20474            )?;
20475            return e.matmul(&fa.wo, &attn, 1);
20476        }
20477        let (off_tok, t_kv) = if swa && kvl.len > win {
20478            (kvl.len - win, win)
20479        } else {
20480            (0, kvl.len)
20481        };
20482        let k_view = e.view_u8_range(
20483            &kvl.k,
20484            off_tok * kvl.k_tok_bytes,
20485            (off_tok + t_kv) * kvl.k_tok_bytes,
20486        );
20487        let v_view = e.view_u8_range(
20488            &kvl.v,
20489            off_tok * kvl.v_tok_bytes,
20490            (off_tok + t_kv) * kvl.v_tok_bytes,
20491        );
20492        e.fa_decode_kvmod(
20493            &q,
20494            &k_view,
20495            &v_view,
20496            &mut attn,
20497            hd,
20498            nh,
20499            nkv,
20500            t_kv,
20501            scale,
20502            kvl.k_tok_bytes,
20503            kvl.v_tok_bytes,
20504            swa && crate::Engine::wkv_on(),
20505        )?;
20506        e.matmul(&fa.wo, &attn, 1)
20507    }
20508
20509    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
20510    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
20511    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
20512    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
20513    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
20514    /// in-graph; the driver gates).
20515    #[allow(clippy::too_many_arguments)]
20516    pub fn gemma4_decode_step_dc(
20517        &self,
20518        e: &Engine,
20519        token_d: &CudaSlice<u32>,
20520        pos_d: &mut CudaSlice<i32>,
20521        embd_gpu: &CudaSlice<u8>,
20522        embd_qt: i32,
20523        embd_rb: usize,
20524        cache: &mut Cache,
20525        n_vocab: usize,
20526        cap_bucket_max: Option<(usize, usize)>,
20527    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
20528        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
20529        self.gemma4_decode_step_dc_into(
20530            e,
20531            token_d,
20532            pos_d,
20533            embd_gpu,
20534            embd_qt,
20535            embd_rb,
20536            cache,
20537            n_vocab,
20538            cap_bucket_max,
20539            &mut tok_out,
20540        )?;
20541        Ok(tok_out)
20542    }
20543
20544    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
20545    /// every replay; pass `token_d` itself for the self-feeding graph loop).
20546    #[allow(clippy::too_many_arguments)]
20547    pub fn gemma4_decode_step_dc_into(
20548        &self,
20549        e: &Engine,
20550        token_d: &CudaSlice<u32>,
20551        pos_d: &mut CudaSlice<i32>,
20552        embd_gpu: &CudaSlice<u8>,
20553        embd_qt: i32,
20554        embd_rb: usize,
20555        cache: &mut Cache,
20556        n_vocab: usize,
20557        cap_bucket_max: Option<(usize, usize)>,
20558        tok_out: &mut CudaSlice<u32>,
20559    ) -> Result<(), Box<dyn std::error::Error>> {
20560        let n_embd = self.cfg.n_embd as usize;
20561        let eps = self.cfg.rms_eps;
20562        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
20563        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20564        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20565        let n_layers = self.layers.len();
20566        for (il, layer) in self.layers.iter().enumerate() {
20567            let (hq, hdq) = match h_carry.take() {
20568                Some(p) => p,
20569                None => {
20570                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
20571                }
20572            };
20573            let Mixer::Full(fa) = &layer.mixer else {
20574                panic!("gemma4 layer {il} not full-attn")
20575            };
20576            let o =
20577                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
20578            let next_norm = if il + 1 < n_layers {
20579                Some(self.layers[il + 1].attn_norm.float_data())
20580            } else {
20581                None
20582            };
20583            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20584            x = xn;
20585            h_carry = hn;
20586        }
20587        let mut hn = e.uninit(n_embd)?;
20588        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20589        let mut logits = e.matmul(&self.output, &hn, 1)?;
20590        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
20591        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
20592        e.inc_seqlen(pos_d)?;
20593        if cap_bucket_max.is_none() {
20594            cache.pos += 1;
20595        }
20596        Ok(())
20597    }
20598
20599    // Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
20600    // every buffer the step produces per token lives here, allocated ONCE pre-capture, so
20601    // the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
20602    // osrt 2026-07-23). Sized for the model's max per-layer shapes.
20603
20604    /// Build the slot set (call OUTSIDE any capture).
20605    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
20606        let n_embd = self.cfg.n_embd as usize;
20607        let n_vocab = self.output.out_features();
20608        let n_layers = self.layers.len();
20609        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
20610        for il in 0..n_layers {
20611            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
20612            qmax = qmax.max(nh * hd);
20613            kvmax = kvmax.max(nkv * hd);
20614            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
20615                ffmax = ffmax.max(ffn_gate.out_features());
20616            }
20617        }
20618        Ok(G4DcSlots {
20619            x: e.uninit(n_embd)?,
20620            xn: e.uninit(n_embd)?,
20621            cur: e.uninit(n_embd)?,
20622            hq: e.alloc_i8_uninit(n_embd)?,
20623            hd_: e.uninit(n_embd / 32)?,
20624            q0: e.uninit(qmax)?,
20625            k0: e.uninit(kvmax)?,
20626            v0: e.uninit(kvmax)?,
20627            q: e.uninit(qmax)?,
20628            k: e.uninit(kvmax)?,
20629            v: e.uninit(kvmax)?,
20630            attn: e.uninit(qmax)?,
20631            o: e.uninit(n_embd)?,
20632            attn_out: e.uninit(n_embd)?,
20633            zsh: e.uninit(n_embd)?,
20634            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
20635            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
20636            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
20637            zd: e.uninit(n_embd.max(qmax) / 32)?,
20638            gate: e.uninit(ffmax)?,
20639            up: e.uninit(ffmax)?,
20640            act: e.uninit(ffmax)?,
20641            actq: e.alloc_i8_uninit(ffmax)?,
20642            actd: e.uninit(ffmax / 32)?,
20643            f0: e.uninit(n_embd)?,
20644            sn: e.uninit(n_embd)?,
20645            hn: e.uninit(n_embd)?,
20646            logits: e.uninit(n_vocab)?,
20647        })
20648    }
20649
20650    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
20651    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
20652    fn g4_matvec_m1_into(
20653        &self,
20654        e: &Engine,
20655        w: &crate::model::GpuTensor,
20656        aq: &CudaSlice<i8>,
20657        ad: &CudaSlice<f32>,
20658        y: &mut CudaSlice<f32>,
20659    ) -> Result<(), Box<dyn std::error::Error>> {
20660        use crate::model::GpuTensor;
20661        let (bytes, qtype, row_bytes, scale, rp) = match w {
20662            GpuTensor::Quant {
20663                bytes,
20664                qtype,
20665                row_bytes,
20666                scale,
20667                rp,
20668                ..
20669            } => (bytes, *qtype, *row_bytes, *scale, *rp),
20670            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
20671        };
20672        let (mbytes, mrp) = match w {
20673            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
20674            _ => (bytes, rp),
20675        };
20676        e.qmatvec_mmvq_into(
20677            mbytes,
20678            aq,
20679            ad,
20680            1,
20681            w.in_features(),
20682            w.out_features(),
20683            qtype,
20684            row_bytes,
20685            scale,
20686            mrp,
20687            y,
20688        )
20689    }
20690
20691    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
20692    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
20693    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
20694    #[allow(clippy::too_many_arguments)]
20695    pub fn gemma4_decode_step_dc_slotted(
20696        &self,
20697        e: &Engine,
20698        token_d: &CudaSlice<u32>,
20699        pos_d: &mut CudaSlice<i32>,
20700        embd_gpu: &CudaSlice<u8>,
20701        embd_qt: i32,
20702        embd_rb: usize,
20703        cache: &mut Cache,
20704        n_vocab: usize,
20705        cap_bucket_max: Option<(usize, usize)>,
20706        sl: &mut G4DcSlots,
20707        tok_out: &mut CudaSlice<u32>,
20708        ring: Option<(&mut CudaSlice<u32>, usize)>,
20709    ) -> Result<(), Box<dyn std::error::Error>> {
20710        let n_embd = self.cfg.n_embd as usize;
20711        let eps = self.cfg.rms_eps;
20712        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
20713        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
20714        let n_layers = self.layers.len();
20715        let mut has_carry = false;
20716        for il in 0..n_layers {
20717            if !has_carry {
20718                e.rms_norm_q8_1_into(
20719                    &sl.x,
20720                    self.layers[il].attn_norm.float_data(),
20721                    n_embd,
20722                    1,
20723                    eps,
20724                    &mut sl.hq,
20725                    &mut sl.hd_,
20726                )?;
20727            }
20728            has_carry = true;
20729            let layer = &self.layers[il];
20730            let Mixer::Full(fa) = &layer.mixer else {
20731                panic!("gemma4 layer {il} not full-attn")
20732            };
20733            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
20734            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
20735            // the standalone norm only survives on the unfused seam arm.
20736            if !Engine::g4_pnfold_on() {
20737                e.rms_norm(
20738                    &sl.o,
20739                    layer.post_attn_norm.float_data(),
20740                    &mut sl.cur,
20741                    n_embd,
20742                    1,
20743                    eps,
20744                )?;
20745            }
20746            let next_norm = if il + 1 < n_layers {
20747                Some(self.layers[il + 1].attn_norm.float_data())
20748            } else {
20749                None
20750            };
20751            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
20752            std::mem::swap(&mut sl.x, &mut sl.xn);
20753        }
20754        e.rms_norm(
20755            &sl.x,
20756            self.output_norm.float_data(),
20757            &mut sl.hn,
20758            n_embd,
20759            1,
20760            eps,
20761        )?;
20762        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
20763        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
20764        {
20765            let (zq, zd) = (&sl.zq, &sl.zd);
20766            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
20767            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
20768            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
20769        }
20770        self.gemma4_suppress(e, &mut sl.logits, 1)?;
20771        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
20772        if let Some((ring, base)) = ring {
20773            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
20774            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
20775            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
20776            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
20777        }
20778        e.inc_seqlen(pos_d)?;
20779        if cap_bucket_max.is_none() {
20780            cache.pos += 1;
20781        }
20782        Ok(())
20783    }
20784
20785    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
20786    #[allow(clippy::too_many_arguments)]
20787    fn gemma4_decode_attn_dc_slotted(
20788        &self,
20789        e: &Engine,
20790        fa: &crate::hybrid::FullAttnLayer,
20791        il: usize,
20792        pos_d: &CudaSlice<i32>,
20793        cache: &mut Cache,
20794        cap_bucket_max: Option<(usize, usize)>,
20795        sl: &mut G4DcSlots,
20796    ) -> Result<(), Box<dyn std::error::Error>> {
20797        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20798        let eps = self.cfg.rms_eps;
20799        let aux = self.gemma4_aux.as_ref().unwrap();
20800        let ones = aux.ones(e);
20801        #[cfg(debug_assertions)]
20802        crate::debug_assert_tensor_stream_device(
20803            ones,
20804            &e.stream(),
20805            "gemma4_decode_attn_dc_slotted.ones",
20806        );
20807        {
20808            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
20809            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
20810            if swa {
20811                if !e.matmul_q4_fused3_into(
20812                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
20813                )? {
20814                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
20815                    // (q,k) pair, v through the generic m1 slot matvec — the same two
20816                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
20817                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
20818                    {
20819                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
20820                    } else {
20821                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
20822                    }
20823                }
20824            } else {
20825                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
20826                    && !e
20827                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
20828                {
20829                    return Err("slotted step: fused2 unavailable".into());
20830                }
20831                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
20832                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
20833            }
20834        }
20835        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
20836        // kernel-for-kernel (graph stream-identity gate).
20837        let ff = if swa {
20838            None
20839        } else {
20840            Some(
20841                aux.rope_freqs(e)
20842                    .expect("gemma4 global rope needs rope_freqs.weight"),
20843            )
20844        };
20845        #[cfg(debug_assertions)]
20846        if let Some(ff) = ff {
20847            crate::debug_assert_tensor_stream_device(
20848                ff,
20849                &e.stream(),
20850                "gemma4_decode_attn_dc_slotted.rope_freqs",
20851            );
20852        }
20853        let kvl = cache.kv[il].as_mut().unwrap();
20854        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
20855        if crate::Engine::qkv_append_on() {
20856            // append fold (2026-07-23): mirrors dc_into.
20857            e.rms_norm_qkv_rope_append_dc(
20858                &sl.q0,
20859                &sl.k0,
20860                &sl.v0,
20861                fa.q_norm.float_data(),
20862                fa.k_norm.float_data(),
20863                ones,
20864                &mut sl.q,
20865                &mut sl.k,
20866                &mut sl.v,
20867                hd,
20868                self.gemma4_rope_dims(il),
20869                nh,
20870                nkv,
20871                pos_d,
20872                nh,
20873                nkv,
20874                base,
20875                1.0,
20876                ff,
20877                eps,
20878                &mut kvl.k,
20879                &mut kvl.v,
20880                &kvl.len_d,
20881                kvl.k_tok_bytes,
20882                kvl.v_tok_bytes,
20883                kv_fp8,
20884            )?;
20885        } else {
20886            e.rms_norm_qkv_rope(
20887                &sl.q0,
20888                &sl.k0,
20889                &sl.v0,
20890                fa.q_norm.float_data(),
20891                fa.k_norm.float_data(),
20892                ones,
20893                &mut sl.q,
20894                &mut sl.k,
20895                &mut sl.v,
20896                hd,
20897                self.gemma4_rope_dims(il),
20898                nh,
20899                nkv,
20900                pos_d,
20901                nh,
20902                nkv,
20903                base,
20904                1.0,
20905                ff,
20906                eps,
20907            )?;
20908            e.append_kv_quantized_dc(
20909                &sl.k,
20910                &sl.v,
20911                &mut kvl.k,
20912                &mut kvl.v,
20913                &kvl.len_d,
20914                kvl.kv_dim_k,
20915                kvl.kv_dim_v,
20916                kvl.k_tok_bytes,
20917                kvl.v_tok_bytes,
20918                kv_fp8,
20919            )?;
20920        }
20921        e.inc_seqlen(&mut kvl.len_d)?;
20922        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
20923        let k_view = e.view_u8(&kvl.k, kvl.k.len());
20924        let v_view = e.view_u8(&kvl.v, kvl.v.len());
20925        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
20926        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20927        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
20928        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
20929        // the dc_into arm branch-for-branch (stream gate).
20930        let mut fa_q8 = false;
20931        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
20932            e.fa_decode_rows(
20933                &sl.q,
20934                &k_view,
20935                &v_view,
20936                &mut sl.attn,
20937                hd,
20938                nh,
20939                nkv,
20940                b_glob - 1,
20941                1,
20942                scale,
20943                kvl.k_tok_bytes,
20944                kvl.v_tok_bytes,
20945                Some((&kvl.len_d, -1)),
20946                false,
20947                false,
20948                Some((&mut sl.zq, &mut sl.zd)),
20949            )?;
20950            fa_q8 = true;
20951        } else if swa && b_swa > win && hd == 256 && rows_on {
20952            e.fa_decode_rows_w(
20953                &sl.q,
20954                &k_view,
20955                &v_view,
20956                &mut sl.attn,
20957                hd,
20958                nh,
20959                nkv,
20960                &kvl.len_d,
20961                -1,
20962                1,
20963                scale,
20964                win,
20965                kvl.k_tok_bytes,
20966                kvl.v_tok_bytes,
20967                Some((&mut sl.zq, &mut sl.zd)),
20968            )?;
20969            fa_q8 = true;
20970        } else {
20971            let b = if swa { b_swa } else { b_glob };
20972            e.fa_decode_dc(
20973                &sl.q,
20974                &k_view,
20975                &v_view,
20976                &mut sl.attn,
20977                hd,
20978                nh,
20979                nkv,
20980                &kvl.len_d,
20981                b,
20982                scale,
20983                kvl.k_tok_bytes,
20984                kvl.v_tok_bytes,
20985                swa && crate::Engine::wkv_on(),
20986            )?;
20987        }
20988        if !fa_q8 {
20989            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
20990            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
20991        }
20992        {
20993            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
20994            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
20995            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
20996        }
20997        Ok(())
20998    }
20999
21000    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
21001    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
21002    fn gemma4_layer_tail_slotted(
21003        &self,
21004        e: &Engine,
21005        layer: &crate::hybrid::HybridLayer,
21006        next_norm: Option<&CudaSlice<f32>>,
21007        sl: &mut G4DcSlots,
21008    ) -> Result<(), Box<dyn std::error::Error>> {
21009        let n_embd = self.cfg.n_embd as usize;
21010        let eps = self.cfg.rms_eps;
21011        let bits = layer.gemma4.as_ref().unwrap();
21012        let crate::hybrid::Ffn::Dense {
21013            ffn_gate,
21014            ffn_up,
21015            ffn_down,
21016        } = &layer.ffn
21017        else {
21018            return Err("slotted tail: dense ffn only".into());
21019        };
21020        let pnfold = Engine::g4_pnfold_on();
21021        if pnfold {
21022            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
21023            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
21024            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
21025            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
21026            e.rms_pre_add_rms_norm_q8z_into(
21027                or,
21028                layer.post_attn_norm.float_data(),
21029                xr,
21030                bits.ffn_norm.float_data(),
21031                &mut sl.attn_out,
21032                &mut sl.zsh,
21033                n_embd,
21034                1,
21035                eps,
21036                &mut sl.zq,
21037                &mut sl.zd,
21038            )?;
21039        } else {
21040            e.add_rms_norm(
21041                &sl.cur,
21042                &sl.x,
21043                bits.ffn_norm.float_data(),
21044                &mut sl.attn_out,
21045                &mut sl.zsh,
21046                n_embd,
21047                1,
21048                eps,
21049            )?;
21050        }
21051        let n_ff = ffn_gate.out_features();
21052        if !pnfold {
21053            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
21054            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
21055        }
21056        {
21057            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
21058            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
21059            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
21060                && !e.matmul_nvfp4_fused2_into(
21061                    ffn_gate,
21062                    ffn_up,
21063                    zq,
21064                    zd,
21065                    &mut sl.gate,
21066                    &mut sl.up,
21067                )?
21068            {
21069                return Err("slotted tail: ffn fused2 unavailable".into());
21070            }
21071        }
21072        debug_assert!(e.uses_q8_1_fast(ffn_down));
21073        {
21074            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
21075            let upv = e.view(upr, n_ff);
21076            let up_all = upv.slice(0..n_ff);
21077            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
21078            e.gelu_tanh_mul_q8_1_into(
21079                gr,
21080                &up_all,
21081                &mut sl.act,
21082                n_ff,
21083                1,
21084                &mut sl.actq,
21085                &mut sl.actd,
21086            )?;
21087        }
21088        {
21089            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
21090            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
21091            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
21092        }
21093        if pnfold {
21094            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
21095            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
21096            if let Some(w) = next_norm {
21097                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
21098                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
21099                e.rms_pre_add_scale_rms_norm_q8_1_into(
21100                    f0r,
21101                    bits.post_ffw_norm.float_data(),
21102                    aor,
21103                    bits.layer_scale,
21104                    w,
21105                    &mut sl.xn,
21106                    n_embd,
21107                    1,
21108                    eps,
21109                    &mut sl.hq,
21110                    &mut sl.hd_,
21111                )?;
21112                return Ok(());
21113            }
21114        }
21115        e.rms_norm(
21116            &sl.f0,
21117            bits.post_ffw_norm.float_data(),
21118            &mut sl.sn,
21119            n_embd,
21120            1,
21121            eps,
21122        )?;
21123        match next_norm {
21124            Some(w) => {
21125                e.add_scale_rms_norm_q8_1_into(
21126                    &sl.sn,
21127                    &sl.attn_out,
21128                    bits.layer_scale,
21129                    w,
21130                    &mut sl.xn,
21131                    n_embd,
21132                    1,
21133                    eps,
21134                    &mut sl.hq,
21135                    &mut sl.hd_,
21136                )?;
21137            }
21138            None => {
21139                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
21140            }
21141        }
21142        Ok(())
21143    }
21144
21145    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
21146    #[allow(clippy::too_many_arguments)]
21147    fn gemma4_decode_attn_dc(
21148        &self,
21149        e: &Engine,
21150        fa: &crate::hybrid::FullAttnLayer,
21151        il: usize,
21152        hq: &CudaSlice<i8>,
21153        hdq: &CudaSlice<f32>,
21154        pos_d: &CudaSlice<i32>,
21155        cache: &mut Cache,
21156        cap_bucket_max: Option<(usize, usize)>,
21157    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21158        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
21159        let eps = self.cfg.rms_eps;
21160        let aux = self.gemma4_aux.as_ref().unwrap();
21161        let ones = aux.ones(e);
21162        #[cfg(debug_assertions)]
21163        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
21164        let (q0, k0, v0) = if swa {
21165            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
21166                Some(t3) => t3,
21167                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
21168                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
21169                    Some((q0, k0)) => {
21170                        let h0 = e.zeros(0)?;
21171                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
21172                        (q0, k0, v0)
21173                    }
21174                    None => {
21175                        let h0 = e.zeros(0)?;
21176                        (
21177                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
21178                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
21179                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
21180                        )
21181                    }
21182                },
21183            }
21184        } else {
21185            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
21186                Some(p) => p,
21187                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
21188                    Some(p) => p,
21189                    None => {
21190                        let h0 = e.zeros(0)?;
21191                        (
21192                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
21193                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
21194                        )
21195                    }
21196                },
21197            };
21198            let v0 = e.clone_dtod(&k0)?;
21199            (q0, k0, v0)
21200        };
21201        let mut q = e.uninit(nh * hd)?;
21202        let mut k = e.uninit(nkv * hd)?;
21203        let mut v = e.uninit(nkv * hd)?;
21204        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
21205        let ff = if swa {
21206            None
21207        } else {
21208            Some(
21209                aux.rope_freqs(e)
21210                    .expect("gemma4 global rope needs rope_freqs.weight"),
21211            )
21212        };
21213        #[cfg(debug_assertions)]
21214        if let Some(ff) = ff {
21215            crate::debug_assert_tensor_stream_device(
21216                ff,
21217                &e.stream(),
21218                "gemma4_decode_attn_dc.rope_freqs",
21219            );
21220        }
21221        let kvl = cache.kv[il].as_mut().unwrap();
21222        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
21223        if crate::Engine::qkv_append_on() {
21224            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
21225            e.rms_norm_qkv_rope_append_dc(
21226                &q0,
21227                &k0,
21228                &v0,
21229                fa.q_norm.float_data(),
21230                fa.k_norm.float_data(),
21231                ones,
21232                &mut q,
21233                &mut k,
21234                &mut v,
21235                hd,
21236                self.gemma4_rope_dims(il),
21237                nh,
21238                nkv,
21239                pos_d,
21240                nh,
21241                nkv,
21242                base,
21243                1.0,
21244                ff,
21245                eps,
21246                &mut kvl.k,
21247                &mut kvl.v,
21248                &kvl.len_d,
21249                kvl.k_tok_bytes,
21250                kvl.v_tok_bytes,
21251                kv_fp8,
21252            )?;
21253        } else {
21254            e.rms_norm_qkv_rope(
21255                &q0,
21256                &k0,
21257                &v0,
21258                fa.q_norm.float_data(),
21259                fa.k_norm.float_data(),
21260                ones,
21261                &mut q,
21262                &mut k,
21263                &mut v,
21264                hd,
21265                self.gemma4_rope_dims(il),
21266                nh,
21267                nkv,
21268                pos_d,
21269                nh,
21270                nkv,
21271                base,
21272                1.0,
21273                ff,
21274                eps,
21275            )?;
21276            e.append_kv_quantized_dc(
21277                &k,
21278                &v,
21279                &mut kvl.k,
21280                &mut kvl.v,
21281                &kvl.len_d,
21282                kvl.kv_dim_k,
21283                kvl.kv_dim_v,
21284                kvl.k_tok_bytes,
21285                kvl.v_tok_bytes,
21286                kv_fp8,
21287            )?;
21288        }
21289        e.inc_seqlen(&mut kvl.len_d)?;
21290        let mut attn = e.uninit(nh * hd)?;
21291        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
21292        // rides g4_matvec_m1_into instead of matmul's internal quantize.
21293        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
21294        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
21295        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
21296        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
21297        // (gemma4_e4b_attn, +0.65% valid window).
21298        match cap_bucket_max {
21299            None => {
21300                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
21301                // decode (SWA layers attend the last `sliding_window` keys); the device
21302                // counters carry only the append slot + the graph seam.
21303                kvl.len += 1;
21304                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
21305                if !swa
21306                    && hd == 512
21307                    && kvl.len >= crate::fa512_min_tkv()
21308                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
21309                {
21310                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
21311                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
21312                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
21313                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
21314                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
21315                    e.fa_decode_rows(
21316                        &q,
21317                        &kp,
21318                        &vp,
21319                        &mut attn,
21320                        hd,
21321                        nh,
21322                        nkv,
21323                        kvl.len - 1,
21324                        1,
21325                        scale,
21326                        kvl.k_tok_bytes,
21327                        kvl.v_tok_bytes,
21328                        Some((&kvl.len_d, -1)),
21329                        false,
21330                        false,
21331                        Some((&mut aq8, &mut ad8)),
21332                    )?;
21333                    fa_q8 = Some((aq8, ad8));
21334                } else if swa
21335                    && kvl.len > win
21336                    && hd == 256
21337                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
21338                {
21339                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
21340                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
21341                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
21342                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
21343                    e.fa_decode_rows_w(
21344                        &q,
21345                        &kp,
21346                        &vp,
21347                        &mut attn,
21348                        hd,
21349                        nh,
21350                        nkv,
21351                        &kvl.len_d,
21352                        -1,
21353                        1,
21354                        scale,
21355                        win,
21356                        kvl.k_tok_bytes,
21357                        kvl.v_tok_bytes,
21358                        Some((&mut aq8, &mut ad8)),
21359                    )?;
21360                    fa_q8 = Some((aq8, ad8));
21361                } else {
21362                    let (off_tok, t_kv) = if swa && kvl.len > win {
21363                        (kvl.len - win, win)
21364                    } else {
21365                        (0, kvl.len)
21366                    };
21367                    let k_view = e.view_u8_range(
21368                        &kvl.k,
21369                        off_tok * kvl.k_tok_bytes,
21370                        (off_tok + t_kv) * kvl.k_tok_bytes,
21371                    );
21372                    let v_view = e.view_u8_range(
21373                        &kvl.v,
21374                        off_tok * kvl.v_tok_bytes,
21375                        (off_tok + t_kv) * kvl.v_tok_bytes,
21376                    );
21377                    e.fa_decode_kvmod(
21378                        &q,
21379                        &k_view,
21380                        &v_view,
21381                        &mut attn,
21382                        hd,
21383                        nh,
21384                        nkv,
21385                        t_kv,
21386                        scale,
21387                        kvl.k_tok_bytes,
21388                        kvl.v_tok_bytes,
21389                        swa && crate::Engine::wkv_on(),
21390                    )?;
21391                }
21392            }
21393            Some((b_swa, b_glob)) => {
21394                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
21395                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
21396                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
21397                // the RUNG max for the rows family (kernels derive per-replay splits from
21398                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
21399                let k_view = e.view_u8(&kvl.k, kvl.k.len());
21400                let v_view = e.view_u8(&kvl.v, kvl.v.len());
21401                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
21402                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
21403                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
21404                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
21405                    e.fa_decode_rows(
21406                        &q,
21407                        &k_view,
21408                        &v_view,
21409                        &mut attn,
21410                        hd,
21411                        nh,
21412                        nkv,
21413                        b_glob - 1,
21414                        1,
21415                        scale,
21416                        kvl.k_tok_bytes,
21417                        kvl.v_tok_bytes,
21418                        Some((&kvl.len_d, -1)),
21419                        false,
21420                        false,
21421                        Some((&mut aq8, &mut ad8)),
21422                    )?;
21423                    fa_q8 = Some((aq8, ad8));
21424                } else if swa && b_swa > win && hd == 256 && rows_on {
21425                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
21426                    e.fa_decode_rows_w(
21427                        &q,
21428                        &k_view,
21429                        &v_view,
21430                        &mut attn,
21431                        hd,
21432                        nh,
21433                        nkv,
21434                        &kvl.len_d,
21435                        -1,
21436                        1,
21437                        scale,
21438                        win,
21439                        kvl.k_tok_bytes,
21440                        kvl.v_tok_bytes,
21441                        Some((&mut aq8, &mut ad8)),
21442                    )?;
21443                    fa_q8 = Some((aq8, ad8));
21444                } else {
21445                    let b = if swa { b_swa } else { b_glob };
21446                    e.fa_decode_dc(
21447                        &q,
21448                        &k_view,
21449                        &v_view,
21450                        &mut attn,
21451                        hd,
21452                        nh,
21453                        nkv,
21454                        &kvl.len_d,
21455                        b,
21456                        scale,
21457                        kvl.k_tok_bytes,
21458                        kvl.v_tok_bytes,
21459                        swa && crate::Engine::wkv_on(),
21460                    )?;
21461                }
21462            }
21463        }
21464        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
21465        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
21466        if let Some((aq8, ad8)) = fa_q8 {
21467            let mut y = e.uninit(fa.wo.out_features())?;
21468            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
21469            return Ok(y);
21470        }
21471        e.matmul(&fa.wo, &attn, 1)
21472    }
21473
21474    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
21475    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
21476    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
21477    /// views in-graph); caller gates and falls back to the dc-eager loop.
21478    #[allow(clippy::too_many_arguments)]
21479    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
21480    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
21481    pub fn gemma4_generate_graph(
21482        &self,
21483        e: &Engine,
21484        prompt_pos: usize,
21485        first_token: u32,
21486        cache: &mut Cache,
21487        max_new: usize,
21488        eos: &[u32],
21489        mut on_token: impl FnMut(u32) -> bool,
21490    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
21491        if self.is_gemma4_e4b() {
21492            return Err(
21493                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
21494                    .into(),
21495            );
21496        }
21497        use crate::decode::StopReason;
21498        let n_vocab = self.output.out_features();
21499        let n_embd = self.cfg.n_embd as usize;
21500        let embd_gpu = self
21501            .embd_gpu
21502            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
21503        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
21504        for kvl in cache.kv.iter_mut().flatten() {
21505            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
21506        }
21507        let mut token_d = e.stream().clone_htod(&[first_token])?;
21508        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
21509        let g4 = self.cfg.gemma4.as_ref().unwrap();
21510        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
21511        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
21512        let nkv_s = g4
21513            .head_count_kv
21514            .iter()
21515            .zip(g4.swa_pattern.iter())
21516            .find(|p| *p.1)
21517            .map(|p| *p.0 as usize)
21518            .unwrap_or(8);
21519        let nkv_g = g4
21520            .head_count_kv
21521            .iter()
21522            .zip(g4.swa_pattern.iter())
21523            .find(|p| !*p.1)
21524            .map(|p| *p.0 as usize)
21525            .unwrap_or(2);
21526        #[allow(clippy::type_complexity)]
21527        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
21528        let mut graphs: std::collections::HashMap<
21529            ((bool, usize), (bool, usize), bool, bool),
21530            (
21531                cudarc::driver::CudaGraph,
21532                Vec<Box<dyn std::any::Any + Send>>,
21533            ),
21534        > = Default::default();
21535        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
21536        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
21537        let mut slots = self.g4_dc_slots(e)?;
21538        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
21539        // baked at the door entry (the modulo keeps every capture valid indefinitely).
21540        const RING: usize = 64;
21541        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
21542        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
21543        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
21544        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
21545        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
21546        const DRAIN: usize = 1;
21547        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
21548        let ring_base = prompt_pos;
21549        let mut out = Vec::with_capacity(max_new);
21550        let mut reason = StopReason::MaxNew;
21551        let mut next = first_token;
21552        let mut captures = 0usize;
21553        for _ in 0..max_new {
21554            out.push(next);
21555            if eos.contains(&next) {
21556                reason = StopReason::Eos;
21557                break;
21558            }
21559            if !on_token(next) {
21560                reason = StopReason::Callback;
21561                break;
21562            }
21563            let t_kv = cache.pos + 1;
21564            // Bucket key per ARM (graph arc step 3):
21565            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
21566            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
21567            //    the component collapses to a single marker).
21568            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
21569            //    at/above it — the kernel derives splits from len_d per replay, so buckets
21570            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
21571            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
21572            let f512 = crate::fa512_min_tkv();
21573            let key_s = if t_kv > win {
21574                (true, usize::MAX)
21575            } else {
21576                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
21577            };
21578            let (key_g, rung_end) = if t_kv >= f512 {
21579                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
21580                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
21581                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
21582                ((true, end), end)
21583            } else {
21584                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
21585            };
21586            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
21587            if !graphs.contains_key(&key) {
21588                let bucket_max = (t_kv, rung_end);
21589                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
21590                let snap = cache.snapshot(e)?;
21591                let pos_save = e.dtoh_i32_one(&pos_d)?;
21592                let len_save: Vec<Option<i32>> = cache
21593                    .kv
21594                    .iter()
21595                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
21596                    .collect();
21597                let tok_save = e.dtoh_u32_one(&token_d)?;
21598                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
21599                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
21600                // regression class, and this door's measured -8.8%. The keeper pins warmup
21601                // transients so the captured graph holds kernel nodes only.
21602                let graph = {
21603                    let tok_ref = &mut token_d;
21604                    let pos_ref = &mut pos_d;
21605                    let cache_ref = &mut *cache;
21606                    let slots_ref = &mut slots;
21607                    let ring_ref = &mut ring;
21608                    e.capture_graph_retained_flags(
21609                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
21610                        |e| {
21611                        // self-feeding: the argmax writes token_d itself.
21612                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
21613                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
21614                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
21615                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
21616                                                           cache_ref, n_vocab, Some(bucket_max),
21617                                                           sl, tok_ref, Some((rg, ring_base)))
21618                    })?
21619                };
21620                cache.rollback(e, &snap, 0)?;
21621                e.set_i32_one(&mut pos_d, pos_save)?;
21622                for (il, ls) in len_save.iter().enumerate() {
21623                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
21624                        e.set_i32_one(&mut kvl.len_d, *v)?;
21625                    }
21626                }
21627                e.set_u32_one(&mut token_d, tok_save)?;
21628                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
21629                    && let Ok(c) = crate::graph_update::node_census(&graph.0)
21630                {
21631                    eprintln!("[graph-census] {c:?}");
21632                }
21633                graphs.insert(key, graph);
21634                captures += 1;
21635            }
21636            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
21637            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
21638            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
21639            // the budget; capture warmups already emitted their tokens through the ring.
21640            let mut chunk = 1usize;
21641            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
21642                .ok()
21643                .and_then(|v| v.parse().ok())
21644                .unwrap_or(DRAIN);
21645            while chunk < drain_cap && out.len() + chunk < max_new {
21646                let t_next = cache.pos + 1 + chunk;
21647                let key_s2 = if t_next > win {
21648                    (true, usize::MAX)
21649                } else {
21650                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
21651                };
21652                let key_g2 = if t_next >= f512 {
21653                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
21654                } else {
21655                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
21656                };
21657                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
21658                    break;
21659                }
21660                chunk += 1;
21661            }
21662            let g = &graphs.get(&key).unwrap().0;
21663            for _ in 0..chunk {
21664                g.launch()?;
21665            }
21666            e.stream().synchronize()?;
21667            let ringh = e.dtoh_u32(&ring)?;
21668            for j in 0..chunk {
21669                let pos_j = cache.pos + j;
21670                let tok_j = ringh[(pos_j - ring_base) % RING];
21671                cache.pos += 0; // advanced below in one shot
21672                if j + 1 == chunk {
21673                    next = tok_j;
21674                } else {
21675                    out.push(tok_j);
21676                    if eos.contains(&tok_j) || !on_token(tok_j) {
21677                        reason = if eos.contains(&tok_j) {
21678                            StopReason::Eos
21679                        } else {
21680                            StopReason::Callback
21681                        };
21682                        // roll device/host state back to the stop point.
21683                        let keep = cache.pos + j + 1;
21684                        e.set_i32_one(&mut pos_d, keep as i32)?;
21685                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
21686                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
21687                            kvl.len = keep;
21688                        }
21689                        cache.pos = keep;
21690                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
21691                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
21692                        }
21693                        return Ok((out, reason));
21694                    }
21695                }
21696            }
21697            cache.pos += chunk;
21698            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
21699                kvl.len += chunk;
21700            }
21701        }
21702        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
21703            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
21704        }
21705        Ok((out, reason))
21706    }
21707
21708    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
21709    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
21710    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
21711    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
21712    /// logits (host) + advances cache.pos by t.
21713    pub(crate) fn gemma4_decode_step_t(
21714        &self,
21715        e: &Engine,
21716        tokens: &[u32],
21717        pos0: usize,
21718        cache: &mut Cache,
21719    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
21720        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
21721    }
21722
21723    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
21724    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
21725    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
21726    pub(crate) fn gemma4_decode_step_t_am(
21727        &self,
21728        e: &Engine,
21729        tokens: &[u32],
21730        pos0: usize,
21731        cache: &mut Cache,
21732    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21733        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
21734        let t = tokens.len();
21735        let n_vocab = self.output.out_features();
21736        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
21737        for i in 0..t {
21738            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
21739        }
21740        Ok((e.dtoh_u32(&toks)?, hn))
21741    }
21742
21743    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
21744    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
21745    pub(crate) fn gemma4_decode_step_t_am_dev(
21746        &self,
21747        e: &Engine,
21748        tok_d: &CudaSlice<u32>,
21749        t: usize,
21750        pos0: usize,
21751        cache: &mut Cache,
21752    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21753        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
21754        let n_vocab = self.output.out_features();
21755        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
21756        for i in 0..t {
21757            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
21758        }
21759        Ok((vam, hn))
21760    }
21761
21762    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
21763    /// llama's h_nextn convention).
21764    pub(crate) fn gemma4_decode_step_t_h(
21765        &self,
21766        e: &Engine,
21767        tokens: &[u32],
21768        pos0: usize,
21769        cache: &mut Cache,
21770    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21771        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
21772        let t = tokens.len();
21773        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
21774        e.softcap(&mut ld, cap, t * self.output.out_features())?;
21775        Ok((e.dtoh(&ld)?, hn))
21776    }
21777
21778    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
21779    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
21780    pub(crate) fn verify_stream_scratch(
21781        &self,
21782        e: &Engine,
21783        cap: usize,
21784    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
21785        Ok(VerifyStreamScratch {
21786            pos_d: e.htod_i32(&vec![0i32; cap])?,
21787            row_ctrs: (0..cap)
21788                .map(|_| e.htod_i32(&[0]))
21789                .collect::<Result<_, _>>()?,
21790        })
21791    }
21792
21793    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
21794    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
21795    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
21796    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
21797    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
21798    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
21799    /// sync, exactly the turnaround the burst exists to remove.
21800    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
21801    pub(crate) fn gemma4_verify_t_am_stream(
21802        &self,
21803        e: &Engine,
21804        tok_d: &CudaSlice<u32>,
21805        t: usize,
21806        ctr: &CudaSlice<i32>,
21807        hint: usize,
21808        cache: &mut Cache,
21809        scr: &mut VerifyStreamScratch,
21810    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21811        let n_embd = self.cfg.n_embd as usize;
21812        let eps = self.cfg.rms_eps;
21813        assert!(t <= scr.row_ctrs.len() && t <= 64);
21814        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
21815        for i in 0..t {
21816            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
21817        }
21818        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
21819        let embd_gpu = self
21820            .embd_gpu
21821            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
21822        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
21823        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
21824        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
21825        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
21826        let n_layers = self.layers.len();
21827        for (il, layer) in self.layers.iter().enumerate() {
21828            let (hq, hdq) = match h_carry.take() {
21829                Some(p) => p,
21830                None => {
21831                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
21832                }
21833            };
21834            let Mixer::Full(fa) = &layer.mixer else {
21835                panic!("gemma4 layer {il} not full-attn")
21836            };
21837            let o = self
21838                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
21839            let next_norm = if il + 1 < n_layers {
21840                Some(self.layers[il + 1].attn_norm.float_data())
21841            } else {
21842                None
21843            };
21844            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
21845            x = xn;
21846            h_carry = hn;
21847            self.dflash_tap(e, cache, il, &x, t)?;
21848        }
21849        let mut hn = e.uninit(t * n_embd)?;
21850        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
21851        let ld = e.matmul(&self.output, &hn, t)?;
21852        let n_vocab = self.output.out_features();
21853        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
21854        for i in 0..t {
21855            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
21856        }
21857        Ok((vam, hn))
21858    }
21859
21860    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
21861    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
21862    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
21863    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
21864    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
21865    /// kernel later if it shows in the profile).
21866    pub(crate) fn dflash_tap(
21867        &self,
21868        e: &Engine,
21869        cache: &mut Cache,
21870        il: usize,
21871        x: &CudaSlice<f32>,
21872        t: usize,
21873    ) -> Result<(), Box<dyn std::error::Error>> {
21874        let Some(taps) = cache.dflash_taps.as_mut() else {
21875            return Ok(());
21876        };
21877        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
21878            return Ok(());
21879        };
21880        let h = taps.hidden;
21881        let n_taps = taps.layer_ids.len();
21882        let base = taps.base;
21883        debug_assert!(
21884            base + t <= taps.t,
21885            "tap window {base}+{t} exceeds sink {}",
21886            taps.t
21887        );
21888        let xv = e.view(x, t * h);
21889        for r in 0..t {
21890            let row = xv.slice(r * h..(r + 1) * h);
21891            e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
21892        }
21893        Ok(())
21894    }
21895
21896    fn gemma4_verify_trunk(
21897        &self,
21898        e: &Engine,
21899        tokens: &[u32],
21900        pos0: usize,
21901        cache: &mut Cache,
21902        tok_dev: Option<&CudaSlice<u32>>,
21903    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
21904        let n_embd = self.cfg.n_embd as usize;
21905        let eps = self.cfg.rms_eps;
21906        let t = tokens.len();
21907        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
21908        let pos_d = e.htod_i32(&pos)?;
21909        let mut x = match tok_dev {
21910            Some(td) => {
21911                let embd_gpu = self
21912                    .embd_gpu
21913                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
21914                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
21915                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
21916            }
21917            None => e.htod(&self.embd.try_gather(n_embd, tokens)?)?,
21918        };
21919        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
21920        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
21921        let n_layers = self.layers.len();
21922        for (il, layer) in self.layers.iter().enumerate() {
21923            let (hq, hdq) = match h_carry.take() {
21924                Some(p) => p,
21925                None => {
21926                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
21927                }
21928            };
21929            let Mixer::Full(fa) = &layer.mixer else {
21930                panic!("gemma4 layer {il} not full-attn")
21931            };
21932            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
21933            let next_norm = if il + 1 < n_layers {
21934                Some(self.layers[il + 1].attn_norm.float_data())
21935            } else {
21936                None
21937            };
21938            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
21939            x = xn;
21940            h_carry = hn;
21941            self.dflash_tap(e, cache, il, &x, t)?;
21942        }
21943        let mut hn = e.uninit(t * n_embd)?;
21944        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
21945        let mut ld = e.matmul(&self.output, &hn, t)?;
21946        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
21947        cache.pos += t;
21948        Ok((ld, hn))
21949    }
21950
21951    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
21952    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
21953    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
21954    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
21955    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
21956    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
21957    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
21958    #[allow(clippy::too_many_arguments)]
21959    fn gemma4_verify_attn_stream(
21960        &self,
21961        e: &Engine,
21962        fa: &crate::hybrid::FullAttnLayer,
21963        il: usize,
21964        hq: &CudaSlice<i8>,
21965        hdq: &CudaSlice<f32>,
21966        pos_d: &CudaSlice<i32>,
21967        t: usize,
21968        cache: &mut Cache,
21969        hint: usize,
21970        row_ctrs: &[CudaSlice<i32>],
21971    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21972        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
21973        let eps = self.cfg.rms_eps;
21974        let aux = self.gemma4_aux.as_ref().unwrap();
21975        let ones = aux.ones(e);
21976        #[cfg(debug_assertions)]
21977        crate::debug_assert_tensor_stream_device(
21978            ones,
21979            &e.stream(),
21980            "gemma4_verify_attn_stream.ones",
21981        );
21982        let h0 = e.zeros(0)?;
21983        let h = &h0;
21984        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
21985        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
21986        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21987        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
21988        let fused_qkv = if f2b {
21989            if swa {
21990                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
21991                    .map(|(a, b, c)| (a, b, Some(c)))
21992            } else {
21993                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
21994                    .map(|(a, b)| (a, b, None))
21995            }
21996        } else {
21997            None
21998        };
21999        let (q0, k0, v0) = match fused_qkv {
22000            Some((a, b, cv)) => {
22001                let v = match cv {
22002                    Some(c) => c,
22003                    None => e.clone_dtod(&b)?,
22004                };
22005                (a, b, v)
22006            }
22007            None => {
22008                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
22009                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
22010                let v0 = if swa {
22011                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
22012                } else {
22013                    e.clone_dtod(&k0)?
22014                };
22015                (q0, k0, v0)
22016            }
22017        };
22018        let mut q = e.uninit(t * nh * hd)?;
22019        let mut k = e.uninit(t * nkv * hd)?;
22020        let mut v = e.uninit(t * nkv * hd)?;
22021        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
22022        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
22023        let ff = if swa {
22024            None
22025        } else {
22026            Some(
22027                aux.rope_freqs(e)
22028                    .expect("gemma4 global rope needs rope_freqs.weight"),
22029            )
22030        };
22031        #[cfg(debug_assertions)]
22032        if let Some(ff) = ff {
22033            crate::debug_assert_tensor_stream_device(
22034                ff,
22035                &e.stream(),
22036                "gemma4_verify_attn_stream.rope_freqs",
22037            );
22038        }
22039        e.rms_norm_qkv_rope(
22040            &q0,
22041            &k0,
22042            &v0,
22043            fa.q_norm.float_data(),
22044            fa.k_norm.float_data(),
22045            ones,
22046            &mut q,
22047            &mut k,
22048            &mut v,
22049            hd,
22050            self.gemma4_rope_dims(il),
22051            nh * t,
22052            nkv * t,
22053            pos_d,
22054            nh,
22055            nkv,
22056            base,
22057            1.0,
22058            ff,
22059            eps,
22060        )?;
22061        let kvl = cache.kv[il].as_mut().unwrap();
22062        // append at the DEVICE slot; the counter advances by t on-device.
22063        e.append_kv_quantized_rows_dc(
22064            &k,
22065            &v,
22066            &mut kvl.k,
22067            &mut kvl.v,
22068            &kvl.len_d,
22069            t,
22070            kvl.kv_dim_k,
22071            kvl.kv_dim_v,
22072            kvl.k_tok_bytes,
22073            kvl.v_tok_bytes,
22074            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
22075        )?;
22076        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
22077        // the sole len writer after this round's attention (base stays = old len, plus = 0).
22078        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
22079        let mut attn = e.uninit(t * nh * hd)?;
22080        let k_view = e.view_u8(&kvl.k, kvl.k.len());
22081        let v_view = e.view_u8(&kvl.v, kvl.v.len());
22082        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
22083        // and a stable window regime — the same rung/regime keys as the draft graph).
22084        if swa && hint + 1 >= win {
22085            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
22086            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
22087            e.fa_decode_rows_w(
22088                &q,
22089                &k_view,
22090                &v_view,
22091                &mut attn,
22092                hd,
22093                nh,
22094                nkv,
22095                &kvl.len_d,
22096                0,
22097                t,
22098                scale,
22099                win,
22100                kvl.k_tok_bytes,
22101                kvl.v_tok_bytes,
22102                None,
22103            )?;
22104        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
22105            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
22106            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
22107            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
22108            // Burst entry gates the horizon onto one side of the crossover, so hint decides
22109            // for every row.
22110            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
22111            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
22112            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
22113            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
22114            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
22115            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
22116            // any bucket >= the live length is exact.
22117            let bucket = (hint + t + 2)
22118                .next_power_of_two()
22119                .min(crate::fa512_min_tkv().saturating_sub(1));
22120            let qv = e.view(&q, t * nh * hd);
22121            #[allow(clippy::needless_range_loop)]
22122            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
22123            for i in 0..t {
22124                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
22125                let mut q_one = e.uninit(nh * hd)?;
22126                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
22127                let mut a_one = e.uninit(nh * hd)?;
22128                e.fa_decode_dc(
22129                    &q_one,
22130                    &k_view,
22131                    &v_view,
22132                    &mut a_one,
22133                    hd,
22134                    nh,
22135                    nkv,
22136                    &row_ctrs[i],
22137                    bucket,
22138                    scale,
22139                    kvl.k_tok_bytes,
22140                    kvl.v_tok_bytes,
22141                    false,
22142                )?;
22143                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
22144            }
22145        } else if hd == 512 {
22146            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
22147            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
22148            e.fa_decode_rows(
22149                &q,
22150                &k_view,
22151                &v_view,
22152                &mut attn,
22153                hd,
22154                nh,
22155                nkv,
22156                hint,
22157                t,
22158                scale,
22159                kvl.k_tok_bytes,
22160                kvl.v_tok_bytes,
22161                Some((&kvl.len_d, 0)),
22162                false,
22163                false,
22164                None,
22165            )?;
22166        } else {
22167            // hd256 under-window: v4 device-len rows twin.
22168            e.fa_decode_rows_dc(
22169                &q,
22170                &k_view,
22171                &v_view,
22172                &mut attn,
22173                hd,
22174                nh,
22175                nkv,
22176                &kvl.len_d,
22177                hint + t,
22178                t,
22179                scale,
22180                kvl.k_tok_bytes,
22181                kvl.v_tok_bytes,
22182                0,
22183                swa && crate::Engine::wkv_on(),
22184            )?;
22185        }
22186        e.matmul(&fa.wo, &attn, t)
22187    }
22188
22189    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
22190    fn gemma4_verify_attn(
22191        &self,
22192        e: &Engine,
22193        fa: &crate::hybrid::FullAttnLayer,
22194        il: usize,
22195        hq: &CudaSlice<i8>,
22196        hdq: &CudaSlice<f32>,
22197        pos_d: &CudaSlice<i32>,
22198        t: usize,
22199        cache: &mut Cache,
22200    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22201        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
22202        let eps = self.cfg.rms_eps;
22203        let aux = self.gemma4_aux.as_ref().unwrap();
22204        let ones = aux.ones(e);
22205        #[cfg(debug_assertions)]
22206        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
22207        let n_embd = self.cfg.n_embd as usize;
22208        let _ = n_embd;
22209
22210        let h0 = e.zeros(0)?;
22211        let h = &h0;
22212        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
22213        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
22214        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22215        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
22216        let fused_qkv = if f2b {
22217            if swa {
22218                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
22219                    .map(|(a, b, c)| (a, b, Some(c)))
22220            } else {
22221                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
22222                    .map(|(a, b)| (a, b, None))
22223            }
22224        } else {
22225            None
22226        };
22227        let (q0, k0, v0) = match fused_qkv {
22228            Some((a, b, cv)) => {
22229                let v = match cv {
22230                    Some(c) => c,
22231                    None => e.clone_dtod(&b)?,
22232                };
22233                (a, b, v)
22234            }
22235            None => {
22236                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
22237                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
22238                let v0 = if swa {
22239                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
22240                } else {
22241                    e.clone_dtod(&k0)?
22242                };
22243                (q0, k0, v0)
22244            }
22245        };
22246        let mut q = e.uninit(t * nh * hd)?;
22247        let mut k = e.uninit(t * nkv * hd)?;
22248        let mut v = e.uninit(t * nkv * hd)?;
22249        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
22250        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
22251        let ff = if swa {
22252            None
22253        } else {
22254            Some(
22255                aux.rope_freqs(e)
22256                    .expect("gemma4 global rope needs rope_freqs.weight"),
22257            )
22258        };
22259        #[cfg(debug_assertions)]
22260        if let Some(ff) = ff {
22261            crate::debug_assert_tensor_stream_device(
22262                ff,
22263                &e.stream(),
22264                "gemma4_verify_attn.rope_freqs",
22265            );
22266        }
22267        e.rms_norm_qkv_rope(
22268            &q0,
22269            &k0,
22270            &v0,
22271            fa.q_norm.float_data(),
22272            fa.k_norm.float_data(),
22273            ones,
22274            &mut q,
22275            &mut k,
22276            &mut v,
22277            hd,
22278            self.gemma4_rope_dims(il),
22279            nh * t,
22280            nkv * t,
22281            pos_d,
22282            nh,
22283            nkv,
22284            base,
22285            1.0,
22286            ff,
22287            eps,
22288        )?;
22289        let kvl = cache.kv[il].as_mut().unwrap();
22290        let base_len = kvl.len;
22291        e.append_kv_quantized_rows(
22292            &k,
22293            &v,
22294            &mut kvl.k,
22295            &mut kvl.v,
22296            base_len,
22297            t,
22298            kvl.kv_dim_k,
22299            kvl.kv_dim_v,
22300            kvl.k_tok_bytes,
22301            kvl.v_tok_bytes,
22302            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
22303        )?;
22304        kvl.len += t;
22305        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
22306        let mut attn = e.uninit(t * nh * hd)?;
22307        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
22308        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
22309        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
22310            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
22311            // decode rides the SAME symbol at t=1 (parity law).
22312            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
22313        if rows_ok && (!swa || base_len + t <= win) {
22314            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
22315            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
22316            if hd == 512 {
22317                // device-len twin: sync the counter to the verify base (async arg-store).
22318                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
22319                e.fa_decode_rows(
22320                    &q,
22321                    &k_view,
22322                    &v_view,
22323                    &mut attn,
22324                    hd,
22325                    nh,
22326                    nkv,
22327                    base_len,
22328                    t,
22329                    scale,
22330                    kvl.k_tok_bytes,
22331                    kvl.v_tok_bytes,
22332                    Some((&kvl.len_d, 0)),
22333                    false,
22334                    swa && crate::Engine::wkv_on(),
22335                    None,
22336                )?;
22337            } else {
22338                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
22339                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
22340                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
22341                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
22342                e.fa_decode_rows_dc(
22343                    &q,
22344                    &k_view,
22345                    &v_view,
22346                    &mut attn,
22347                    hd,
22348                    nh,
22349                    nkv,
22350                    &kvl.len_d,
22351                    base_len + t,
22352                    t,
22353                    scale,
22354                    kvl.k_tok_bytes,
22355                    kvl.v_tok_bytes,
22356                    0,
22357                    swa && crate::Engine::wkv_on(),
22358                )?;
22359            }
22360            return e.matmul(&fa.wo, &attn, t);
22361        }
22362        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
22363        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
22364        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
22365        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
22366        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
22367        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
22368        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
22369        if hd == 256
22370            && swa
22371            && base_len + 1 >= win
22372            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
22373        {
22374            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
22375            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
22376            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
22377            e.fa_decode_rows_w(
22378                &q,
22379                &k_view,
22380                &v_view,
22381                &mut attn,
22382                hd,
22383                nh,
22384                nkv,
22385                &kvl.len_d,
22386                0,
22387                t,
22388                scale,
22389                win,
22390                kvl.k_tok_bytes,
22391                kvl.v_tok_bytes,
22392                None,
22393            )?;
22394            return e.matmul(&fa.wo, &attn, t);
22395        }
22396        for i in 0..t {
22397            let avail = base_len + i + 1;
22398            let (off_tok, t_kv) = if swa && avail > win {
22399                (avail - win, win)
22400            } else {
22401                (0, avail)
22402            };
22403            let k_view = e.view_u8_range(
22404                &kvl.k,
22405                off_tok * kvl.k_tok_bytes,
22406                (off_tok + t_kv) * kvl.k_tok_bytes,
22407            );
22408            let v_view = e.view_u8_range(
22409                &kvl.v,
22410                off_tok * kvl.v_tok_bytes,
22411                (off_tok + t_kv) * kvl.v_tok_bytes,
22412            );
22413            let qi = e.view(&q, t * nh * hd);
22414            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
22415            let mut q_one = e.uninit(nh * hd)?;
22416            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
22417            let mut a_one = e.uninit(nh * hd)?;
22418            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
22419            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
22420            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
22421            if swa
22422                && avail > win
22423                && hd == 256
22424                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
22425            {
22426                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
22427                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
22428                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
22429                e.fa_decode_rows_w(
22430                    &q_one,
22431                    &kp,
22432                    &vp,
22433                    &mut a_one,
22434                    hd,
22435                    nh,
22436                    nkv,
22437                    &kvl.len_d,
22438                    0,
22439                    1,
22440                    scale,
22441                    win,
22442                    kvl.k_tok_bytes,
22443                    kvl.v_tok_bytes,
22444                    None,
22445                )?;
22446            } else if !swa
22447                && hd == 512
22448                && avail >= crate::fa512_min_tkv()
22449                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
22450            {
22451                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
22452                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
22453                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
22454                e.fa_decode_rows(
22455                    &q_one,
22456                    &kp,
22457                    &vp,
22458                    &mut a_one,
22459                    hd,
22460                    nh,
22461                    nkv,
22462                    avail - 1,
22463                    1,
22464                    scale,
22465                    kvl.k_tok_bytes,
22466                    kvl.v_tok_bytes,
22467                    Some((&kvl.len_d, 0)),
22468                    false,
22469                    false,
22470                    None,
22471                )?;
22472            } else {
22473                e.fa_decode_kvmod(
22474                    &q_one,
22475                    &k_view,
22476                    &v_view,
22477                    &mut a_one,
22478                    hd,
22479                    nh,
22480                    nkv,
22481                    t_kv,
22482                    scale,
22483                    kvl.k_tok_bytes,
22484                    kvl.v_tok_bytes,
22485                    swa && crate::Engine::wkv_on(),
22486                )?;
22487            }
22488            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
22489        }
22490        e.matmul(&fa.wo, &attn, t)
22491    }
22492
22493    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
22494    /// h_seed = pre-output_norm hidden). Advances cache.pos.
22495    pub(crate) fn gemma4_decode_step_h(
22496        &self,
22497        e: &Engine,
22498        token: u32,
22499        cache: &mut Cache,
22500    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22501        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
22502        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
22503        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
22504        // unsplit rather than guessing a fence.
22505        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
22506            let rt = crate::pp::Pp2Rt::get(e)?;
22507            let _walk = rt.acquire_walk("gemma4_decode_step_h_pp2")?;
22508            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
22509        }
22510        if crate::pp::pp_cuts(self.layers.len()).is_some() {
22511            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
22512        }
22513        let n_embd = self.cfg.n_embd as usize;
22514        let eps = self.cfg.rms_eps;
22515        let pos_d = e.htod_i32(&[cache.pos as i32])?;
22516        let mut x = e.htod(&self.embd.try_gather(n_embd, &[token])?)?;
22517        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
22518        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
22519        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
22520        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
22521        let n_layers = self.layers.len();
22522        for (il, layer) in self.layers.iter().enumerate() {
22523            let (hq, hdq) = match h_carry.take() {
22524                Some(p) => p,
22525                None => {
22526                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
22527                }
22528            };
22529            let Mixer::Full(fa) = &layer.mixer else {
22530                panic!("gemma4 layer {il} not full-attn")
22531            };
22532            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
22533            let next_norm = if il + 1 < n_layers {
22534                Some(self.layers[il + 1].attn_norm.float_data())
22535            } else {
22536                None
22537            };
22538            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
22539            x = xn;
22540            h_carry = hn;
22541        }
22542        let mut hn = e.uninit(n_embd)?;
22543        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
22544        let h_seed = e.clone_dtod(&x)?;
22545        let mut ld = e.matmul(&self.output, &hn, 1)?;
22546        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
22547        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
22548        self.gemma4_suppress(e, &mut ld, 1)?;
22549        let logits = e.dtoh(&ld)?;
22550        cache.pos += 1;
22551        Ok((logits, h_seed))
22552    }
22553
22554    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
22555    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
22556    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
22557    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
22558    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
22559    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
22560    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
22561    fn gemma4_decode_layers(
22562        &self,
22563        e: &Engine,
22564        mut x: CudaSlice<f32>,
22565        lo: usize,
22566        hi: usize,
22567        pos_d: &CudaSlice<i32>,
22568        cache: &mut Cache,
22569    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22570        let n_embd = self.cfg.n_embd as usize;
22571        let eps = self.cfg.rms_eps;
22572        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
22573        for il in lo..hi {
22574            let layer = &self.layers[il];
22575            let (hq, hdq) = match h_carry.take() {
22576                Some(p) => p,
22577                // range head: il == lo — norm against THIS layer's attn_norm.
22578                None => {
22579                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
22580                }
22581            };
22582            let Mixer::Full(fa) = &layer.mixer else {
22583                panic!("gemma4 layer {il} not full-attn")
22584            };
22585            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
22586            let next_norm = if il + 1 < hi {
22587                Some(self.layers[il + 1].attn_norm.float_data())
22588            } else {
22589                None
22590            };
22591            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
22592            x = xn;
22593            h_carry = hn;
22594        }
22595        Ok(x)
22596    }
22597
22598    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
22599    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
22600    /// boundary handoff — same choreography as the generic arm (decode.rs), same
22601    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
22602    /// stage 1 = layers [split, n) + output_norm + softcapped head.
22603    /// Each stage uploads its own copy of the step's position scalar on its own stream.
22604    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
22605    fn gemma4_decode_step_h_pp2(
22606        &self,
22607        e: &Engine,
22608        token: u32,
22609        cache: &mut Cache,
22610        split: usize,
22611    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22612        if crate::pp::pp2_streams_off() {
22613            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
22614        }
22615        let rt = crate::pp::Pp2Rt::get(e)?;
22616        let e0 = rt.engine(0, e);
22617        let e1 = rt.engine(1, e);
22618        let n_embd = self.cfg.n_embd as usize;
22619        let eps = self.cfg.rms_eps;
22620        let pos = cache.pos as i32;
22621
22622        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
22623        let slot = {
22624            let _st0 = rt.enter(0);
22625            let pos_d = e0.htod_i32(&[pos])?;
22626            #[cfg(debug_assertions)]
22627            crate::debug_assert_tensor_stream_device(
22628                &pos_d,
22629                &e0.stream(),
22630                "gemma4_decode_step_h_pp2.stage0.pos_d",
22631            );
22632            let mut x = e0.htod(&self.embd.try_gather(n_embd, &[token])?)?;
22633            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
22634            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
22635            rt.tx(0, &x, n_embd)?
22636        };
22637
22638        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
22639        let _st1 = rt.enter(1);
22640        let pos_d = e1.htod_i32(&[pos])?;
22641        #[cfg(debug_assertions)]
22642        crate::debug_assert_tensor_stream_device(
22643            &pos_d,
22644            &e1.stream(),
22645            "gemma4_decode_step_h_pp2.stage1.pos_d",
22646        );
22647        let x = rt.rx(0, slot, n_embd)?;
22648        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
22649
22650        let mut hn = e1.uninit(n_embd)?;
22651        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
22652        let h_seed = e1.clone_dtod(&x)?;
22653        let mut ld = e1.matmul(&self.output, &hn, 1)?;
22654        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
22655        e1.softcap(&mut ld, cap, self.output.out_features())?;
22656        self.gemma4_suppress(e1, &mut ld, 1)?;
22657        let logits = e1.dtoh(&ld)?;
22658        cache.pos += 1;
22659        Ok((logits, h_seed))
22660    }
22661
22662    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
22663    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
22664    fn gemma4_decode_step_h_pp2_samestream(
22665        &self,
22666        e: &Engine,
22667        token: u32,
22668        cache: &mut Cache,
22669        split: usize,
22670    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
22671        let n_embd = self.cfg.n_embd as usize;
22672        let eps = self.cfg.rms_eps;
22673        let pos_d = e.htod_i32(&[cache.pos as i32])?;
22674
22675        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
22676        let mut x = e.htod(&self.embd.try_gather(n_embd, &[token])?)?;
22677        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
22678        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
22679
22680        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
22681        let boundary_tx = e.clone_dtod(&x)?;
22682        let boundary_rx = e.clone_dtod(&boundary_tx)?;
22683
22684        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
22685        let x =
22686            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
22687
22688        let mut hn = e.uninit(n_embd)?;
22689        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
22690        let h_seed = e.clone_dtod(&x)?;
22691        let mut ld = e.matmul(&self.output, &hn, 1)?;
22692        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
22693        e.softcap(&mut ld, cap, self.output.out_features())?;
22694        self.gemma4_suppress(e, &mut ld, 1)?;
22695        let logits = e.dtoh(&ld)?;
22696        cache.pos += 1;
22697        Ok((logits, h_seed))
22698    }
22699}
22700
22701// ============================ step35 (Step-3.7-Flash) ==================================
22702// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
22703// FAMILY and not a few branches inside the generic `full_attn*` chain:
22704//
22705//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
22706//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
22707//      shapes and the FA head counts would be wrong on 33 of 45 layers.
22708//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
22709//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
22710//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
22711//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
22712//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
22713//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
22714//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
22715//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
22716//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
22717//
22718// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
22719impl HybridModel {
22720    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
22721    /// synthesize a drafter or trunk layer from a neighboring class.
22722    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
22723        let geometry = self
22724            .cfg
22725            .layer_geometry(il as u32)
22726            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
22727        debug_assert_eq!(
22728            geometry.attention_gate,
22729            memra_gguf::config::AttentionGateKind::SeparateHead
22730        );
22731        geometry
22732    }
22733
22734    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
22735    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
22736    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
22737    ///
22738    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
22739    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
22740    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
22741    /// `cache`:
22742    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
22743    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
22744    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
22745    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
22746    ///     contract, lane/chunkinv-flip).
22747    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
22748    ///     q/k/v, no cache side effect.
22749    ///
22750    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
22751    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
22752    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
22753    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
22754    /// still contains must be masked per query. memra's window convention
22755    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
22756    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
22757    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
22758    ///
22759    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
22760    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
22761    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
22762    ///
22763    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
22764    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
22765    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
22766    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
22767    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
22768    /// hidden rows, and the generated text — a function of the chunk size:
22769    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
22770    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
22771    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
22772    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
22773    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
22774    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
22775    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
22776    ///   one-token change in a documented machine-config knob changed the answer.
22777    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
22778    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
22779    /// the same rows moves the logits by ~1.8.
22780    ///
22781    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
22782    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
22783    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
22784    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
22785    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
22786    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
22787    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
22788    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
22789    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
22790    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
22791    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
22792    /// those with t_kv <= win = 512.
22793    #[allow(clippy::too_many_arguments)]
22794    fn step35_attn_pre_wo(
22795        &self,
22796        e: &Engine,
22797        fa: &FullAttnLayer,
22798        mut g3: Vec<CudaSlice<f32>>,
22799        hg: Option<&CudaSlice<f32>>,
22800        gt_pre: Option<&CudaSlice<f32>>,
22801        pos_d: &CudaSlice<i32>,
22802        t: usize,
22803        cache: Option<&mut Cache>,
22804        il: usize,
22805        seq_end: usize,
22806    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22807        let geometry = self.cfg.full_attention_geometry_at(il as u32);
22808        let hd = geometry.head_dim_k as usize;
22809        let nkv = geometry.n_head_kv as usize;
22810        let nh = geometry.n_head as usize;
22811        let rbase = geometry.rope_base;
22812        let scale = geometry.attention_scale();
22813        let swa = geometry.window.is_some();
22814        let eps = self.cfg.rms_eps;
22815        let win = geometry.window.unwrap_or(0) as usize;
22816        let n_rot = geometry.n_rot as usize;
22817
22818        let v = g3.pop().unwrap();
22819        let k0 = g3.pop().unwrap();
22820        let q0 = g3.pop().unwrap();
22821
22822        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
22823        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
22824        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
22825        let mut q = e.uninit(t * nh * hd)?;
22826        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
22827        let mut k = e.uninit(t * nkv * hd)?;
22828        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
22829        let ff = if geometry.rope_factors {
22830            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
22831        } else {
22832            None
22833        };
22834        #[cfg(debug_assertions)]
22835        if let Some(ff) = ff {
22836            crate::debug_assert_tensor_stream_device(
22837                ff,
22838                &e.stream(),
22839                "step35_attn_pre_wo.rope_freqs",
22840            );
22841        }
22842        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
22843
22844        let mut attn = e.uninit(t * nh * hd)?;
22845        match cache {
22846            Some(cache) => {
22847                let base_len = cache.kv[il].as_ref().unwrap().len;
22848                // Read per layer call, never in a measured default.
22849                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
22850                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
22851                let off = if swa {
22852                    let raw = base_len.saturating_sub(win - 1);
22853                    if legacy_tkv || legacy_calllocal {
22854                        raw
22855                    } else {
22856                        raw & !31usize
22857                    }
22858                } else {
22859                    0
22860                };
22861                {
22862                    let kvl = cache.kv[il].as_mut().unwrap();
22863                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
22864                    let write_row = e.prepare_kv_append(kvl, off, t)?;
22865                    e.append_kv_quantized_rows(
22866                        &k,
22867                        &v,
22868                        &mut kvl.k,
22869                        &mut kvl.v,
22870                        write_row,
22871                        t,
22872                        kvl.kv_dim_k,
22873                        kvl.kv_dim_v,
22874                        kvl.k_tok_bytes,
22875                        kvl.v_tok_bytes,
22876                        crate::Engine::kv_fp8_on(),
22877                    )?;
22878                    kvl.len += t;
22879                    let new_len = kvl.len as i32;
22880                    e.set_i32_one(&mut kvl.len_d, new_len)?;
22881                }
22882                let kvl = cache.kv[il].as_ref().unwrap();
22883                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
22884                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
22885                // unaligned view offset here. Both halves are load-bearing for the canaries:
22886                // on the FA default the predicate arms agree bitwise wherever they can differ
22887                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
22888                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
22889                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
22890                // on the current FA path: its tile grid starts at the chunk/call boundary.
22891                // SWA: trim the view to the oldest key any query in this chunk can reach —
22892                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
22893                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
22894                // kernel's online-softmax recurrence groups keys into BK tiles relative to
22895                // the VIEW START — so an unaligned off regroups the same absolute keys into
22896                // different tiles at different chunk sizes = different (m,l) rounding =
22897                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
22898                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
22899                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
22900                // size; the <=31 extra leading keys are older than EVERY query's window
22901                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
22902                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
22903                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
22904                // the floor arm's bits do not move either (gated: G2f, battery 2).
22905                let t_kv = base_len + t - off;
22906                let physical = kvl.physical_rows(off, off + t_kv)?;
22907                let k_view = e.view_u8_range(
22908                    &kvl.k,
22909                    physical.start * kvl.k_tok_bytes,
22910                    physical.end * kvl.k_tok_bytes,
22911                );
22912                let v_view = e.view_u8_range(
22913                    &kvl.v,
22914                    physical.start * kvl.v_tok_bytes,
22915                    physical.end * kvl.v_tok_bytes,
22916                );
22917                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
22918                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
22919                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
22920                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
22921                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
22922                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
22923                // construction, so the invariance assertion MUST break under it (the seam whose
22924                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
22925                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
22926                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
22927                // cached (probes flip it in-process). Never on in a measured default run.
22928                let swa_naive = if legacy_tkv {
22929                    t_kv > win
22930                } else {
22931                    seq_end > win
22932                };
22933                if swa && swa_naive {
22934                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
22935                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
22936                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
22937                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
22938                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
22939                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
22940                    // identically to the unwindowed one modulo the mask, which is the point.
22941                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
22942                    // selected on `seq_end` like every arm here, so the class is uniform for
22943                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
22944                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
22945                    // the f32 floor (the previous numeric config, kept as the A/B seam).
22946                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
22947                        e.sdpa_naive_w_quantized_view(
22948                            &q,
22949                            &k_view,
22950                            &v_view,
22951                            &mut attn,
22952                            hd,
22953                            nh,
22954                            nkv,
22955                            t,
22956                            t_kv,
22957                            scale,
22958                            true,
22959                            win,
22960                            kvl.k_tok_bytes,
22961                            kvl.v_tok_bytes,
22962                        )?;
22963                    } else {
22964                        e.fa_prefill_view_ws_w_hd128(
22965                            &q,
22966                            &k_view,
22967                            &v_view,
22968                            &mut attn,
22969                            hd,
22970                            nh,
22971                            nkv,
22972                            t,
22973                            t_kv,
22974                            scale,
22975                            true,
22976                            win,
22977                            kvl.k_tok_bytes,
22978                            kvl.v_tok_bytes,
22979                        )?;
22980                    }
22981                } else if std::env::var("MEMRA_NOFA").is_ok() {
22982                    e.sdpa_naive_quantized_view(
22983                        &q,
22984                        &k_view,
22985                        &v_view,
22986                        &mut attn,
22987                        hd,
22988                        nh,
22989                        nkv,
22990                        t,
22991                        t_kv,
22992                        scale,
22993                        true,
22994                        kvl.k_tok_bytes,
22995                        kvl.v_tok_bytes,
22996                    )?;
22997                } else {
22998                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
22999                    // reach past the window, so the window mask is a no-op under causal and every
23000                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
23001                    // request either way, which is what makes the chunk size arithmetic-free.
23002                    e.fa_prefill_view_ws(
23003                        &q,
23004                        &k_view,
23005                        &v_view,
23006                        &mut attn,
23007                        hd,
23008                        nh,
23009                        nkv,
23010                        t,
23011                        t_kv,
23012                        scale,
23013                        true,
23014                        kvl.k_tok_bytes,
23015                        kvl.v_tok_bytes,
23016                        crate::Engine::kv_fp8_on(),
23017                    )?;
23018                }
23019            }
23020            None => {
23021                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
23022                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
23023                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
23024                // seq_end here too or it re-opens the same door.
23025                debug_assert_eq!(
23026                    seq_end, t,
23027                    "step35 cacheless prefill is monolithic (seq_end == t)"
23028                );
23029                if swa && seq_end > win {
23030                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
23031                } else if std::env::var("MEMRA_NOFA").is_ok() {
23032                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
23033                } else {
23034                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
23035                }
23036            }
23037        }
23038
23039        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
23040        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
23041        let gw = fa
23042            .attn_gate
23043            .as_ref()
23044            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
23045        let gt_owned = if gt_pre.is_none() {
23046            Some(e.matmul(
23047                gw,
23048                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
23049                t,
23050            )?)
23051        } else {
23052            None
23053        };
23054        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
23055        let mut ag = e.uninit(t * nh * hd)?;
23056        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
23057        Ok(ag)
23058    }
23059
23060    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
23061    /// `forward_last`, t2probe). Post-`wo`.
23062    pub(crate) fn step35_attn(
23063        &self,
23064        e: &Engine,
23065        fa: &FullAttnLayer,
23066        h: &CudaSlice<f32>,
23067        pos_d: &CudaSlice<i32>,
23068        t: usize,
23069        il: usize,
23070    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23071        let g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
23072            Some(g3) => g3,
23073            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
23074        };
23075        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
23076        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
23077        self.full_attn_o(e, fa, &ag, t)
23078    }
23079
23080    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
23081    /// resident quantized cache, attend through the cache view). Post-`wo`.
23082    ///
23083    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
23084    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
23085    /// own extent.
23086    #[allow(clippy::too_many_arguments)]
23087    pub(crate) fn step35_attn_prime(
23088        &self,
23089        e: &Engine,
23090        fa: &FullAttnLayer,
23091        h: &CudaSlice<f32>,
23092        hx: Option<&CudaSlice<u8>>,
23093        pos_d: &CudaSlice<i32>,
23094        t: usize,
23095        cache: &mut Cache,
23096        il: usize,
23097        seq_end: usize,
23098    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23099        if step_tp_prefill_enabled()? && fa.step_tp_qkv.is_some() {
23100            if hx.is_some() {
23101                return Err(
23102                    "rank-local Step prefill preserves BF16 activations and refuses the q8_1 \
23103                     pre-quantized prime path"
23104                        .into(),
23105                );
23106            }
23107            return self.step35_tp_prefill_attn_resident(e, fa, il, h, pos_d, t, cache, seq_end);
23108        }
23109        let g3 = if fa.step_tp_qkv.is_some() {
23110            if hx.is_some() {
23111                return Err(
23112                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
23113                     pre-quantized prime path"
23114                        .into(),
23115                );
23116            }
23117            self.full_attn_tp_qkv(e, fa, h, t)?
23118                .expect("Step Q/K/V TP disappeared after the presence check")
23119        } else {
23120            match hx {
23121                Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
23122                None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
23123            }
23124        };
23125        let ag =
23126            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
23127        self.full_attn_o(e, fa, &ag, t)
23128    }
23129
23130    fn ensure_step_tp_kv_cache(
23131        &self,
23132        e: &Engine,
23133        fa: &FullAttnLayer,
23134        il: usize,
23135        cache: &mut Cache,
23136    ) -> Result<bool, Box<dyn std::error::Error>> {
23137        let tp = fa
23138            .step_tp_qkv
23139            .as_ref()
23140            .ok_or("Step TP cache hydration lost its resident projections")?;
23141        let geometry = self.cfg.full_attention_geometry_at(il as u32);
23142        let window = geometry.window.map(|window| window as usize);
23143        let ranks = tp.runtime.devices().len();
23144        let head_dim = geometry.head_dim_k as usize;
23145        let kv_heads = geometry.n_head_kv as usize;
23146        let max_ctx = cache.max_ctx;
23147
23148        if cache.tp_kv[il].is_some() {
23149            return Ok(false);
23150        }
23151        let local = cache.kv[il]
23152            .as_ref()
23153            .ok_or_else(|| format!("Step TP layer {il} has no owning-stage KV cache"))?;
23154        if local.kv_dim_k != kv_heads * head_dim || local.kv_dim_v != kv_heads * head_dim {
23155            return Err(format!(
23156                "Step TP layer {il} local KV geometry k={} v={} != {}",
23157                local.kv_dim_k,
23158                local.kv_dim_v,
23159                kv_heads * head_dim
23160            )
23161            .into());
23162        }
23163        let resident_start = window
23164            .map(|window| local.len.saturating_sub(window.saturating_sub(1)) & !31usize)
23165            .unwrap_or(0);
23166        let resident_rows = local.len - resident_start;
23167        let physical = local.physical_rows(resident_start, local.len)?;
23168        let k_rows = if resident_rows == 0 {
23169            Vec::new()
23170        } else {
23171            e.dtoh_u8_view(&e.view_u8_range(
23172                &local.k,
23173                physical.start * local.k_tok_bytes,
23174                physical.end * local.k_tok_bytes,
23175            ))?
23176        };
23177        let v_rows = if resident_rows == 0 {
23178            Vec::new()
23179        } else {
23180            e.dtoh_u8_view(&e.view_u8_range(
23181                &local.v,
23182                physical.start * local.v_tok_bytes,
23183                physical.end * local.v_tok_bytes,
23184            ))?
23185        };
23186        let mut distributed = match window {
23187            Some(window) => tp.runtime.allocate_tp_swa_kv_cache(
23188                kv_heads * head_dim,
23189                kv_heads * head_dim,
23190                max_ctx,
23191                window,
23192            )?,
23193            None => tp.runtime.allocate_tp_kv_cache(
23194                kv_heads * head_dim,
23195                kv_heads * head_dim,
23196                max_ctx,
23197            )?,
23198        };
23199        if distributed.k_tok_bytes() * ranks != local.k_tok_bytes
23200            || distributed.v_tok_bytes() * ranks != local.v_tok_bytes
23201        {
23202            return Err(format!(
23203                "Step TP layer {il} distributed/local KV token bytes disagree: \
23204                 k={}x{ranks}/{} v={}x{ranks}/{}",
23205                distributed.k_tok_bytes(),
23206                local.k_tok_bytes,
23207                distributed.v_tok_bytes(),
23208                local.v_tok_bytes,
23209            )
23210            .into());
23211        }
23212        tp.runtime.hydrate_tp_kv_cache_from(
23213            &mut distributed,
23214            local.len,
23215            resident_start,
23216            &k_rows,
23217            &v_rows,
23218        )?;
23219        cache.tp_kv[il] = Some(distributed);
23220        Ok(true)
23221    }
23222
23223    #[allow(clippy::too_many_arguments)]
23224    #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
23225    fn step35_tp_prefill_attn_resident(
23226        &self,
23227        e: &Engine,
23228        fa: &FullAttnLayer,
23229        il: usize,
23230        h: &CudaSlice<f32>,
23231        pos_d: &CudaSlice<i32>,
23232        tokens: usize,
23233        cache: &mut Cache,
23234        seq_end: usize,
23235    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23236        let tp = fa
23237            .step_tp_qkv
23238            .as_ref()
23239            .ok_or("Step TP prefill lost its resident projections")?;
23240        let attention = tp
23241            .attention
23242            .as_ref()
23243            .ok_or("Step TP prefill lost its resident attention auxiliaries")?;
23244        let ranks = tp.runtime.devices().len();
23245        if !step_tp_prefill_shape(
23246            true,
23247            tokens,
23248            ranks,
23249            tp.runtime.native_p2p(),
23250            true,
23251            crate::Engine::kv_fp8_on(),
23252        ) {
23253            return Err(format!(
23254                "rank-local Step prefill requires tokens>={PRIME_MIN_T}, TP2/TP4 native P2P, \
23255                 rank-local attention, and q8_0/q5_1 KV; got tokens={tokens} ranks={ranks} \
23256                 native_p2p={} fp8_kv={}",
23257                tp.runtime.native_p2p(),
23258                crate::Engine::kv_fp8_on(),
23259            )
23260            .into());
23261        }
23262        for seam in [
23263            "MEMRA_STEP35_SWA_TKV",
23264            "MEMRA_PRIME_CALLLOCAL",
23265            "MEMRA_PRIME_F32CHUNK0",
23266        ] {
23267            if std::env::var(seam).as_deref() == Ok("1") {
23268                return Err(format!(
23269                    "rank-local Step prefill has not qualified the legacy seam {seam}=1"
23270                )
23271                .into());
23272            }
23273        }
23274
23275        let geometry = self.cfg.full_attention_geometry_at(il as u32);
23276        let window = geometry.window.map(|window| window as usize);
23277        let head_dim = geometry.head_dim_k as usize;
23278        let heads = geometry.n_head as usize;
23279        let kv_heads = geometry.n_head_kv as usize;
23280        if heads % ranks != 0 || kv_heads % ranks != 0 {
23281            return Err(format!(
23282                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
23283            )
23284            .into());
23285        }
23286        let local_heads = heads / ranks;
23287        let local_kv_heads = kv_heads / ranks;
23288        let local_kv_dim = local_kv_heads * head_dim;
23289        let hidden = self.cfg.n_embd as usize;
23290        let expected_input = tokens
23291            .checked_mul(hidden)
23292            .ok_or("Step TP prefill input size overflow")?;
23293        if h.len() < expected_input {
23294            return Err(format!(
23295                "Step TP prefill input {} is shorter than {tokens}x{hidden}",
23296                h.len()
23297            )
23298            .into());
23299        }
23300        let positions = e.dtoh_i32(pos_d)?;
23301        if positions.len() != tokens {
23302            return Err(format!(
23303                "rank-local Step prefill positions {} != tokens {tokens}",
23304                positions.len()
23305            )
23306            .into());
23307        }
23308
23309        let mut active_input = e.uninit(expected_input)?;
23310        e.copy_view_into(
23311            &mut active_input,
23312            0,
23313            &h.slice(0..expected_input),
23314            expected_input,
23315        )?;
23316        let mut input = tp.runtime.allocate_replicated_device_rows(tokens, hidden)?;
23317        // PRODUCER FENCE (2026-08-20 flake fix): active_input was written on the MODEL engine's
23318        // stream; the refresh below reads it from the runtime root engine's stream (same device,
23319        // different stream). Unfenced, the peer read can overtake the in-flight copy — the
23320        // layer-count-amplified arm of the boot flake.
23321        e.stream().synchronize()?;
23322        tp.runtime
23323            .refresh_replicated_device_rows_from_root(&mut input, &active_input)?;
23324        let q_raw = tp
23325            .runtime
23326            .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &input)?;
23327        let k_raw = tp
23328            .runtime
23329            .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &input)?;
23330        let v_raw = tp
23331            .runtime
23332            .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &input)?;
23333        let mut q = Vec::with_capacity(ranks);
23334        let mut k = Vec::with_capacity(ranks);
23335        for rank in 0..ranks {
23336            let engine = tp
23337                .runtime
23338                .rank_engine(rank)
23339                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23340            let _main = engine.gpu.enter_main()?;
23341            let mut q_rank = engine.uninit(tokens * local_heads * head_dim)?;
23342            engine.rms_norm(
23343                &q_raw[rank],
23344                &attention.q_norm[rank],
23345                &mut q_rank,
23346                head_dim,
23347                tokens * local_heads,
23348                self.cfg.rms_eps,
23349            )?;
23350            let mut k_rank = engine.uninit(tokens * local_kv_dim)?;
23351            engine.rms_norm(
23352                &k_raw[rank],
23353                &attention.k_norm[rank],
23354                &mut k_rank,
23355                head_dim,
23356                tokens * local_kv_heads,
23357                self.cfg.rms_eps,
23358            )?;
23359            let position = engine.htod_i32(&positions)?;
23360            let rope_freqs = if geometry.rope_factors {
23361                self.step35_aux
23362                    .as_ref()
23363                    .and_then(|aux| aux.rope_freqs(engine))
23364            } else {
23365                None
23366            };
23367            engine.rope_neox2(
23368                &mut q_rank,
23369                &mut k_rank,
23370                &position,
23371                head_dim,
23372                geometry.n_rot as usize,
23373                local_heads,
23374                local_kv_heads,
23375                tokens,
23376                geometry.rope_base,
23377                1.0,
23378                rope_freqs,
23379            )?;
23380            q.push(q_rank);
23381            k.push(k_rank);
23382        }
23383
23384        let gate_weight = fa
23385            .attn_gate
23386            .as_ref()
23387            .ok_or("step35 layer is missing attn_gate.weight")?;
23388        let gate = e.dtoh(&e.matmul(gate_weight, h, tokens)?)?;
23389        if gate.len() != tokens * heads {
23390            return Err(format!(
23391                "Step TP layer {il} gate output {} != {tokens}x{heads}",
23392                gate.len()
23393            )
23394            .into());
23395        }
23396
23397        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
23398        let base_len = cache.kv[il]
23399            .as_ref()
23400            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
23401            .len;
23402        let distributed = cache.tp_kv[il]
23403            .as_ref()
23404            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
23405        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
23406            return Err(format!(
23407                "Step TP layer {il} cache lengths diverged before prefill: \
23408                 local={base_len} distributed={}/{}",
23409                distributed.committed_len(),
23410                distributed.staged_len()
23411            )
23412            .into());
23413        }
23414        let target_len = base_len
23415            .checked_add(tokens)
23416            .ok_or("Step TP prefill cache length overflow")?;
23417        if target_len > cache.max_ctx {
23418            return Err(format!(
23419                "Step TP layer {il} prefill exceeds cache: {base_len}+{tokens}>{}",
23420                cache.max_ctx
23421            )
23422            .into());
23423        }
23424        if seq_end < target_len {
23425            return Err(format!(
23426                "Step TP layer {il} request end {seq_end} precedes chunk end {target_len}"
23427            )
23428            .into());
23429        }
23430
23431        let transaction = cache.tp_kv[il]
23432            .as_mut()
23433            .expect("distributed cache checked above")
23434            .begin_transaction()?;
23435        if let Err(error) = tp.runtime.append_tp_kv_transaction(
23436            cache.tp_kv[il]
23437                .as_mut()
23438                .expect("distributed cache checked above"),
23439            transaction,
23440            &k,
23441            &v_raw,
23442            tokens,
23443        ) {
23444            let _ = tp.runtime.rollback_tp_kv_transaction(
23445                cache.tp_kv[il]
23446                    .as_mut()
23447                    .expect("distributed cache checked above"),
23448                transaction,
23449            );
23450            return Err(error);
23451        }
23452
23453        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23454            let distributed = cache.tp_kv[il]
23455                .as_ref()
23456                .expect("distributed cache checked above");
23457            let staged_len = distributed.staged_len();
23458            let view_start = window
23459                .map(|window| base_len.saturating_sub(window.saturating_sub(1)) & !31usize)
23460                .unwrap_or(0);
23461            let physical = distributed.physical_range(view_start, staged_len)?;
23462            let t_kv = staged_len - view_start;
23463            let swa_naive = window.is_some_and(|window| seq_end > window);
23464            let mut gated = Vec::with_capacity(ranks);
23465            #[allow(clippy::needless_range_loop)]
23466            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
23467            for rank in 0..ranks {
23468                let engine = tp
23469                    .runtime
23470                    .rank_engine(rank)
23471                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23472                let _main = engine.gpu.enter_main()?;
23473                let rank_cache = distributed
23474                    .rank(rank)
23475                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23476                let k_view = engine.view_u8_range(
23477                    rank_cache.k(),
23478                    physical.start * distributed.k_tok_bytes(),
23479                    physical.end * distributed.k_tok_bytes(),
23480                );
23481                let v_view = engine.view_u8_range(
23482                    rank_cache.v(),
23483                    physical.start * distributed.v_tok_bytes(),
23484                    physical.end * distributed.v_tok_bytes(),
23485                );
23486                let mut attention_out = engine.uninit(tokens * local_heads * head_dim)?;
23487                if swa_naive {
23488                    let window = window.expect("SWA predicate requires a window");
23489                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
23490                        engine.sdpa_naive_w_quantized_view(
23491                            &q[rank],
23492                            &k_view,
23493                            &v_view,
23494                            &mut attention_out,
23495                            head_dim,
23496                            local_heads,
23497                            local_kv_heads,
23498                            tokens,
23499                            t_kv,
23500                            geometry.attention_scale(),
23501                            true,
23502                            window,
23503                            distributed.k_tok_bytes(),
23504                            distributed.v_tok_bytes(),
23505                        )?;
23506                    } else {
23507                        engine.fa_prefill_view_ws_w_hd128(
23508                            &q[rank],
23509                            &k_view,
23510                            &v_view,
23511                            &mut attention_out,
23512                            head_dim,
23513                            local_heads,
23514                            local_kv_heads,
23515                            tokens,
23516                            t_kv,
23517                            geometry.attention_scale(),
23518                            true,
23519                            window,
23520                            distributed.k_tok_bytes(),
23521                            distributed.v_tok_bytes(),
23522                        )?;
23523                    }
23524                } else if std::env::var("MEMRA_NOFA").is_ok() {
23525                    engine.sdpa_naive_quantized_view(
23526                        &q[rank],
23527                        &k_view,
23528                        &v_view,
23529                        &mut attention_out,
23530                        head_dim,
23531                        local_heads,
23532                        local_kv_heads,
23533                        tokens,
23534                        t_kv,
23535                        geometry.attention_scale(),
23536                        true,
23537                        distributed.k_tok_bytes(),
23538                        distributed.v_tok_bytes(),
23539                    )?;
23540                } else {
23541                    engine.fa_prefill_view_ws(
23542                        &q[rank],
23543                        &k_view,
23544                        &v_view,
23545                        &mut attention_out,
23546                        head_dim,
23547                        local_heads,
23548                        local_kv_heads,
23549                        tokens,
23550                        t_kv,
23551                        geometry.attention_scale(),
23552                        true,
23553                        distributed.k_tok_bytes(),
23554                        distributed.v_tok_bytes(),
23555                        false,
23556                    )?;
23557                }
23558
23559                let gate_start = rank * local_heads;
23560                let mut gate_rank = Vec::with_capacity(tokens * local_heads);
23561                for token in 0..tokens {
23562                    let start = token * heads + gate_start;
23563                    gate_rank.extend_from_slice(&gate[start..start + local_heads]);
23564                }
23565                let gate_rank = engine.htod(&gate_rank)?;
23566                let mut gated_rank = engine.uninit(tokens * local_heads * head_dim)?;
23567                engine.attn_head_gate(
23568                    &attention_out,
23569                    &gate_rank,
23570                    &mut gated_rank,
23571                    None,
23572                    head_dim,
23573                    local_heads,
23574                    tokens,
23575                )?;
23576                gated.push(gated_rank);
23577            }
23578            for rank in 1..ranks {
23579                let engine = tp
23580                    .runtime
23581                    .rank_engine(rank)
23582                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23583                let _main = engine.gpu.enter_main()?;
23584                engine.stream().synchronize()?;
23585            }
23586
23587            let (output, k_shadow, v_shadow) = if tp.runtime.bulk_p2p() {
23588                let output = tp
23589                    .runtime
23590                    .step_bf16_row_parallel_resident_root_device(&tp.o, &gated, tokens)?;
23591                let k_shadow =
23592                    tp.runtime
23593                        .gather_native_column_shards_device(&k, tokens, local_kv_dim)?;
23594                let v_shadow =
23595                    tp.runtime
23596                        .gather_native_column_shards_device(&v_raw, tokens, local_kv_dim)?;
23597                let root = tp
23598                    .runtime
23599                    .rank_engine(0)
23600                    .ok_or("Step TP prefill lost its root engine")?;
23601                let _main = root.gpu.enter_main()?;
23602                root.stream().synchronize()?;
23603                (output, k_shadow, v_shadow)
23604            } else {
23605                let attention = tp.runtime.gather_native_column_shards(
23606                    &gated,
23607                    tokens,
23608                    local_heads * head_dim,
23609                )?;
23610                let output = tp
23611                    .runtime
23612                    .step_bf16_row_parallel_resident_native(&tp.o, &attention, tokens)?;
23613                let k_shadow = tp
23614                    .runtime
23615                    .gather_native_column_shards(&k, tokens, local_kv_dim)?;
23616                let v_shadow =
23617                    tp.runtime
23618                        .gather_native_column_shards(&v_raw, tokens, local_kv_dim)?;
23619                (e.htod(&output)?, e.htod(&k_shadow)?, e.htod(&v_shadow)?)
23620            };
23621            let local = cache.kv[il]
23622                .as_mut()
23623                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
23624            if local.len != base_len {
23625                return Err(format!(
23626                    "Step TP layer {il} local cache changed during prefill: \
23627                     len={} base={base_len}",
23628                    local.len
23629                )
23630                .into());
23631            }
23632            let retain_from = window
23633                .map(|window| {
23634                    let staged_retain = staged_len.saturating_sub(window) & !31usize;
23635                    let rollback_retain =
23636                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
23637                    staged_retain.min(rollback_retain)
23638                })
23639                .unwrap_or(0);
23640            let write_row = e.prepare_kv_append(local, retain_from, tokens)?;
23641            e.append_kv_quantized_rows(
23642                &k_shadow,
23643                &v_shadow,
23644                &mut local.k,
23645                &mut local.v,
23646                write_row,
23647                tokens,
23648                local.kv_dim_k,
23649                local.kv_dim_v,
23650                local.k_tok_bytes,
23651                local.v_tok_bytes,
23652                false,
23653            )?;
23654            local.len = staged_len;
23655            e.set_i32_one(&mut local.len_d, staged_len as i32)?;
23656            Ok(output)
23657        })();
23658
23659        let output = match staged {
23660            Ok(output) => output,
23661            Err(error) => {
23662                let _ = tp.runtime.rollback_tp_kv_transaction(
23663                    cache.tp_kv[il]
23664                        .as_mut()
23665                        .expect("distributed cache checked above"),
23666                    transaction,
23667                );
23668                if let Some(local) = cache.kv[il].as_mut() {
23669                    local.len = base_len;
23670                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
23671                }
23672                return Err(error);
23673            }
23674        };
23675        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
23676            cache.tp_kv[il]
23677                .as_mut()
23678                .expect("distributed cache checked above"),
23679            transaction,
23680            tokens,
23681        ) {
23682            let _ = tp.runtime.rollback_tp_kv_transaction(
23683                cache.tp_kv[il]
23684                    .as_mut()
23685                    .expect("distributed cache checked above"),
23686                transaction,
23687            );
23688            let local = cache.kv[il].as_mut().expect("local cache checked above");
23689            local.len = base_len;
23690            e.set_i32_one(&mut local.len_d, base_len as i32)?;
23691            return Err(error);
23692        }
23693
23694        let committed = cache.tp_kv[il]
23695            .as_ref()
23696            .expect("distributed cache checked above")
23697            .committed_len();
23698        let local_len = cache.kv[il]
23699            .as_ref()
23700            .expect("local cache checked above")
23701            .len;
23702        if committed != local_len {
23703            return Err(format!(
23704                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
23705            )
23706            .into());
23707        }
23708        eprintln!(
23709            "[step-tp-prefill-attn] execute layer={} devices={:?} tokens={tokens} \
23710             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
23711             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
23712             attention_scope={} input_path=root-device-replicated gate_tensor_parallel=false \
23713             gate_shards=host-canonical o_tensor_parallel=true local_cache_shadow=true \
23714             cache_commit=chunk transport={} native_p2p=true bulk_p2p={} \
23715             output={} performance_claim=false",
23716            tp.layer,
23717            tp.devices,
23718            hydrated,
23719            if window.is_some() {
23720                "rank-local-swa-ring"
23721            } else {
23722                "rank-local-global"
23723            },
23724            tp.runtime.transport_label(),
23725            tp.runtime.bulk_p2p(),
23726            if tp.runtime.bulk_p2p() {
23727                "root-device"
23728            } else {
23729                "root-readback"
23730            },
23731        );
23732        Ok(output)
23733    }
23734
23735    fn step35_tp_decode_attn_resident(
23736        &self,
23737        e: &Engine,
23738        fa: &FullAttnLayer,
23739        il: usize,
23740        h: &CudaSlice<f32>,
23741        pos_d: &CudaSlice<i32>,
23742        cache: &mut Cache,
23743    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23744        // MEMRA_STEP_TP_TIMING=1: cumulative wall-clock of the rank-local TP attention decode,
23745        // printed every ~10 decode steps' worth of layers — the wall-decomposition twin of the
23746        // nvfp4-dev-routes counter.
23747        static ATTN_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23748        static ATTN_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23749        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
23750        let started = timing.then(std::time::Instant::now);
23751        let result = if crate::tp::step_tp_decode_v2_enabled()? {
23752            self.step35_tp_decode_attn_resident_v2(e, fa, il, h, pos_d, cache)
23753        } else {
23754            self.step35_tp_decode_attn_resident_inner(e, fa, il, h, pos_d, cache)
23755        };
23756        if let Some(started) = started {
23757            use std::sync::atomic::Ordering;
23758            let ns = ATTN_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
23759                + started.elapsed().as_nanos() as u64;
23760            let calls = ATTN_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
23761            if calls.is_multiple_of(430) {
23762                eprintln!(
23763                    "[step-tp-attn-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
23764                    ns as f64 / 1.0e6,
23765                    ns as f64 / calls as f64 / 1.0e3,
23766                );
23767            }
23768        }
23769        result
23770    }
23771
23772    #[allow(clippy::too_many_arguments)]
23773    fn step35_tp_decode_attn_resident_inner(
23774        &self,
23775        e: &Engine,
23776        fa: &FullAttnLayer,
23777        il: usize,
23778        h: &CudaSlice<f32>,
23779        pos_d: &CudaSlice<i32>,
23780        cache: &mut Cache,
23781    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23782        // MEMRA_STEP_TP_TIMING=1 phase decomposition of the 1550us/layer decode wall. Each lap
23783        // drains every stream so queued async work is billed to the phase that queued it — the
23784        // drains perturb absolute wall, but v1 already ends most phases on a host sync, so the
23785        // relative split is honest. Timing OFF is the measured configuration: zero extra syncs.
23786        static T_POS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23787        static T_QKV: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23788        static T_NORMROPE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23789        static T_GATE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23790        static T_APPEND: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23791        static T_ATTN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23792        static T_OPROJ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23793        static T_SHADOW: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23794        static T_PHASE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23795        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
23796        #[allow(clippy::manual_is_multiple_of)] // allow: divisor is runtime-derived; the modulo form keeps a zero divisor loud (a panic), where is_multiple_of would return false silently
23797        fn lap(
23798            runtime: &crate::tp::TpE4m3HostBounce,
23799            e: &Engine,
23800            timer: &std::sync::atomic::AtomicU64,
23801            started: &mut Option<std::time::Instant>,
23802        ) -> Result<(), Box<dyn std::error::Error>> {
23803            let Some(start) = started.as_mut() else {
23804                return Ok(());
23805            };
23806            for rank in 0..runtime.devices().len() {
23807                if let Some(engine) = runtime.rank_engine(rank) {
23808                    let _main = engine.gpu.enter_main()?;
23809                    engine.stream().synchronize()?;
23810                }
23811            }
23812            e.stream().synchronize()?;
23813            timer.fetch_add(
23814                start.elapsed().as_nanos() as u64,
23815                std::sync::atomic::Ordering::Relaxed,
23816            );
23817            *start = std::time::Instant::now();
23818            Ok(())
23819        }
23820        let tp = fa
23821            .step_tp_qkv
23822            .as_ref()
23823            .ok_or("Step TP decode lost its resident projections")?;
23824        let attention = tp
23825            .attention
23826            .as_ref()
23827            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
23828        if !tp.runtime.native_p2p() {
23829            return Err("rank-local Step attention requires native P2P".into());
23830        }
23831        if crate::Engine::kv_fp8_on() {
23832            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
23833        }
23834
23835        let geometry = self.step35_geom(il);
23836        let window = geometry.window.map(|window| window as usize);
23837        let ranks = tp.runtime.devices().len();
23838        let head_dim = geometry.head_dim_k as usize;
23839        let heads = geometry.n_head as usize;
23840        let kv_heads = geometry.n_head_kv as usize;
23841        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
23842            return Err(format!(
23843                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
23844            )
23845            .into());
23846        }
23847        let local_heads = heads / ranks;
23848        let local_kv_heads = kv_heads / ranks;
23849        let local_kv_dim = local_kv_heads * head_dim;
23850        let max_ctx = cache.max_ctx;
23851
23852        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
23853
23854        let base_len = cache.kv[il]
23855            .as_ref()
23856            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
23857            .len;
23858        let distributed = cache.tp_kv[il]
23859            .as_ref()
23860            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
23861        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
23862            return Err(format!(
23863                "Step TP layer {il} cache lengths diverged before decode: \
23864                 local={base_len} distributed={}/{}",
23865                distributed.committed_len(),
23866                distributed.staged_len()
23867            )
23868            .into());
23869        }
23870
23871        let mut lap_start = timing.then(std::time::Instant::now);
23872        let positions = e.dtoh_i32(pos_d)?;
23873        if positions.len() != 1 {
23874            return Err(format!(
23875                "rank-local Step decode requires one position, got {}",
23876                positions.len()
23877            )
23878            .into());
23879        }
23880        lap(&tp.runtime, e, &T_POS, &mut lap_start)?;
23881        let (q_raw, k_raw, v_raw, input_path) = if let Some(decode_input) =
23882            attention.decode_input.as_ref()
23883        {
23884            let mut decode_input = decode_input
23885                .lock()
23886                .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
23887            // PRODUCER FENCE (2026-08-20 flake fix): h is the layer hidden written on the MODEL
23888            // engine's stream; the refresh reads it from the runtime root engine's stream. This
23889            // per-layer-per-token seam is the layer-count-amplified arm of the boot flake.
23890            e.stream().synchronize()?;
23891            tp.runtime
23892                .refresh_replicated_device_rows_from_root(&mut decode_input, h)?;
23893            let q_raw = tp
23894                .runtime
23895                .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &decode_input)?;
23896            let k_raw = tp
23897                .runtime
23898                .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &decode_input)?;
23899            let v_raw = tp
23900                .runtime
23901                .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &decode_input)?;
23902            (q_raw, k_raw, v_raw, "root-device-replicated")
23903        } else {
23904            let activation = e.dtoh(h)?;
23905            let q_raw =
23906                tp.runtime
23907                    .bf16_column_parallel_resident_device_shards(&tp.q, &activation, 1)?;
23908            let k_raw =
23909                tp.runtime
23910                    .bf16_column_parallel_resident_device_shards(&tp.k, &activation, 1)?;
23911            let v_raw =
23912                tp.runtime
23913                    .bf16_column_parallel_resident_device_shards(&tp.v, &activation, 1)?;
23914            (q_raw, k_raw, v_raw, "host-replicated")
23915        };
23916        lap(&tp.runtime, e, &T_QKV, &mut lap_start)?;
23917        let mut q = Vec::with_capacity(ranks);
23918        let mut k = Vec::with_capacity(ranks);
23919        for rank in 0..ranks {
23920            let engine = tp
23921                .runtime
23922                .rank_engine(rank)
23923                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23924            let _main = engine.gpu.enter_main()?;
23925            let mut q_rank = engine.uninit(local_heads * head_dim)?;
23926            engine.rms_norm(
23927                &q_raw[rank],
23928                &attention.q_norm[rank],
23929                &mut q_rank,
23930                head_dim,
23931                local_heads,
23932                self.cfg.rms_eps,
23933            )?;
23934            let mut k_rank = engine.uninit(local_kv_dim)?;
23935            engine.rms_norm(
23936                &k_raw[rank],
23937                &attention.k_norm[rank],
23938                &mut k_rank,
23939                head_dim,
23940                local_kv_heads,
23941                self.cfg.rms_eps,
23942            )?;
23943            let position = engine.htod_i32(&positions)?;
23944            let rope_freqs = if geometry.rope_factors {
23945                self.step35_aux
23946                    .as_ref()
23947                    .and_then(|aux| aux.rope_freqs(engine))
23948            } else {
23949                None
23950            };
23951            engine.rope_neox2(
23952                &mut q_rank,
23953                &mut k_rank,
23954                &position,
23955                head_dim,
23956                geometry.n_rot as usize,
23957                local_heads,
23958                local_kv_heads,
23959                1,
23960                geometry.rope_base,
23961                1.0,
23962                rope_freqs,
23963            )?;
23964            q.push(q_rank);
23965            k.push(k_rank);
23966        }
23967        lap(&tp.runtime, e, &T_NORMROPE, &mut lap_start)?;
23968
23969        let gate_weight = fa
23970            .attn_gate
23971            .as_ref()
23972            .ok_or("step35 layer is missing attn_gate.weight")?;
23973        let gate = e.matmul(gate_weight, h, 1)?;
23974        let gate = e.dtoh(&gate)?;
23975        if gate.len() != heads {
23976            return Err(format!("Step TP layer {il} gate output {} != {heads}", gate.len()).into());
23977        }
23978        lap(&tp.runtime, e, &T_GATE, &mut lap_start)?;
23979
23980        let transaction = cache.tp_kv[il]
23981            .as_mut()
23982            .expect("distributed cache checked above")
23983            .begin_transaction()?;
23984        if let Err(error) = tp.runtime.append_tp_kv_transaction(
23985            cache.tp_kv[il]
23986                .as_mut()
23987                .expect("distributed cache checked above"),
23988            transaction,
23989            &k,
23990            &v_raw,
23991            1,
23992        ) {
23993            let _ = tp.runtime.rollback_tp_kv_transaction(
23994                cache.tp_kv[il]
23995                    .as_mut()
23996                    .expect("distributed cache checked above"),
23997                transaction,
23998            );
23999            return Err(error);
24000        }
24001        lap(&tp.runtime, e, &T_APPEND, &mut lap_start)?;
24002
24003        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24004            let distributed = cache.tp_kv[il]
24005                .as_ref()
24006                .expect("distributed cache checked above");
24007            let staged_len = distributed.staged_len();
24008            let view_start = window
24009                .map(|window| staged_len.saturating_sub(window))
24010                .unwrap_or(0);
24011            let physical = distributed.physical_range(view_start, staged_len)?;
24012            let t_kv = staged_len - view_start;
24013            let mut gated = Vec::with_capacity(ranks);
24014            #[allow(clippy::needless_range_loop)]
24015            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
24016            for rank in 0..ranks {
24017                let engine = tp
24018                    .runtime
24019                    .rank_engine(rank)
24020                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
24021                let _main = engine.gpu.enter_main()?;
24022                let rank_cache = distributed
24023                    .rank(rank)
24024                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
24025                let k_view = engine.view_u8_range(
24026                    rank_cache.k(),
24027                    physical.start * distributed.k_tok_bytes(),
24028                    physical.end * distributed.k_tok_bytes(),
24029                );
24030                let v_view = engine.view_u8_range(
24031                    rank_cache.v(),
24032                    physical.start * distributed.v_tok_bytes(),
24033                    physical.end * distributed.v_tok_bytes(),
24034                );
24035                let mut attention_out = engine.uninit(local_heads * head_dim)?;
24036                engine.fa_decode_kvmod(
24037                    &q[rank],
24038                    &k_view,
24039                    &v_view,
24040                    &mut attention_out,
24041                    head_dim,
24042                    local_heads,
24043                    local_kv_heads,
24044                    t_kv,
24045                    geometry.attention_scale(),
24046                    distributed.k_tok_bytes(),
24047                    distributed.v_tok_bytes(),
24048                    false,
24049                )?;
24050                let gate_start = rank * local_heads;
24051                let gate_rank = engine.htod(&gate[gate_start..gate_start + local_heads])?;
24052                let mut gated_rank = engine.uninit(local_heads * head_dim)?;
24053                engine.attn_head_gate(
24054                    &attention_out,
24055                    &gate_rank,
24056                    &mut gated_rank,
24057                    None,
24058                    head_dim,
24059                    local_heads,
24060                    1,
24061                )?;
24062                gated.push(gated_rank);
24063            }
24064            lap(&tp.runtime, e, &T_ATTN, &mut lap_start)?;
24065
24066            let gathered =
24067                tp.runtime
24068                    .gather_native_column_shards(&gated, 1, local_heads * head_dim)?;
24069            let output = tp
24070                .runtime
24071                .step_bf16_row_parallel_resident_native(&tp.o, &gathered, 1)?;
24072            let output = e.htod(&output)?;
24073            lap(&tp.runtime, e, &T_OPROJ, &mut lap_start)?;
24074
24075            let k_shadow = tp
24076                .runtime
24077                .gather_native_column_shards(&k, 1, local_kv_dim)?;
24078            let v_shadow = tp
24079                .runtime
24080                .gather_native_column_shards(&v_raw, 1, local_kv_dim)?;
24081            let k_shadow = e.htod(&k_shadow)?;
24082            let v_shadow = e.htod(&v_shadow)?;
24083            let local = cache.kv[il]
24084                .as_mut()
24085                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
24086            if local.len != base_len || base_len + 1 > max_ctx {
24087                return Err(format!(
24088                    "Step TP layer {il} local cache changed during decode: \
24089                     len={} base={base_len} max={max_ctx}",
24090                    local.len
24091                )
24092                .into());
24093            }
24094            let retain_from = window
24095                .map(|window| {
24096                    let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
24097                    let rollback_retain =
24098                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
24099                    staged_retain.min(rollback_retain)
24100                })
24101                .unwrap_or(0);
24102            let write_row = e.prepare_kv_append(local, retain_from, 1)?;
24103            e.append_kv_quantized(
24104                &k_shadow,
24105                &v_shadow,
24106                &mut local.k,
24107                &mut local.v,
24108                write_row,
24109                local.kv_dim_k,
24110                local.kv_dim_v,
24111                local.k_tok_bytes,
24112                local.v_tok_bytes,
24113                false,
24114            )?;
24115            local.len = base_len + 1;
24116            e.set_i32_one(&mut local.len_d, local.len as i32)?;
24117            Ok(output)
24118        })();
24119
24120        let output = match staged {
24121            Ok(output) => output,
24122            Err(error) => {
24123                let _ = tp.runtime.rollback_tp_kv_transaction(
24124                    cache.tp_kv[il]
24125                        .as_mut()
24126                        .expect("distributed cache checked above"),
24127                    transaction,
24128                );
24129                if let Some(local) = cache.kv[il].as_mut() {
24130                    local.len = base_len;
24131                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
24132                }
24133                return Err(error);
24134            }
24135        };
24136        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
24137            cache.tp_kv[il]
24138                .as_mut()
24139                .expect("distributed cache checked above"),
24140            transaction,
24141            1,
24142        ) {
24143            let _ = tp.runtime.rollback_tp_kv_transaction(
24144                cache.tp_kv[il]
24145                    .as_mut()
24146                    .expect("distributed cache checked above"),
24147                transaction,
24148            );
24149            let local = cache.kv[il].as_mut().expect("local cache checked above");
24150            local.len = base_len;
24151            e.set_i32_one(&mut local.len_d, base_len as i32)?;
24152            return Err(error);
24153        }
24154
24155        let committed = cache.tp_kv[il]
24156            .as_ref()
24157            .expect("distributed cache checked above")
24158            .committed_len();
24159        let local_len = cache.kv[il]
24160            .as_ref()
24161            .expect("local cache checked above")
24162            .len;
24163        if committed != local_len {
24164            return Err(format!(
24165                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
24166            )
24167            .into());
24168        }
24169        lap(&tp.runtime, e, &T_SHADOW, &mut lap_start)?;
24170        if timing {
24171            use std::sync::atomic::Ordering;
24172            let calls = T_PHASE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
24173            if calls.is_multiple_of(430) {
24174                let avg = |t: &std::sync::atomic::AtomicU64| {
24175                    t.load(Ordering::Relaxed) as f64 / calls as f64 / 1.0e3
24176                };
24177                eprintln!(
24178                    "[step-tp-attn-phase] calls={calls} avg_us pos={:.1} qkv={:.1} \
24179                     normrope={:.1} gate={:.1} append={:.1} attn={:.1} oproj={:.1} shadow={:.1}",
24180                    avg(&T_POS),
24181                    avg(&T_QKV),
24182                    avg(&T_NORMROPE),
24183                    avg(&T_GATE),
24184                    avg(&T_APPEND),
24185                    avg(&T_ATTN),
24186                    avg(&T_OPROJ),
24187                    avg(&T_SHADOW),
24188                );
24189            }
24190        }
24191        eprintln!(
24192            "[step-tp-attn] execute layer={} devices={:?} tokens=1 \
24193             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
24194             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
24195             attention_scope={} input_path={} kv_physical_rows={} \
24196             gate_tensor_parallel=false gate_shards=host-canonical o_tensor_parallel=true \
24197             local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
24198             bulk_p2p={} output=root-readback performance_claim=false",
24199            tp.layer,
24200            tp.devices,
24201            hydrated,
24202            if window.is_some() {
24203                "rank-local-swa-ring"
24204            } else {
24205                "rank-local-global"
24206            },
24207            input_path,
24208            cache.tp_kv[il]
24209                .as_ref()
24210                .expect("distributed cache checked above")
24211                .physical_capacity(),
24212            tp.runtime.transport_label(),
24213            tp.runtime.bulk_p2p(),
24214        );
24215        Ok(output)
24216    }
24217
24218    /// v2 rank-local decode attention (MEMRA_STEP_TP_DECODE_V2): the same kernels, operand
24219    /// values, and canonical reduction order as `step35_tp_decode_attn_resident_inner`,
24220    /// restructured onto a persistent per-runtime workspace with evented cross-stream ordering.
24221    /// Per layer per token this path performs exactly one cuMemAlloc (the returned e-context
24222    /// output row), no host round-trip, and no host stream synchronize — the phase timers
24223    /// measured v1 spending 81% of its 1550us/layer wall on those three classes.
24224    #[allow(clippy::too_many_arguments)]
24225    /// T-COLUMN verify precompute for layer `il`: weight-amortized QKV(+gate) over the T
24226    /// verify columns into the ws slabs (per-column rope/append/fa run later through the
24227    /// unmodified t=1 program via the col-select door). Ok(false) when the layer is not on
24228    /// the resident fused TP2 class (caller falls back to the per-row walk).
24229    pub(crate) fn step35_verify_qkv_precompute(
24230        &self,
24231        e: &Engine,
24232        il: usize,
24233        h_t: &CudaSlice<f32>,
24234        t: usize,
24235    ) -> Result<bool, Box<dyn std::error::Error>> {
24236        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24237            return Ok(false);
24238        };
24239        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24240            return Ok(false);
24241        };
24242        let Some(attention) = tp.attention.as_ref() else {
24243            return Ok(false);
24244        };
24245        if !tp.runtime.native_p2p() || !crate::tp::step_tp_qkv_fused_enabled()? {
24246            return Ok(false);
24247        }
24248        let geometry = self.step35_geom(il);
24249        let heads = geometry.n_head as usize;
24250        let ws_index = tp
24251            .runtime
24252            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
24253        let gate_shards = attention
24254            .gate_shards_bf16
24255            .as_deref()
24256            .map(crate::tp::StepTpGateShards::Bf16);
24257        tp.runtime.decode_v2_input_qkv_tcol(
24258            ws_index,
24259            e,
24260            h_t,
24261            t,
24262            &tp.q,
24263            &tp.k,
24264            &tp.v,
24265            gate_shards,
24266        )?;
24267        Ok(true)
24268    }
24269
24270    /// MEMRA_TCOL_OPROJ join for the verify walk: after every column of layer `il`
24271    /// stashed its `gated` rows, produce the [t, o_out] `mixed` slab on `e` via the
24272    /// weight-amortized b4_tcol + slab join. Callers only reach this after the stash
24273    /// flag confirmed the defer engaged for every column.
24274    pub(crate) fn step35_verify_oproj_tcol(
24275        &self,
24276        e: &Engine,
24277        il: usize,
24278        t: usize,
24279    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24280        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24281            return Err("tcol o_proj join expects full attention".into());
24282        };
24283        let tp = fa
24284            .step_tp_qkv
24285            .as_ref()
24286            .ok_or("tcol o_proj join lost its resident projections")?;
24287        let heads = self.step35_geom(il).n_head as usize;
24288        let ws_index = tp
24289            .runtime
24290            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
24291        tp.runtime.decode_v2_oproj_tcol(ws_index, e, &tp.o, t)
24292    }
24293
24294    /// MEMRA_SPEC_FA2 precheck: decide BEFORE arming the defer whether both verify
24295    /// columns of layer `il` will take the dcw arm AND the T=2 launch is bit-safe —
24296    /// stashing is unrecoverable (no per-column output exists), so every dynamic input
24297    /// to the engine-side dcw decision is evaluated here, plus the equal-partition
24298    /// guard fa_decode_dcw2's contract requires. Boundary rounds return false and the
24299    /// walk runs the ordinary per-column program.
24300    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
24301    pub(crate) fn step35_spec_fa2_precheck(
24302        &self,
24303        cache: &Cache,
24304        il: usize,
24305        pos0: usize,
24306    ) -> Result<bool, Box<dyn std::error::Error>> {
24307        // MEMRA_SPEC_FA2_DEBUG=1: print the first failing clause once per clause id —
24308        // a silently-vacuous door is indistinguishable from a slow one without this.
24309        fn nope(clause: &str, il: usize, pos0: usize) -> bool {
24310            static DBG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24311            static SEEN: std::sync::Mutex<Vec<&'static str>> = std::sync::Mutex::new(Vec::new());
24312            if *DBG.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1")) {
24313                let mut seen = SEEN.lock().unwrap();
24314                if !seen.contains(&clause) {
24315                    // leak: bounded by the clause-id set
24316                    seen.push(Box::leak(clause.to_string().into_boxed_str()));
24317                    eprintln!("[spec-fa2] precheck FAIL clause={clause} il={il} pos0={pos0}");
24318                }
24319            }
24320            false
24321        }
24322        // MEMRA_SPEC_FA2_LAYER=<il>: engage on ONE layer only (divergence bisection).
24323        static ONLY: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
24324        if let Some(only) =
24325            ONLY.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_LAYER").ok()?.parse().ok())
24326            && *only != il
24327        {
24328            return Ok(false);
24329        }
24330        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24331            return Ok(nope("mixer", il, pos0));
24332        };
24333        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24334            return Ok(nope("step_tp", il, pos0));
24335        };
24336        let Some(attention) = tp.attention.as_ref() else {
24337            return Ok(nope("attention", il, pos0));
24338        };
24339        if !tp.runtime.native_p2p()
24340            || crate::Engine::kv_fp8_on()
24341            || !crate::tp::step_tp_dcw_enabled()?
24342            || !crate::tp::step_tp_qkv_fused_enabled()?
24343            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
24344        {
24345            return Ok(nope("runtime-doors", il, pos0));
24346        }
24347        let geometry = self.step35_geom(il);
24348        let head_dim = geometry.head_dim_k as usize;
24349        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
24350            return Ok(nope("fa-class", il, pos0));
24351        }
24352        let Some(distributed) = cache.tp_kv[il].as_ref() else {
24353            return Ok(nope("tp-kv", il, pos0));
24354        };
24355        if distributed.staged_len() != pos0 {
24356            return Ok(nope("staged-len", il, pos0));
24357        }
24358        // Both appends must land without a ring rebase (rebase columns take the
24359        // host-row path, which cannot stash).
24360        let (_, would_rebase) = distributed.peek_append_ring(2)?;
24361        if would_rebase {
24362            return Ok(nope("rebase", il, pos0));
24363        }
24364        let window = geometry.window.map(|w| w as usize);
24365        // Capped SWA is REDUCTION-CLASS in the joined kernel (the two rows' windows
24366        // shift by one key, so one shared tile grid cannot reproduce both rows'
24367        // per-column FP grouping) — and drifted verify logits change accept decisions,
24368        // breaking the spec==target contract. Engage only when BOTH rows' views start
24369        // at 0 (global, or SWA still inside its window): bitwise per row under the
24370        // partition guard below. At agentic ctx this keeps the global layers — ~3/4 of
24371        // the per-key fa work — and leaves capped-SWA layers on the per-column program.
24372        if let Some(w) = window
24373            && pos0 + 2 > w
24374        {
24375            return Ok(nope("swa-capped", il, pos0));
24376        }
24377        // Row r's own per-column launch sees the POST-append view: T_r = pos0 + 1 + r
24378        // (kernel T_kv = len_dev - lstart; the host bucket matches it — the one-partition
24379        // law). Both dcw eligibility (t_kv_eff >= 96) and the vec floor key off T0.
24380        let (t0, t1) = (pos0 + 1, pos0 + 2);
24381        if t0 < 96 {
24382            return Ok(nope("dcw-floor", il, pos0));
24383        }
24384        if std::env::var("MEMRA_NO_FA_VEC").is_ok() || t0 < crate::fa_vec_min_tkv() {
24385            return Ok(nope("vec-floor", il, pos0));
24386        }
24387        // Equal-partition guard, on the KERNEL's derivation: split width (sp), effective
24388        // count (ns = ceil(T/sp)) and stride (per = ceil(T/ns)) must all match between
24389        // the two rows' own launches — the joined kernel derives one grid from T1 and
24390        // row0 inherits it, so any difference shifts row0's split boundaries and changes
24391        // the combine's merge rounding. Boundary rounds fall back per column.
24392        let ranks = tp.runtime.devices().len();
24393        let local_kv_heads = (geometry.n_head_kv as usize / ranks).max(1);
24394        let sp0 = crate::fa_split_keys_pub(t0, local_kv_heads);
24395        let sp1 = crate::fa_split_keys_pub(t1, local_kv_heads);
24396        if sp0 != sp1 {
24397            return Ok(nope("partition-sp", il, pos0));
24398        }
24399        let (ns0, ns1) = (t0.div_ceil(sp0), t1.div_ceil(sp1));
24400        if ns0 != ns1 {
24401            return Ok(nope("partition-ns", il, pos0));
24402        }
24403        if t0.div_ceil(ns0) != t1.div_ceil(ns1) {
24404            return Ok(nope("partition-per", il, pos0));
24405        }
24406        Ok(true)
24407    }
24408
24409    /// T-ROW fa precheck (the rows kernel supersedes the dcw2 pair-join): every dynamic
24410    /// input of the engine-side dcw decision must hold for EVERY row — stashing is
24411    /// unrecoverable — plus the rows-launcher guards (big-rig ladder, no env split
24412    /// overrides). No partition or capped-SWA clauses: each row derives its OWN geometry.
24413    pub(crate) fn step35_fa_rows_precheck(
24414        &self,
24415        cache: &Cache,
24416        il: usize,
24417        pos0: usize,
24418        t: usize,
24419    ) -> Result<bool, Box<dyn std::error::Error>> {
24420        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24421            return Ok(false);
24422        };
24423        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24424            return Ok(false);
24425        };
24426        let Some(attention) = tp.attention.as_ref() else {
24427            return Ok(false);
24428        };
24429        if !tp.runtime.native_p2p()
24430            || crate::Engine::kv_fp8_on()
24431            || !crate::tp::step_tp_dcw_enabled()?
24432            || !crate::tp::step_tp_qkv_fused_enabled()?
24433            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
24434        {
24435            return Ok(false);
24436        }
24437        let geometry = self.step35_geom(il);
24438        let head_dim = geometry.head_dim_k as usize;
24439        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
24440            return Ok(false);
24441        }
24442        if crate::fa_sm_count() < 128
24443            || std::env::var("MEMRA_FA_SPLIT").is_ok()
24444            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24445            || std::env::var("MEMRA_FA_SP16").is_ok()
24446            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
24447        {
24448            return Ok(false);
24449        }
24450        let Some(distributed) = cache.tp_kv[il].as_ref() else {
24451            return Ok(false);
24452        };
24453        if distributed.staged_len() != pos0 {
24454            return Ok(false);
24455        }
24456        if distributed.peek_append_ring(t).is_err() {
24457            return Ok(false);
24458        }
24459        if distributed.ring_base().is_none() && pos0 + t > distributed.physical_capacity() {
24460            return Ok(false);
24461        }
24462        // Row 0 sees the smallest view: its post-append effective t_kv must clear both
24463        // the dcw floor and the vec-class floor (later rows only grow).
24464        let window = geometry.window.map(|w| w as usize);
24465        let t0 = window.map(|w| (pos0 + 1).min(w)).unwrap_or(pos0 + 1);
24466        if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
24467            return Ok(false);
24468        }
24469        Ok(true)
24470    }
24471
24472    /// T-ROW fa join for the verify walk (same-session rows: shared ring/len with
24473    /// len_back = t-1-r). Tables stage once per (layer, rank, ring, t) and live on the
24474    /// owning rank.
24475    pub(crate) fn step35_verify_fa_rows_join(
24476        &self,
24477        e: &Engine,
24478        il: usize,
24479        cache: &Cache,
24480        pos0: usize,
24481        t: usize,
24482    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24483        use cudarc::driver::DevicePtr;
24484        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24485            return Err("fa rows join expects full attention".into());
24486        };
24487        let tp = fa
24488            .step_tp_qkv
24489            .as_ref()
24490            .ok_or("fa rows join lost its resident projections")?;
24491        let geometry = self.step35_geom(il);
24492        let heads = geometry.n_head as usize;
24493        let head_dim = geometry.head_dim_k as usize;
24494        let window = geometry.window.map(|w| w as usize);
24495        let distributed = cache.tp_kv[il]
24496            .as_ref()
24497            .ok_or("fa rows join lost its distributed KV cache")?;
24498        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
24499        // Host mirror of the kernel's big-rig ladder (launcher-guarded identical).
24500        let ladder = |t_kv: usize| -> usize {
24501            if t_kv <= 2048 {
24502                16
24503            } else if t_kv <= 16384 {
24504                64
24505            } else {
24506                128
24507            }
24508        };
24509        let mut max_ns = 1usize;
24510        for r in 0..t {
24511            let t_kv = window
24512                .map(|w| (pos0 + r + 1).min(w))
24513                .unwrap_or(pos0 + r + 1);
24514            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
24515        }
24516        // Rebuild the tiny raw-pointer table from the live distributed cache immediately
24517        // before launch. A process-lifetime map cannot prove allocation generation: CUDA may
24518        // recycle len/base independently of the large K/V rings, making a pointer-key cache
24519        // hit refer to another session (Hermes `11339f5cd3c132a3`).
24520        let ranks = tp.runtime.devices().len();
24521        let mut tables = Vec::with_capacity(ranks);
24522        for rank in 0..ranks {
24523            let engine = tp
24524                .runtime
24525                .rank_engine(rank)
24526                .ok_or("fa rows join lost a rank engine")?;
24527            let rank_cache = distributed
24528                .rank(rank)
24529                .ok_or("fa rows join lost a KV cache rank")?;
24530            let _main = engine.gpu.enter_main()?;
24531            let s = engine.stream();
24532            let (kp, _g0) = rank_cache.k().device_ptr(&s);
24533            let (vp, _g1) = rank_cache.v().device_ptr(&s);
24534            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
24535            let bp = match rank_cache.base_d() {
24536                Some(b) => {
24537                    let (p, _g) = b.device_ptr(&s);
24538                    p
24539                }
24540                None => 0u64,
24541            };
24542            let mut host = Vec::with_capacity(t * 6);
24543            for r in 0..t {
24544                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, (t - 1 - r) as u64]);
24545            }
24546            tables.push(engine.stream().clone_htod(&host)?);
24547        }
24548        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
24549        let ws_index = tp
24550            .runtime
24551            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
24552        tp.runtime.decode_v2_fa_rows_join(
24553            ws_index,
24554            e,
24555            &tp.o,
24556            &tabs,
24557            t,
24558            head_dim,
24559            window.unwrap_or(0),
24560            max_ns,
24561            geometry.attention_scale(),
24562            k_tok_bytes,
24563            v_tok_bytes,
24564        )
24565    }
24566
24567    /// Multi-session t-row fa precheck (the batched serving walk): the static doors of
24568    /// the rows kernel plus per-SESSION dynamic checks — every row's own cache must be
24569    /// hydrated, in sync, rebase-free and above both floors.
24570    pub(crate) fn step35_batch_fa_rows_precheck(
24571        &self,
24572        caches: &[&mut Cache],
24573        row_to_cache: impl Fn(usize) -> usize,
24574        positions: &[i32],
24575        il: usize,
24576    ) -> Result<bool, Box<dyn std::error::Error>> {
24577        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24578            return Ok(false);
24579        };
24580        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24581            return Ok(false);
24582        };
24583        let Some(attention) = tp.attention.as_ref() else {
24584            return Ok(false);
24585        };
24586        if !tp.runtime.native_p2p()
24587            || crate::Engine::kv_fp8_on()
24588            || !crate::tp::step_tp_dcw_enabled()?
24589            || !crate::tp::step_tp_qkv_fused_enabled()?
24590            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
24591        {
24592            return Ok(false);
24593        }
24594        let geometry = self.step35_geom(il);
24595        let head_dim = geometry.head_dim_k as usize;
24596        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
24597            return Ok(false);
24598        }
24599        if crate::fa_sm_count() < 128
24600            || std::env::var("MEMRA_FA_SPLIT").is_ok()
24601            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24602            || std::env::var("MEMRA_FA_SP16").is_ok()
24603            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
24604        {
24605            return Ok(false);
24606        }
24607        let window = geometry.window.map(|w| w as usize);
24608        for (r, &pos) in positions.iter().enumerate() {
24609            let cache = &caches[row_to_cache(r)];
24610            let Some(distributed) = cache.tp_kv[il].as_ref() else {
24611                return Ok(false);
24612            };
24613            if distributed.staged_len() != pos as usize {
24614                return Ok(false);
24615            }
24616            if distributed.peek_append_ring(1)?.1 {
24617                return Ok(false);
24618            }
24619            let t0 = window
24620                .map(|w| (pos as usize + 1).min(w))
24621                .unwrap_or(pos as usize + 1);
24622            if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
24623                return Ok(false);
24624            }
24625        }
24626        Ok(true)
24627    }
24628
24629    /// FULL t-row attention pass for the VERIFY walk (same-session rows): rope/append +
24630    /// fa + combine + o_proj join in 3 launches/rank/layer. Row r appends at slot
24631    /// len-base+r and one last block advances len by t; the fa rows read len_back =
24632    /// t-1-r. Returns None when the fused-rope class does not hold (the walk keeps the
24633    /// per-column stash flow). Caller has passed `step35_fa_rows_precheck`.
24634    pub(crate) fn step35_verify_rope_fa_pass(
24635        &self,
24636        e: &Engine,
24637        il: usize,
24638        cache: &Cache,
24639        pos0: usize,
24640        t: usize,
24641        stage_pos: bool,
24642    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
24643        use cudarc::driver::DevicePtr;
24644        if !crate::tp::fuse_rope_append_on() {
24645            return Ok(None);
24646        }
24647        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24648            return Ok(None);
24649        };
24650        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24651            return Ok(None);
24652        };
24653        let Some(attention) = tp.attention.as_ref() else {
24654            return Ok(None);
24655        };
24656        let geometry = self.step35_geom(il);
24657        let head_dim = geometry.head_dim_k as usize;
24658        if head_dim != 128 {
24659            return Ok(None);
24660        }
24661        let heads = geometry.n_head as usize;
24662        let window = geometry.window.map(|w| w as usize);
24663        let ranks = tp.runtime.devices().len();
24664        let Some(distributed) = cache.tp_kv[il].as_ref() else {
24665            return Ok(None);
24666        };
24667        if distributed.kv_dim_k() != distributed.kv_dim_v() {
24668            return Ok(None);
24669        }
24670        {
24671            let rank0 = distributed.rank(0).ok_or("verify rope pass lost rank 0")?;
24672            if rank0.base_d().is_none()
24673                && distributed.staged_len() > distributed.physical_capacity()
24674            {
24675                return Ok(None);
24676            }
24677        }
24678        let mut rope_freqs = Vec::with_capacity(ranks);
24679        for rank in 0..ranks {
24680            let engine = tp
24681                .runtime
24682                .rank_engine(rank)
24683                .ok_or("verify rope pass lost a rank engine")?;
24684            rope_freqs.push(if geometry.rope_factors {
24685                match self
24686                    .step35_aux
24687                    .as_ref()
24688                    .and_then(|aux| aux.rope_freqs(engine))
24689                {
24690                    Some(f) => Some(f),
24691                    None => return Ok(None),
24692                }
24693            } else {
24694                None
24695            });
24696        }
24697        let ladder = |t_kv: usize| -> usize {
24698            if t_kv <= 2048 {
24699                16
24700            } else if t_kv <= 16384 {
24701                64
24702            } else {
24703                128
24704            }
24705        };
24706        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
24707        let mut max_ns = 1usize;
24708        let mut positions = Vec::with_capacity(t);
24709        for r in 0..t {
24710            positions.push((pos0 + r) as i32);
24711            let t_kv = window
24712                .map(|w| (pos0 + r + 1).min(w))
24713                .unwrap_or(pos0 + r + 1);
24714            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
24715        }
24716        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
24717        let mut tab_keys = vec![0u64; ranks];
24718        for rank in 0..ranks {
24719            let engine = tp
24720                .runtime
24721                .rank_engine(rank)
24722                .ok_or("verify rope pass lost a rank engine")?;
24723            let rank_cache = distributed
24724                .rank(rank)
24725                .ok_or("verify rope pass lost a KV cache rank")?;
24726            let _main = engine.gpu.enter_main()?;
24727            let s = engine.stream();
24728            let (kp, _g0) = rank_cache.k().device_ptr(&s);
24729            let (vp, _g1) = rank_cache.v().device_ptr(&s);
24730            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
24731            let bp = match rank_cache.base_d() {
24732                Some(b) => {
24733                    let (p, _g) = b.device_ptr(&s);
24734                    p
24735                }
24736                None => 0u64,
24737            };
24738            tab_keys[rank] = kp
24739                .rotate_left(17)
24740                .wrapping_add(bp)
24741                .wrapping_add((il as u64) << 32)
24742                .wrapping_add(t as u64)
24743                .wrapping_add(1 << 63);
24744            for _r in 0..t {
24745                session_parts[rank].push([kp, vp, lp, bp]);
24746            }
24747        }
24748        let ws_index = tp
24749            .runtime
24750            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
24751        tp.runtime
24752            .decode_v2_rope_fa_rows(
24753                ws_index,
24754                e,
24755                &tp.o,
24756                &session_parts,
24757                &tab_keys,
24758                &positions,
24759                stage_pos,
24760                true,
24761                &attention.q_norm,
24762                &attention.k_norm,
24763                &rope_freqs,
24764                t,
24765                head_dim,
24766                geometry.n_rot as usize,
24767                window.unwrap_or(0),
24768                max_ns,
24769                geometry.attention_scale(),
24770                k_tok_bytes,
24771                v_tok_bytes,
24772                self.cfg.rms_eps,
24773                geometry.rope_base,
24774            )
24775            .map(Some)
24776    }
24777
24778    /// FULL t-row attention pass for the batched walk (rope/append + fa + combine +
24779    /// o_proj join, 3 launches/rank/layer): returns None when the fused-rope class does
24780    /// not hold — the caller falls back to the per-row stash flow. The caller has
24781    /// already passed `step35_batch_fa_rows_precheck`.
24782    #[allow(clippy::too_many_arguments)]
24783    pub(crate) fn step35_batch_rope_fa_pass(
24784        &self,
24785        e: &Engine,
24786        il: usize,
24787        caches: &[&mut Cache],
24788        row_to_cache: impl Fn(usize) -> usize,
24789        positions: &[i32],
24790        t: usize,
24791        stage_pos: bool,
24792    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
24793        use cudarc::driver::DevicePtr;
24794        if !crate::tp::fuse_rope_append_on() {
24795            return Ok(None);
24796        }
24797        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24798            return Ok(None);
24799        };
24800        let Some(tp) = fa.step_tp_qkv.as_ref() else {
24801            return Ok(None);
24802        };
24803        let Some(attention) = tp.attention.as_ref() else {
24804            return Ok(None);
24805        };
24806        let geometry = self.step35_geom(il);
24807        let head_dim = geometry.head_dim_k as usize;
24808        if head_dim != 128 {
24809            return Ok(None);
24810        }
24811        let heads = geometry.n_head as usize;
24812        let window = geometry.window.map(|w| w as usize);
24813        let ranks = tp.runtime.devices().len();
24814        // The rows kernels never arm base_d; refuse once a ring could have rebased
24815        // without an armed base (the table would read base=0 after a real rebase).
24816        for r in 0..t {
24817            let cache = &caches[row_to_cache(r)];
24818            let Some(distributed) = cache.tp_kv[il].as_ref() else {
24819                return Ok(None);
24820            };
24821            if distributed.kv_dim_k() != distributed.kv_dim_v() {
24822                return Ok(None);
24823            }
24824            let rank0 = distributed.rank(0).ok_or("rope fa pass lost rank 0")?;
24825            if rank0.base_d().is_none()
24826                && distributed.staged_len() + t > distributed.physical_capacity()
24827            {
24828                return Ok(None);
24829            }
24830        }
24831        let mut rope_freqs = Vec::with_capacity(ranks);
24832        for rank in 0..ranks {
24833            let engine = tp
24834                .runtime
24835                .rank_engine(rank)
24836                .ok_or("rope fa pass lost a rank engine")?;
24837            rope_freqs.push(if geometry.rope_factors {
24838                match self
24839                    .step35_aux
24840                    .as_ref()
24841                    .and_then(|aux| aux.rope_freqs(engine))
24842                {
24843                    Some(f) => Some(f),
24844                    None => return Ok(None),
24845                }
24846            } else {
24847                None
24848            });
24849        }
24850        let ladder = |t_kv: usize| -> usize {
24851            if t_kv <= 2048 {
24852                16
24853            } else if t_kv <= 16384 {
24854                64
24855            } else {
24856                128
24857            }
24858        };
24859        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
24860        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
24861        let mut tab_keys = vec![0u64; ranks];
24862        for (r, &pos) in positions.iter().enumerate().take(t) {
24863            let cache = &caches[row_to_cache(r)];
24864            let distributed = cache.tp_kv[il]
24865                .as_ref()
24866                .ok_or("rope fa pass lost a distributed KV cache")?;
24867            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
24868            let t_kv = window
24869                .map(|w| (pos as usize + 1).min(w))
24870                .unwrap_or(pos as usize + 1);
24871            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
24872            for rank in 0..ranks {
24873                let engine = tp
24874                    .runtime
24875                    .rank_engine(rank)
24876                    .ok_or("rope fa pass lost a rank engine")?;
24877                let rank_cache = distributed
24878                    .rank(rank)
24879                    .ok_or("rope fa pass lost a KV cache rank")?;
24880                let _main = engine.gpu.enter_main()?;
24881                let s = engine.stream();
24882                let (kp, _g0) = rank_cache.k().device_ptr(&s);
24883                let (vp, _g1) = rank_cache.v().device_ptr(&s);
24884                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
24885                let bp = match rank_cache.base_d() {
24886                    Some(b) => {
24887                        let (p, _g) = b.device_ptr(&s);
24888                        p
24889                    }
24890                    None => 0u64,
24891                };
24892                tab_keys[rank] = tab_keys[rank]
24893                    .rotate_left(9)
24894                    .wrapping_add(kp)
24895                    .wrapping_add(bp)
24896                    .wrapping_add(il as u64);
24897                session_parts[rank].push([kp, vp, lp, bp]);
24898            }
24899        }
24900        let ws_index = tp
24901            .runtime
24902            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
24903        tp.runtime
24904            .decode_v2_rope_fa_rows(
24905                ws_index,
24906                e,
24907                &tp.o,
24908                &session_parts,
24909                &tab_keys,
24910                positions,
24911                stage_pos,
24912                false,
24913                &attention.q_norm,
24914                &attention.k_norm,
24915                &rope_freqs,
24916                t,
24917                head_dim,
24918                geometry.n_rot as usize,
24919                window.unwrap_or(0),
24920                max_ns,
24921                geometry.attention_scale(),
24922                k_tok_bytes,
24923                v_tok_bytes,
24924                self.cfg.rms_eps,
24925                geometry.rope_base,
24926            )
24927            .map(Some)
24928    }
24929
24930    /// Multi-session t-row fa join (batched serving): per-row table entries point at
24931    /// each row's OWN session rings/counters (len_back = 0 — every session appended
24932    /// exactly its one row). Tables stage once per (layer, rank, session-set, t).
24933    #[allow(clippy::too_many_arguments)]
24934    pub(crate) fn step35_batch_fa_rows_join(
24935        &self,
24936        e: &Engine,
24937        il: usize,
24938        caches: &[&mut Cache],
24939        row_to_cache: impl Fn(usize) -> usize,
24940        positions: &[i32],
24941        t: usize,
24942    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24943        use cudarc::driver::DevicePtr;
24944        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
24945            return Err("batch fa rows join expects full attention".into());
24946        };
24947        let tp = fa
24948            .step_tp_qkv
24949            .as_ref()
24950            .ok_or("batch fa rows join lost its resident projections")?;
24951        let geometry = self.step35_geom(il);
24952        let heads = geometry.n_head as usize;
24953        let head_dim = geometry.head_dim_k as usize;
24954        let window = geometry.window.map(|w| w as usize);
24955        let ladder = |t_kv: usize| -> usize {
24956            if t_kv <= 2048 {
24957                16
24958            } else if t_kv <= 16384 {
24959                64
24960            } else {
24961                128
24962            }
24963        };
24964        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
24965        for (r, &pos) in positions.iter().enumerate() {
24966            let cache = &caches[row_to_cache(r)];
24967            let distributed = cache.tp_kv[il]
24968                .as_ref()
24969                .ok_or("batch fa rows join lost a distributed KV cache")?;
24970            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
24971            let t_kv = window
24972                .map(|w| (pos as usize + 1).min(w))
24973                .unwrap_or(pos as usize + 1);
24974            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
24975        }
24976        // Multi-session tables also rebuild from every live K/V/len/base tuple. Keeping a
24977        // process-lifetime raw-pointer cache here omitted V and len identity and had no
24978        // allocation generation, so allocator reuse could bind one request to another.
24979        let ranks = tp.runtime.devices().len();
24980        let mut tables = Vec::with_capacity(ranks);
24981        for rank in 0..ranks {
24982            let engine = tp
24983                .runtime
24984                .rank_engine(rank)
24985                .ok_or("batch fa rows join lost a rank engine")?;
24986            let _main = engine.gpu.enter_main()?;
24987            let s = engine.stream();
24988            let mut host = Vec::with_capacity(t * 6);
24989            for r in 0..t {
24990                let cache = &caches[row_to_cache(r)];
24991                let distributed = cache.tp_kv[il]
24992                    .as_ref()
24993                    .ok_or("batch fa rows join lost a distributed KV cache")?;
24994                let rank_cache = distributed
24995                    .rank(rank)
24996                    .ok_or("batch fa rows join lost a KV cache rank")?;
24997                let (kp, _g0) = rank_cache.k().device_ptr(&s);
24998                let (vp, _g1) = rank_cache.v().device_ptr(&s);
24999                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
25000                let bp = match rank_cache.base_d() {
25001                    Some(b) => {
25002                        let (p, _g) = b.device_ptr(&s);
25003                        p
25004                    }
25005                    None => 0u64,
25006                };
25007                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, 0u64]);
25008            }
25009            tables.push(engine.stream().clone_htod(&host)?);
25010        }
25011        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
25012        let ws_index = tp
25013            .runtime
25014            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
25015        tp.runtime.decode_v2_fa_rows_join(
25016            ws_index,
25017            e,
25018            &tp.o,
25019            &tabs,
25020            t,
25021            head_dim,
25022            window.unwrap_or(0),
25023            max_ns,
25024            geometry.attention_scale(),
25025            k_tok_bytes,
25026            v_tok_bytes,
25027        )
25028    }
25029
25030    /// MEMRA_SPEC_FA2 join for the verify walk: both columns stashed; one shared-KV T=2
25031    /// fa per rank + the weight-amortized o_proj join produce the [2, o_out] `mixed`
25032    /// slab on `e`.
25033    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
25034    pub(crate) fn step35_verify_spec_fa2_join(
25035        &self,
25036        e: &Engine,
25037        il: usize,
25038        cache: &Cache,
25039        pos0: usize,
25040    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25041        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
25042            return Err("spec fa2 join expects full attention".into());
25043        };
25044        let tp = fa
25045            .step_tp_qkv
25046            .as_ref()
25047            .ok_or("spec fa2 join lost its resident projections")?;
25048        let geometry = self.step35_geom(il);
25049        let heads = geometry.n_head as usize;
25050        let head_dim = geometry.head_dim_k as usize;
25051        let window = geometry.window.map(|w| w as usize);
25052        // POST-append view of the second row (kernel T1 = len - lstart with len =
25053        // pos0 + 2): sp/ns derive from it, and the precheck proved row0 shares them.
25054        let bucket = window.map(|w| (pos0 + 2).min(w)).unwrap_or(pos0 + 2);
25055        let distributed = cache.tp_kv[il]
25056            .as_ref()
25057            .ok_or("spec fa2 join lost its distributed KV cache")?;
25058        let ws_index = tp
25059            .runtime
25060            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
25061        tp.runtime.decode_v2_spec_fa2_join(
25062            ws_index,
25063            e,
25064            &tp.o,
25065            distributed,
25066            head_dim,
25067            window.unwrap_or(0),
25068            bucket,
25069            geometry.attention_scale(),
25070        )
25071    }
25072
25073    pub(crate) fn step35_tp_decode_attn_resident_v2(
25074        &self,
25075        e: &Engine,
25076        fa: &FullAttnLayer,
25077        il: usize,
25078        h: &CudaSlice<f32>,
25079        pos_d: &CudaSlice<i32>,
25080        cache: &mut Cache,
25081    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25082        let tp = fa
25083            .step_tp_qkv
25084            .as_ref()
25085            .ok_or("Step TP decode lost its resident projections")?;
25086        let attention = tp
25087            .attention
25088            .as_ref()
25089            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
25090        if !tp.runtime.native_p2p() {
25091            return Err("rank-local Step attention requires native P2P".into());
25092        }
25093        if crate::Engine::kv_fp8_on() {
25094            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
25095        }
25096
25097        let geometry = self.cfg.full_attention_geometry_at(il as u32);
25098        let window = geometry.window.map(|window| window as usize);
25099        let ranks = tp.runtime.devices().len();
25100        let head_dim = geometry.head_dim_k as usize;
25101        let heads = geometry.n_head as usize;
25102        let kv_heads = geometry.n_head_kv as usize;
25103        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
25104            return Err(format!(
25105                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
25106            )
25107            .into());
25108        }
25109        let local_heads = heads / ranks;
25110        let local_kv_heads = kv_heads / ranks;
25111        let max_ctx = cache.max_ctx;
25112
25113        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
25114
25115        let base_len = cache.kv[il]
25116            .as_ref()
25117            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
25118            .len;
25119        {
25120            let distributed = cache.tp_kv[il]
25121                .as_ref()
25122                .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
25123            if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
25124                return Err(format!(
25125                    "Step TP layer {il} cache lengths diverged before decode: \
25126                     local={base_len} distributed={}/{}",
25127                    distributed.committed_len(),
25128                    distributed.staged_len()
25129                )
25130                .into());
25131            }
25132        }
25133        if pos_d.len() != 1 {
25134            return Err(format!(
25135                "rank-local Step decode requires one position, got {}",
25136                pos_d.len()
25137            )
25138            .into());
25139        }
25140
25141        let decode_input = attention
25142            .decode_input
25143            .as_ref()
25144            .ok_or("Step TP decode v2 requires the replicated decode input")?;
25145        let mut decode_input = decode_input
25146            .lock()
25147            .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
25148
25149        let has_gate = fa.attn_gate.is_some();
25150        // Gate: with per-rank shards loaded (fused door), the fused QKV+gate kernel computes
25151        // it rank-locally and the model-engine matmul (and its staging copies) disappears.
25152        // Otherwise it queues on e's stream BEFORE decode_v2_input_qkv records the entry
25153        // event, so the rank-stream reads of the staged gate are ordered without a host sync.
25154        let use_gate_shards = has_gate
25155            && (attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some())
25156            && crate::tp::step_tp_qkv_fused_enabled()?;
25157        let gate_raw = if !has_gate || use_gate_shards {
25158            None
25159        } else {
25160            let gate_weight = fa
25161                .attn_gate
25162                .as_ref()
25163                .ok_or("step35 layer is missing attn_gate.weight")?;
25164            let gate_raw = e.matmul(gate_weight, h, 1)?;
25165            if gate_raw.len() != heads {
25166                return Err(format!(
25167                    "Step TP layer {il} gate output {} != {heads}",
25168                    gate_raw.len()
25169                )
25170                .into());
25171            }
25172            Some(gate_raw)
25173        };
25174
25175        let ws_index = tp
25176            .runtime
25177            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
25178        let mut ws_guard = tp
25179            .runtime
25180            .decode_v2_workspace()
25181            .lock()
25182            .map_err(|_| "Step TP decode v2 workspace lock is poisoned")?;
25183        let ws = ws_guard
25184            .get_mut(ws_index)
25185            .ok_or("Step TP decode v2 workspace missing after ensure")?;
25186
25187        let mut rope_freqs = Vec::with_capacity(ranks);
25188        for rank in 0..ranks {
25189            let engine = tp
25190                .runtime
25191                .rank_engine(rank)
25192                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
25193            rope_freqs.push(if geometry.rope_factors {
25194                self.step35_aux
25195                    .as_ref()
25196                    .and_then(|aux| aux.rope_freqs(engine))
25197            } else {
25198                None
25199            });
25200        }
25201        // DCW arm (MEMRA_STEP_TP_DCW=1): device-counter append + counter-derived fa — the
25202        // exact captured-child content, run eagerly. bucket = effective t_kv keeps the fa
25203        // split geometry identical to the kvmod path (one-partition law) -> bit-identical.
25204        // Rebase tokens and sub-vec-floor contexts take the host-row path below.
25205        // (Eligibility computed BEFORE input_qkv so FUSION #1 can defer the norm+rope into
25206        // the fused rope+append+inc launch on dcw tokens.)
25207        let staged_next = base_len + 1;
25208        let t_kv_eff = window
25209            .map(|window| staged_next.min(window))
25210            .unwrap_or(staged_next);
25211        let dcw = crate::tp::step_tp_dcw_enabled()?
25212            && (use_gate_shards || (!has_gate && crate::tp::step_tp_qkv_fused_enabled()?))
25213            && t_kv_eff >= 96
25214            && {
25215                let (write_row, would_rebase) = cache.tp_kv[il]
25216                    .as_ref()
25217                    .expect("distributed cache checked above")
25218                    .peek_append_ring(1)?;
25219                if !would_rebase {
25220                    // Arm the base mirrors on first use: base = logical staged - physical row.
25221                    let base = (base_len - write_row) as i32;
25222                    let distributed = cache.tp_kv[il]
25223                        .as_mut()
25224                        .expect("distributed cache checked above");
25225                    for rank in 0..ranks {
25226                        let engine = tp.runtime.rank_engine(rank).ok_or_else(|| {
25227                            format!("Step TP layer {il} has no engine for rank {rank}")
25228                        })?;
25229                        let _main = engine.gpu.enter_main()?;
25230                        let rank_cache = distributed.rank_mut(rank).ok_or_else(|| {
25231                            format!("Step TP layer {il} has no KV cache rank {rank}")
25232                        })?;
25233                        if rank_cache.base_d().is_none() {
25234                            rank_cache.arm_base_d(engine.htod_i32(&[base])?);
25235                        }
25236                    }
25237                }
25238                !would_rebase
25239            };
25240        let fuse_rope = dcw
25241            && crate::tp::fuse_rope_append_on()
25242            && head_dim == 128
25243            && cache.tp_kv[il]
25244                .as_ref()
25245                .map(|d| d.kv_dim_k() == d.kv_dim_v() && d.kv_dim_k() == local_kv_heads * head_dim)
25246                .unwrap_or(false);
25247
25248        let tcol_col = crate::tp::take_verify_tcol();
25249        // MEMRA_SPEC_FA2 defer: the verify walk armed this column for the shared-KV T=2
25250        // attention. On dcw tokens the per-rank pass still norms/ropes/APPENDS (cache
25251        // state must advance per column) but skips the fa+gate launch; post-rope q and
25252        // gate rows are stashed instead, and ONE fa_decode_dcw2 per rank joins both
25253        // columns after the second append. Non-dcw tokens ignore the defer (the fa runs
25254        // normally and the walk consumes the real output — stash flag stays unset).
25255        let fa2_col = crate::tp::take_spec_fa2_defer();
25256        tp.runtime.decode_v2_input_qkv(
25257            ws,
25258            e,
25259            h,
25260            pos_d,
25261            gate_raw.as_ref(),
25262            if !use_gate_shards {
25263                None
25264            } else if let Some(shards) = attention.gate_shards.as_deref() {
25265                Some(crate::tp::StepTpGateShards::F32(shards))
25266            } else {
25267                attention
25268                    .gate_shards_bf16
25269                    .as_deref()
25270                    .map(crate::tp::StepTpGateShards::Bf16)
25271            },
25272            &mut decode_input,
25273            &tp.q,
25274            &tp.k,
25275            &tp.v,
25276            &attention.q_norm,
25277            &attention.k_norm,
25278            head_dim,
25279            geometry.n_rot as usize,
25280            geometry.rope_base,
25281            &rope_freqs,
25282            self.cfg.rms_eps,
25283            has_gate,
25284            fuse_rope,
25285            tcol_col,
25286        )?;
25287
25288        let transaction = cache.tp_kv[il]
25289            .as_mut()
25290            .expect("distributed cache checked above")
25291            .begin_transaction()?;
25292        let append_result = tp.runtime.append_tp_kv_transaction_inner(
25293            cache.tp_kv[il]
25294                .as_mut()
25295                .expect("distributed cache checked above"),
25296            transaction,
25297            &ws.k,
25298            &ws.v_raw,
25299            1,
25300            dcw,
25301        );
25302        if let Err(error) = append_result {
25303            let _ = tp.runtime.rollback_tp_kv_transaction(
25304                cache.tp_kv[il]
25305                    .as_mut()
25306                    .expect("distributed cache checked above"),
25307                transaction,
25308            );
25309            return Err(error);
25310        }
25311
25312        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25313            // Copy the view scalars out and DROP the shared borrow — the merged dcw arm
25314            // reborrows the cache mutably per rank.
25315            let (staged_len, physical, k_tok_bytes_c, v_tok_bytes_c, capacity) = {
25316                let distributed = cache.tp_kv[il]
25317                    .as_ref()
25318                    .expect("distributed cache checked above");
25319                let staged_len = distributed.staged_len();
25320                let view_start = window
25321                    .map(|window| staged_len.saturating_sub(window))
25322                    .unwrap_or(0);
25323                (
25324                    staged_len,
25325                    distributed.physical_range(view_start, staged_len)?,
25326                    distributed.k_tok_bytes(),
25327                    distributed.v_tok_bytes(),
25328                    distributed.physical_capacity(),
25329                )
25330            };
25331            let view_start = window
25332                .map(|window| staged_len.saturating_sub(window))
25333                .unwrap_or(0);
25334            let t_kv = staged_len - view_start;
25335            for rank in 0..ranks {
25336                let engine = tp
25337                    .runtime
25338                    .rank_engine(rank)
25339                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
25340                let _main = engine.gpu.enter_main()?;
25341                if dcw {
25342                    // MERGED per-rank pass (the capture unit): append + inc + fa + gate on ONE
25343                    // stream visit. distributed is borrowed shared here; the planes need mut —
25344                    // reborrow through the cache Option (the closure holds cache mutably).
25345                    {
25346                        let distributed_mut = cache.tp_kv[il]
25347                            .as_mut()
25348                            .expect("distributed cache checked above");
25349                        let (kv_dim_k, kv_dim_v) =
25350                            (distributed_mut.kv_dim_k(), distributed_mut.kv_dim_v());
25351                        let (k_tok_bytes, v_tok_bytes) =
25352                            (distributed_mut.k_tok_bytes(), distributed_mut.v_tok_bytes());
25353                        let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
25354                            format!("Step TP layer {il} has no KV cache rank {rank}")
25355                        })?;
25356                        let (k_plane, v_plane, len_d, base_d) =
25357                            rank_cache.planes_and_counters_mut();
25358                        if fuse_rope {
25359                            // FUSION #1: norms + rope (deferred out of input_qkv) + append
25360                            // + last-block len inc in ONE launch. Bit-identical bodies.
25361                            let same_dev = engine.ctx().ordinal() == e.ctx().ordinal();
25362                            let crate::tp::StepTpDecodeV2Ws {
25363                                q_raw,
25364                                k_raw,
25365                                v_raw,
25366                                q,
25367                                k,
25368                                pos,
25369                                pos_stage,
25370                                fuse_ctr,
25371                                ..
25372                            } = &mut *ws;
25373                            // Same-device rank: the staged-copy elision leaves pos[rank]
25374                            // stale — read the e-context pos stage directly (mirrors the
25375                            // rope elision in input_qkv_rank).
25376                            let pos_ref: &CudaSlice<i32> = if same_dev {
25377                                pos_stage
25378                                    .as_ref()
25379                                    .ok_or("step TP decode v2 pos stage not armed")?
25380                            } else {
25381                                &pos[rank]
25382                            };
25383                            engine.qk_norm_rope_append_inc_dcw(
25384                                &q_raw[rank],
25385                                &k_raw[rank],
25386                                &v_raw[rank],
25387                                &attention.q_norm[rank],
25388                                &attention.k_norm[rank],
25389                                &mut q[rank],
25390                                &mut k[rank],
25391                                pos_ref,
25392                                k_plane,
25393                                v_plane,
25394                                len_d,
25395                                base_d,
25396                                &mut fuse_ctr[rank],
25397                                kv_dim_k,
25398                                kv_dim_v,
25399                                k_tok_bytes,
25400                                v_tok_bytes,
25401                                head_dim,
25402                                geometry.n_rot as usize,
25403                                local_heads,
25404                                local_kv_heads,
25405                                self.cfg.rms_eps,
25406                                geometry.rope_base,
25407                                1.0,
25408                                rope_freqs[rank],
25409                            )?;
25410                        } else {
25411                            engine.append_kv_quantized_dcw(
25412                                &ws.k[rank],
25413                                &ws.v_raw[rank],
25414                                k_plane,
25415                                v_plane,
25416                                len_d,
25417                                base_d,
25418                                kv_dim_k,
25419                                kv_dim_v,
25420                                k_tok_bytes,
25421                                v_tok_bytes,
25422                            )?;
25423                        }
25424                        if !fuse_rope {
25425                            let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
25426                                format!("Step TP layer {il} has no KV cache rank {rank}")
25427                            })?;
25428                            engine.inc_i32(rank_cache.len_d_mut())?;
25429                        }
25430                    }
25431                    let distributed = cache.tp_kv[il]
25432                        .as_ref()
25433                        .expect("distributed cache checked above");
25434                    let rank_cache = distributed
25435                        .rank(rank)
25436                        .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
25437                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * k_tok_bytes_c);
25438                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * v_tok_bytes_c);
25439                    if fa2_col.is_some() {
25440                        // SPEC_FA2 defer: append landed above; the fa for this column
25441                        // runs in the T=2 joined launch after the pair's second append.
25442                        continue;
25443                    }
25444                    {
25445                        // FUSION #2d: combine + head gate in the dcw tail — `gated` receives
25446                        // the gated output directly (bit-identical; one launch saved).
25447                        let crate::tp::StepTpDecodeV2Ws { q, gate, gated, .. } = &mut *ws;
25448                        engine.fa_decode_dcw(
25449                            &q[rank],
25450                            &k_ring,
25451                            &v_ring,
25452                            &mut gated[rank],
25453                            head_dim,
25454                            local_heads,
25455                            local_kv_heads,
25456                            rank_cache.len_d(),
25457                            rank_cache.base_d(),
25458                            window.unwrap_or(0),
25459                            t_kv,
25460                            geometry.attention_scale(),
25461                            k_tok_bytes_c,
25462                            v_tok_bytes_c,
25463                            has_gate.then_some(&gate[rank]),
25464                        )?;
25465                    }
25466                    continue;
25467                }
25468                let distributed = cache.tp_kv[il]
25469                    .as_ref()
25470                    .expect("distributed cache checked above");
25471                let rank_cache = distributed
25472                    .rank(rank)
25473                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
25474                let k_view = engine.view_u8_range(
25475                    rank_cache.k(),
25476                    physical.start * k_tok_bytes_c,
25477                    physical.end * k_tok_bytes_c,
25478                );
25479                let v_view = engine.view_u8_range(
25480                    rank_cache.v(),
25481                    physical.start * v_tok_bytes_c,
25482                    physical.end * v_tok_bytes_c,
25483                );
25484                if has_gate {
25485                    engine.fa_decode_kvmod(
25486                        &ws.q[rank],
25487                        &k_view,
25488                        &v_view,
25489                        &mut ws.attn_out[rank],
25490                        head_dim,
25491                        local_heads,
25492                        local_kv_heads,
25493                        t_kv,
25494                        geometry.attention_scale(),
25495                        k_tok_bytes_c,
25496                        v_tok_bytes_c,
25497                        false,
25498                    )?;
25499                    engine.attn_head_gate(
25500                        &ws.attn_out[rank],
25501                        &ws.gate[rank],
25502                        &mut ws.gated[rank],
25503                        None,
25504                        head_dim,
25505                        local_heads,
25506                        1,
25507                    )?;
25508                } else {
25509                    engine.fa_decode_kvmod(
25510                        &ws.q[rank],
25511                        &k_view,
25512                        &v_view,
25513                        &mut ws.gated[rank],
25514                        head_dim,
25515                        local_heads,
25516                        local_kv_heads,
25517                        t_kv,
25518                        geometry.attention_scale(),
25519                        k_tok_bytes_c,
25520                        v_tok_bytes_c,
25521                        false,
25522                    )?;
25523                }
25524            }
25525
25526            // MEMRA_TCOL_OPROJ defer: the verify driver armed a column — stash this
25527            // column's `gated` rows and skip the per-column finish choreography entirely
25528            // (the batched b4_tcol + join runs after every column). The returned buffer
25529            // is UNWRITTEN in that mode (oproj-tail precedent); the driver reads the
25530            // stashed flag, never this buffer. Ineligible configs fall back to the
25531            // normal finish and the driver consumes the real `mixed` per column.
25532            let output = if let Some(col) = fa2_col.filter(|_| dcw) {
25533                // SPEC_FA2 stash: q + gate rows to the fa2 slabs; fa, o_proj and the
25534                // finish all run in the joined pass. Returned buffer is UNWRITTEN
25535                // (oproj-defer precedent — the walk reads the stash flag, never this).
25536                tp.runtime.decode_v2_stash_fa2(ws, e, col)?;
25537                crate::tp::set_spec_fa2_stashed();
25538                e.uninit(ws.o_out)?
25539            } else if let Some(col) = crate::tp::take_tcol_oproj_defer() {
25540                if tp.runtime.decode_v2_oproj_tcol_eligible(ws, &tp.o) {
25541                    tp.runtime.decode_v2_stash_gated(ws, e, col)?;
25542                    crate::tp::set_tcol_oproj_stashed();
25543                    e.uninit(ws.o_out)?
25544                } else {
25545                    tp.runtime.decode_v2_finish(ws, e, &tp.o)?
25546                }
25547            } else {
25548                tp.runtime.decode_v2_finish(ws, e, &tp.o)?
25549            };
25550
25551            // Local shadow append: reads ws.k_shadow/ws.v_shadow on e's stream, which
25552            // decode_v2_finish ordered behind the root event. Same math and cache state
25553            // transitions as v1.
25554            let local = cache.kv[il]
25555                .as_mut()
25556                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
25557            if local.len != base_len || base_len + 1 > max_ctx {
25558                return Err(format!(
25559                    "Step TP layer {il} local cache changed during decode: \
25560                     len={} base={base_len} max={max_ctx}",
25561                    local.len
25562                )
25563                .into());
25564            }
25565            if crate::tp::no_local_shadow_on() {
25566                // Lengths advance, contents stay stale (graph-door precedent: decode reads
25567                // only the distributed TP caches; local contents feed spec/MTP scratch).
25568                local.len = base_len + 1;
25569                // MEMRA_LEN_MIRROR_LAZY=1: skip the 4B in-stream htod — nothing reads the
25570                // LOCAL device mirror in TP decode (the dcw fa reads the RANK counters),
25571                // and each tiny copy costs a compute->copy engine turnaround mid-layer.
25572                if !crate::tp::len_mirror_lazy_on() {
25573                    e.set_i32_one(&mut local.len_d, local.len as i32)?;
25574                }
25575            } else {
25576                let retain_from = window
25577                    .map(|window| {
25578                        let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
25579                        let rollback_retain =
25580                            base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
25581                        staged_retain.min(rollback_retain)
25582                    })
25583                    .unwrap_or(0);
25584                let write_row = e.prepare_kv_append(local, retain_from, 1)?;
25585                e.append_kv_quantized(
25586                    &ws.k_shadow,
25587                    &ws.v_shadow,
25588                    &mut local.k,
25589                    &mut local.v,
25590                    write_row,
25591                    local.kv_dim_k,
25592                    local.kv_dim_v,
25593                    local.k_tok_bytes,
25594                    local.v_tok_bytes,
25595                    false,
25596                )?;
25597                local.len = base_len + 1;
25598                e.set_i32_one(&mut local.len_d, local.len as i32)?;
25599            }
25600            Ok(output)
25601        })();
25602
25603        let output = match staged {
25604            Ok(output) => output,
25605            Err(error) => {
25606                let _ = tp.runtime.rollback_tp_kv_transaction(
25607                    cache.tp_kv[il]
25608                        .as_mut()
25609                        .expect("distributed cache checked above"),
25610                    transaction,
25611                );
25612                if let Some(local) = cache.kv[il].as_mut() {
25613                    local.len = base_len;
25614                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
25615                }
25616                return Err(error);
25617            }
25618        };
25619        // MEMRA_LEN_MIRROR_LAZY under FUSE_ROPE_APPEND: the fused append atomicInc owns
25620        // the rank counters (same value as the absolute re-set on full accept), so commit
25621        // host bookkeeping only — kills two 4B in-stream htods per layer. Non-fused dcw
25622        // keeps the absolute set (its appends do NOT inc).
25623        let lazy_commit = fuse_rope && crate::tp::len_mirror_lazy_on();
25624        if lazy_commit {
25625            if let Err(error) = tp.runtime.commit_tp_kv_transaction_external(
25626                cache.tp_kv[il]
25627                    .as_mut()
25628                    .expect("distributed cache checked above"),
25629                transaction,
25630                1,
25631            ) {
25632                let _ = tp.runtime.rollback_tp_kv_transaction(
25633                    cache.tp_kv[il]
25634                        .as_mut()
25635                        .expect("distributed cache checked above"),
25636                    transaction,
25637                );
25638                let local = cache.kv[il].as_mut().expect("local cache checked above");
25639                local.len = base_len;
25640                e.set_i32_one(&mut local.len_d, base_len as i32)?;
25641                return Err(error);
25642            }
25643        } else if let Err(error) = tp.runtime.commit_tp_kv_transaction(
25644            cache.tp_kv[il]
25645                .as_mut()
25646                .expect("distributed cache checked above"),
25647            transaction,
25648            1,
25649        ) {
25650            let _ = tp.runtime.rollback_tp_kv_transaction(
25651                cache.tp_kv[il]
25652                    .as_mut()
25653                    .expect("distributed cache checked above"),
25654                transaction,
25655            );
25656            let local = cache.kv[il].as_mut().expect("local cache checked above");
25657            local.len = base_len;
25658            e.set_i32_one(&mut local.len_d, base_len as i32)?;
25659            return Err(error);
25660        }
25661
25662        let committed = cache.tp_kv[il]
25663            .as_ref()
25664            .expect("distributed cache checked above")
25665            .committed_len();
25666        let local_len = cache.kv[il]
25667            .as_ref()
25668            .expect("local cache checked above")
25669            .len;
25670        if committed != local_len {
25671            return Err(format!(
25672                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
25673            )
25674            .into());
25675        }
25676        static V2_LOGGED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
25677        if !V2_LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) {
25678            eprintln!(
25679                "[step-tp-attn-v2] execute layer={} devices={:?} tokens=1 driver=v2 \
25680                 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
25681                 kv_cache_distributed=true kv_cache_hydrated={hydrated} \
25682                 attention_tensor_parallel=true attention_scope={} \
25683                 input_path=root-device-replicated gate={} gate_tensor_parallel={} \
25684                 gate_shards={} o_tensor_parallel=true o_reduce=root-device \
25685                 local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
25686                 bulk_p2p={} workspace=persistent ordering=evented output=e-device \
25687                 performance_claim=false (logged once; every decode layer runs this driver)",
25688                tp.layer,
25689                tp.devices,
25690                if window.is_some() {
25691                    "rank-local-swa-ring"
25692                } else {
25693                    "rank-local-global"
25694                },
25695                has_gate,
25696                use_gate_shards,
25697                if use_gate_shards {
25698                    "device-staged"
25699                } else if has_gate {
25700                    "root-staged"
25701                } else {
25702                    "none"
25703                },
25704                tp.runtime.transport_label(),
25705                tp.runtime.bulk_p2p(),
25706            );
25707        }
25708        Ok(output)
25709    }
25710
25711    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
25712    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
25713    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
25714    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
25715    /// requiring `attn_gate`).
25716    ///
25717    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
25718    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
25719    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
25720    #[allow(clippy::too_many_arguments)]
25721    pub(crate) fn step35_decode_attn(
25722        &self,
25723        e: &Engine,
25724        fa: &FullAttnLayer,
25725        il: usize,
25726        h: &CudaSlice<f32>,
25727        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
25728        pos_d: &CudaSlice<i32>,
25729        cache: &mut Cache,
25730    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25731        if fa
25732            .step_tp_qkv
25733            .as_ref()
25734            .is_some_and(|tp| tp.attention.is_some())
25735        {
25736            if pre_q.is_some() {
25737                return Err(
25738                    "rank-local Step attention preserves BF16 activations and refuses the q8_1 \
25739                     pre-quantized decode path"
25740                        .into(),
25741                );
25742            }
25743            return self.step35_tp_decode_attn_resident(e, fa, il, h, pos_d, cache);
25744        }
25745
25746        let geometry = self.step35_geom(il);
25747        let hd = geometry.head_dim_k as usize;
25748        let nkv = geometry.n_head_kv as usize;
25749        let nh = geometry.n_head as usize;
25750        let rbase = geometry.rope_base;
25751        let scale = geometry.attention_scale();
25752        let swa = geometry.window.is_some();
25753        let eps = self.cfg.rms_eps;
25754        let win = geometry.window.unwrap_or(0) as usize;
25755        let n_rot = geometry.n_rot as usize;
25756        let n_embd = self.cfg.n_embd as usize;
25757        let gw = fa
25758            .attn_gate
25759            .as_ref()
25760            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
25761
25762        let tp_qkv = if fa.step_tp_qkv.is_some() {
25763            if pre_q.is_some() {
25764                return Err(
25765                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
25766                     pre-quantized decode path"
25767                        .into(),
25768                );
25769            }
25770            self.full_attn_tp_qkv(e, fa, h, 1)?
25771        } else {
25772            None
25773        };
25774
25775        let (q0, k0, v0, gt) = match tp_qkv {
25776            Some(mut g3) => {
25777                let v = g3.pop().unwrap();
25778                let k = g3.pop().unwrap();
25779                let q = g3.pop().unwrap();
25780                let gt = e.matmul(gw, h, 1)?;
25781                (q, k, v, gt)
25782            }
25783            None => match pre_q {
25784                Some((hq, hdq)) => {
25785                    debug_assert!(
25786                        e.uses_q8_1_fast(gw),
25787                        "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
25788                         (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
25789                    );
25790                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
25791                        Some(t3) => t3,
25792                        None => (
25793                            e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
25794                            e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
25795                            e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
25796                        ),
25797                    };
25798                    let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
25799                    (a, b, c, gt)
25800                }
25801                None => {
25802                    if e.uses_q8_1_fast(&fa.wq)
25803                        && e.uses_q8_1_fast(&fa.wk)
25804                        && e.uses_q8_1_fast(&fa.wv)
25805                        && e.uses_q8_1_fast(gw)
25806                    {
25807                        let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
25808                        let (a, b, c) =
25809                            match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
25810                                Some(t3) => t3,
25811                                None => (
25812                                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
25813                                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
25814                                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
25815                                ),
25816                            };
25817                        let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
25818                        (a, b, c, gt)
25819                    } else {
25820                        (
25821                            e.matmul(&fa.wq, h, 1)?,
25822                            e.matmul(&fa.wk, h, 1)?,
25823                            e.matmul(&fa.wv, h, 1)?,
25824                            e.matmul(gw, h, 1)?,
25825                        )
25826                    }
25827                }
25828            },
25829        };
25830
25831        let mut q = e.uninit(nh * hd)?;
25832        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
25833        let mut k = e.uninit(nkv * hd)?;
25834        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
25835        let ff = if swa {
25836            None
25837        } else {
25838            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
25839        };
25840        #[cfg(debug_assertions)]
25841        if let Some(ff) = ff {
25842            crate::debug_assert_tensor_stream_device(
25843                ff,
25844                &e.stream(),
25845                "step35_decode_attn.rope_freqs",
25846            );
25847        }
25848        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
25849
25850        if std::env::var("MEMRA_NOFA").is_ok() {
25851            return Err(
25852                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
25853                        cache; unset MEMRA_NOFA to use fa_decode"
25854                    .into(),
25855            );
25856        }
25857        let kvl = cache.kv[il].as_mut().unwrap();
25858        let next_len = kvl.len + 1;
25859        let (off, t_kv) = if swa && next_len > win {
25860            (next_len - win, win)
25861        } else {
25862            (0, next_len)
25863        };
25864        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
25865        e.append_kv_quantized(
25866            &k,
25867            &v0,
25868            &mut kvl.k,
25869            &mut kvl.v,
25870            write_row,
25871            kvl.kv_dim_k,
25872            kvl.kv_dim_v,
25873            kvl.k_tok_bytes,
25874            kvl.v_tok_bytes,
25875            crate::Engine::kv_fp8_on(),
25876        )?;
25877        kvl.len = next_len;
25878        let physical = kvl.physical_rows(off, off + t_kv)?;
25879        let k_view = e.view_u8_range(
25880            &kvl.k,
25881            physical.start * kvl.k_tok_bytes,
25882            physical.end * kvl.k_tok_bytes,
25883        );
25884        let v_view = e.view_u8_range(
25885            &kvl.v,
25886            physical.start * kvl.v_tok_bytes,
25887            physical.end * kvl.v_tok_bytes,
25888        );
25889        let mut attn = e.uninit(nh * hd)?;
25890        e.fa_decode_kvmod(
25891            &q,
25892            &k_view,
25893            &v_view,
25894            &mut attn,
25895            hd,
25896            nh,
25897            nkv,
25898            t_kv,
25899            scale,
25900            kvl.k_tok_bytes,
25901            kvl.v_tok_bytes,
25902            crate::Engine::kv_fp8_on(),
25903        )?;
25904
25905        let mut ag = e.uninit(nh * hd)?;
25906        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
25907        self.full_attn_o(e, fa, &ag, 1)
25908    }
25909}
25910
25911// ===================================================================================== //
25912//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
25913//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
25914//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
25915//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
25916//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
25917//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
25918// ===================================================================================== //
25919impl HybridModel {
25920    pub fn is_gemma4_e4b(&self) -> bool {
25921        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
25922    }
25923
25924    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
25925    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
25926    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
25927    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
25928        let g = self.cfg.gemma4.as_ref().unwrap();
25929        let swa = g.swa_pattern[il];
25930        let hd = if swa {
25931            g.key_length_swa
25932        } else {
25933            g.key_length_global
25934        } as usize;
25935        let Mixer::Full(fa) = &self.layers[il].mixer else {
25936            panic!("e4b layer {il} not full-attn")
25937        };
25938        let nh = fa.wq.out_features() / hd;
25939        let nkv = fa.wk.out_features() / hd;
25940        (
25941            hd,
25942            nkv,
25943            nh,
25944            if swa {
25945                g.rope_base_swa
25946            } else {
25947                g.rope_base_global
25948            },
25949            1.0,
25950            swa,
25951        )
25952    }
25953
25954    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
25955    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
25956        self.layers[il]
25957            .gemma4
25958            .as_ref()
25959            .and_then(|b| b.e4b.as_ref())
25960            .and_then(|e4| e4.kv_share.map(|t| t as usize))
25961    }
25962
25963    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
25964    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
25965    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
25966    ///     (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
25967    fn gemma4_e4b_inp_pl(
25968        &self,
25969        e: &Engine,
25970        tokens: &[u32],
25971        x_scaled: &CudaSlice<f32>,
25972        t: usize,
25973    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25974        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
25975        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
25976    }
25977
25978    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
25979    fn gemma4_e4b_inp_pl_dev(
25980        &self,
25981        e: &Engine,
25982        tok_d: &CudaSlice<u32>,
25983        x_scaled: &CudaSlice<f32>,
25984        t: usize,
25985    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25986        let aux = self.gemma4_aux.as_ref().unwrap();
25987        let m = aux.e4b.as_ref().unwrap();
25988        let n_embd = self.cfg.n_embd as usize;
25989        let n_layer = self.layers.len();
25990        let width = m.n_epl * n_layer;
25991        let tbl = m.tok_tbl_gpu.get_or_init(|| {
25992            e.upload_u8(&m.tok_embd_bytes)
25993                .expect("e4b per-layer token table upload")
25994        });
25995        let mut a =
25996            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
25997        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
25998        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
25999        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
26000        let mut pn = e.uninit(t * width)?;
26001        e.rms_norm(
26002            &p,
26003            m.proj_norm.float_data(),
26004            &mut pn,
26005            m.n_epl,
26006            t * n_layer,
26007            self.cfg.rms_eps,
26008        )?;
26009        let mut out = e.uninit(t * width)?;
26010        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
26011        Ok(out)
26012    }
26013
26014    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
26015    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
26016    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
26017    /// already holds this forward's rows — the target runs earlier in the stack).
26018    #[allow(clippy::too_many_arguments)]
26019    fn gemma4_e4b_attn(
26020        &self,
26021        e: &Engine,
26022        il: usize,
26023        hq: &CudaSlice<i8>,
26024        hdq: &CudaSlice<f32>,
26025        pos_d: &CudaSlice<i32>,
26026        t: usize,
26027        cache: &mut Cache,
26028        dc_bucket: Option<usize>,
26029    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26030        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
26031        let eps = self.cfg.rms_eps;
26032        let aux = self.gemma4_aux.as_ref().unwrap();
26033        let ones = aux.ones(e);
26034        #[cfg(debug_assertions)]
26035        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
26036        let Mixer::Full(fa) = &self.layers[il].mixer else {
26037            unreachable!()
26038        };
26039        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
26040        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
26041        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
26042        let h0 = e.zeros(0)?;
26043        let h = &h0;
26044
26045        let ff = if swa {
26046            None
26047        } else {
26048            Some(
26049                aux.rope_freqs(e)
26050                    .expect("e4b global rope needs rope_freqs.weight"),
26051            )
26052        };
26053        #[cfg(debug_assertions)]
26054        if let Some(ff) = ff {
26055            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
26056        }
26057        let share = self.gemma4_e4b_kv_target(il);
26058        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
26059        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
26060        let mut q;
26061        if let Some(_tgt) = share {
26062            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
26063            q = e.uninit(t * nh * hd)?;
26064            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
26065            // empty; q0 stands in for the unused k/v pointers).
26066            let mut kdummy = e.uninit(1)?;
26067            let mut vdummy = e.uninit(1)?;
26068            e.rms_norm_qkv_rope(
26069                &q0,
26070                &q0,
26071                &q0,
26072                fa.q_norm.float_data(),
26073                fa.q_norm.float_data(),
26074                ones,
26075                &mut q,
26076                &mut kdummy,
26077                &mut vdummy,
26078                hd,
26079                self.gemma4_rope_dims(il),
26080                nh * t,
26081                0,
26082                pos_d,
26083                nh,
26084                1,
26085                base,
26086                1.0,
26087                ff,
26088                eps,
26089            )?;
26090        } else {
26091            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
26092            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
26093            // q|k|v rows — the cat norm+rope twin consumes it directly.
26094            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
26095            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
26096            q = e.uninit(t * nh * hd)?;
26097            let mut k = e.uninit(t * nkv * hd)?;
26098            let mut v = e.uninit(t * nkv * hd)?;
26099            if t == 1 && cat.is_some() {
26100                #[allow(clippy::unnecessary_unwrap)]
26101                // allow: the Some-guard sits in a multi-clause regime gate; if-let would reshape the arm structure
26102                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
26103                e.rms_norm_qkv_rope_cat(
26104                    &qkv0,
26105                    fa.q_norm.float_data(),
26106                    fa.k_norm.float_data(),
26107                    ones,
26108                    &mut q,
26109                    &mut k,
26110                    &mut v,
26111                    hd,
26112                    self.gemma4_rope_dims(il),
26113                    nh,
26114                    nkv,
26115                    pos_d,
26116                    nh,
26117                    nkv,
26118                    base,
26119                    1.0,
26120                    ff,
26121                    eps,
26122                )?;
26123            } else {
26124                let (q0, k0, v0) = match if t == 1 {
26125                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
26126                } else {
26127                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
26128                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
26129                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26130                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
26131                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
26132                    } else {
26133                        None
26134                    }
26135                } {
26136                    Some(triple) => triple,
26137                    None => (
26138                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
26139                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
26140                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
26141                    ), // E4B: real v (K != V)
26142                };
26143                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
26144                // the normed rows; V ones-rms, never roped).
26145                e.rms_norm_qkv_rope(
26146                    &q0,
26147                    &k0,
26148                    &v0,
26149                    fa.q_norm.float_data(),
26150                    fa.k_norm.float_data(),
26151                    ones,
26152                    &mut q,
26153                    &mut k,
26154                    &mut v,
26155                    hd,
26156                    self.gemma4_rope_dims(il),
26157                    nh * t,
26158                    nkv * t,
26159                    pos_d,
26160                    nh,
26161                    nkv,
26162                    base,
26163                    1.0,
26164                    ff,
26165                    eps,
26166                )?;
26167            }
26168            let kvl = cache.kv[il].as_mut().unwrap();
26169            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
26170            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
26171            // degenerate tok-0 stream, 2026-07-12).
26172            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
26173            if dc_bucket.is_some() {
26174                // DC arm (graph serving): append at the len_d slot, advance the counter
26175                // in-stream — replay-correct, no host len in the launch args. Host mirrors
26176                // are NOT touched here (the replay loop owns them; a bump at capture-record
26177                // time would double-count the capture iteration).
26178                debug_assert!(t == 1);
26179                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
26180                e.append_kv_quantized_row_dc_inc(
26181                    &k,
26182                    &v,
26183                    &mut kvl.k,
26184                    &mut kvl.v,
26185                    &mut kvl.len_d,
26186                    kvl.kv_dim_k,
26187                    kvl.kv_dim_v,
26188                    kvl.k_tok_bytes,
26189                    kvl.v_tok_bytes,
26190                    cls,
26191                )?;
26192            } else {
26193                e.append_kv_quantized_rows(
26194                    &k,
26195                    &v,
26196                    &mut kvl.k,
26197                    &mut kvl.v,
26198                    kvl.len,
26199                    t,
26200                    kvl.kv_dim_k,
26201                    kvl.kv_dim_v,
26202                    kvl.k_tok_bytes,
26203                    kvl.v_tok_bytes,
26204                    cls,
26205                )?;
26206                kvl.len += t;
26207            }
26208            kv_f32 = Some((k, v));
26209        }
26210        // attention: per-row causal fa over the (own or target) quantized cache. The cache
26211        // already contains this forward's rows in both arms; row i attends [.., base+i].
26212        let kvl_idx = share.unwrap_or(il);
26213        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
26214        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
26215        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
26216        let mut attn = e.uninit(t * nh * hd)?;
26217        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
26218        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
26219        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
26220        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
26221        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
26222        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
26223        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
26224        //     rows (the T=K verify kernel; the target appended this forward's rows already).
26225        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
26226        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
26227        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
26228        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
26229            if let Some((kf, vf)) = &kv_f32 {
26230                if hd == 256 && t <= win {
26231                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
26232                    return e.matmul(&fa.wo, &attn, t);
26233                }
26234                if hd == 256 && swa && t > win {
26235                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
26236                    return e.matmul(&fa.wo, &attn, t);
26237                }
26238                if hd == 512 && !swa {
26239                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
26240                    return e.matmul(&fa.wo, &attn, t);
26241                }
26242            } else if share.is_some() {
26243                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
26244                let k_view = e.view_u8(&kvl.k, kvl.k.len());
26245                let v_view = e.view_u8(&kvl.v, kvl.v.len());
26246                if hd == 256 && (!swa || t <= win) {
26247                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
26248                    e.fa_prefill_view(
26249                        &q,
26250                        &k_view,
26251                        &v_view,
26252                        &mut attn,
26253                        hd,
26254                        nh,
26255                        nkv,
26256                        t,
26257                        t,
26258                        scale,
26259                        true,
26260                        kvl.k_tok_bytes,
26261                        kvl.v_tok_bytes,
26262                        g,
26263                    )?;
26264                    return e.matmul(&fa.wo, &attn, t);
26265                }
26266                // remaining shared classes (swa above the window; hd512 globals): dequant
26267                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
26268                let kv_dim = nkv * hd;
26269                let mut kf = e.uninit(t * kv_dim)?;
26270                let mut vf = e.uninit(t * kv_dim)?;
26271                e.fa_dequant_kv_view_f32(
26272                    &k_view,
26273                    &v_view,
26274                    &mut kf,
26275                    &mut vf,
26276                    kv_dim,
26277                    kv_dim,
26278                    t,
26279                    kvl.k_tok_bytes,
26280                    kvl.v_tok_bytes,
26281                    g,
26282                )?;
26283                if hd == 512 {
26284                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
26285                } else {
26286                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
26287                }
26288                return e.matmul(&fa.wo, &attn, t);
26289            }
26290        }
26291        if let Some(bucket) = dc_bucket {
26292            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
26293            // fa_decode_dc over the live counter. len_d already advanced past this token
26294            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
26295            // counter (advanced when the target ran earlier in the stack).
26296            assert!(t == 1);
26297            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
26298            // and under the window every live t_kv sits below it — cap the capture bucket
26299            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
26300            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
26301            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
26302            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
26303                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
26304            } else {
26305                bucket
26306            };
26307            let k_view = e.view_u8(&kvl.k, kvl.k.len());
26308            let v_view = e.view_u8(&kvl.v, kvl.v.len());
26309            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
26310            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
26311            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
26312            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
26313            // captured into the dc graph like any other launch. Extending the cascade to
26314            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
26315            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
26316            // MEMRA_WPF=0 rollback seam.
26317            if crate::Engine::wpf_level() >= 1 {
26318                e.prefetch_weight_l2(&fa.wo)?;
26319            }
26320            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
26321            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
26322            if e.uses_q8_1_fast(&fa.wo) {
26323                let mut oq = e.alloc_i8_uninit(nh * hd)?;
26324                let mut od = e.zeros(nh * hd / 32)?;
26325                e.fa_decode_dc_q8(
26326                    &q,
26327                    &k_view,
26328                    &v_view,
26329                    &mut attn,
26330                    hd,
26331                    nh,
26332                    nkv,
26333                    &kvl.len_d,
26334                    bucket,
26335                    scale,
26336                    kvl.k_tok_bytes,
26337                    kvl.v_tok_bytes,
26338                    g,
26339                    Some((&mut oq, &mut od)),
26340                )?;
26341                return e.matmul_pre(&fa.wo, &oq, &od, &attn, t);
26342            }
26343            e.fa_decode_dc(
26344                &q,
26345                &k_view,
26346                &v_view,
26347                &mut attn,
26348                hd,
26349                nh,
26350                nkv,
26351                &kvl.len_d,
26352                bucket,
26353                scale,
26354                kvl.k_tok_bytes,
26355                kvl.v_tok_bytes,
26356                g,
26357            )?;
26358            return e.matmul(&fa.wo, &attn, t);
26359        }
26360        for i in 0..t {
26361            let avail = base_len + i + 1;
26362            let (off_tok, t_kv) = if swa && avail > win {
26363                (avail - win, win)
26364            } else {
26365                (0, avail)
26366            };
26367            let k_view = e.view_u8_range(
26368                &kvl.k,
26369                off_tok * kvl.k_tok_bytes,
26370                (off_tok + t_kv) * kvl.k_tok_bytes,
26371            );
26372            let v_view = e.view_u8_range(
26373                &kvl.v,
26374                off_tok * kvl.v_tok_bytes,
26375                (off_tok + t_kv) * kvl.v_tok_bytes,
26376            );
26377            let qv = e.view(&q, t * nh * hd);
26378            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
26379            let mut q_one = e.uninit(nh * hd)?;
26380            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
26381            let mut a_one = e.uninit(nh * hd)?;
26382            // read class MUST match the append class (globals are e4m3 under gkv): the
26383            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
26384            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
26385            e.fa_decode_kvmod(
26386                &q_one,
26387                &k_view,
26388                &v_view,
26389                &mut a_one,
26390                hd,
26391                nh,
26392                nkv,
26393                t_kv,
26394                scale,
26395                kvl.k_tok_bytes,
26396                kvl.v_tok_bytes,
26397                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
26398            )?;
26399            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
26400        }
26401        e.matmul(&fa.wo, &attn, t)
26402    }
26403
26404    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
26405    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
26406    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
26407    /// layer; does NOT advance cache.pos (caller owns pos).
26408    fn gemma4_e4b_trunk(
26409        &self,
26410        e: &Engine,
26411        tokens: &[u32],
26412        pos0: usize,
26413        cache: &mut Cache,
26414        head_last: bool,
26415    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26416        let n_embd = self.cfg.n_embd as usize;
26417        let t = tokens.len();
26418        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
26419        let pos_d = e.htod_i32(&pos)?;
26420        let mut x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
26421        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
26422        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
26423        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
26424    }
26425
26426    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
26427    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
26428    /// eager chain by construction: SAME functions, not twins).
26429    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
26430    fn gemma4_e4b_trunk_core(
26431        &self,
26432        e: &Engine,
26433        x_in: CudaSlice<f32>,
26434        inp_pl: CudaSlice<f32>,
26435        pos_d: &CudaSlice<i32>,
26436        t: usize,
26437        cache: &mut Cache,
26438        dc_bucket: Option<usize>,
26439        cap_logits: bool,
26440        head_last: bool,
26441    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26442        let n_embd = self.cfg.n_embd as usize;
26443        let eps = self.cfg.rms_eps;
26444        let n_layer = self.layers.len();
26445        let mut x = x_in;
26446        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
26447        let n_epl = aux_e4b.n_epl;
26448
26449        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
26450        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
26451        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
26452        // head rides matmul_pre too. First layer's pair comes from a standalone fused
26453        // norm+quant.
26454        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
26455        for il in 0..n_layer {
26456            let layer = &self.layers[il];
26457            let (hq, hdq) = match h_carry.take() {
26458                Some(p) => p,
26459                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
26460            };
26461            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
26462            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
26463            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
26464            let bits = layer.gemma4.as_ref().unwrap();
26465            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
26466            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
26467            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
26468            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
26469            // the fused single-phase reduction is NOT FP-order-identical to the unfused
26470            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
26471            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
26472            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
26473            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
26474            // gate dropped, decode AND verify ride the same fused chain — parity by
26475            // construction, VERIFY-GATE 0.000e0.
26476            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
26477            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
26478                e,
26479                layer,
26480                &o,
26481                &x,
26482                t,
26483                Some(layer.post_attn_norm.float_data()),
26484                fuse_exit,
26485            )?;
26486            let mut resid = e.uninit(t * n_embd)?;
26487            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
26488            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
26489            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
26490            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
26491            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
26492            let g = if fuse_exit {
26493                // sn here = RAW f0 (post_ffw deferred).
26494                let (rq, rd) = e.rms_pre_add_q8_1(
26495                    &sn,
26496                    bits.post_ffw_norm.float_data(),
26497                    &attn_out,
26498                    &mut resid,
26499                    n_embd,
26500                    t,
26501                    self.cfg.rms_eps,
26502                )?;
26503                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
26504            } else {
26505                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
26506                e.matmul(&e4b.inp_gate, &resid, t)?
26507            };
26508            let mut act = e.uninit(t * n_epl)?;
26509            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
26510                let ipv = e.view(&inp_pl, n_epl * n_layer);
26511                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
26512                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
26513                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
26514            } else {
26515                let mut inp_this = e.uninit(t * n_epl)?;
26516                e.copy_rows_strided(
26517                    &inp_pl,
26518                    &mut inp_this,
26519                    n_epl,
26520                    t,
26521                    n_epl * n_layer,
26522                    il * n_epl,
26523                )?;
26524                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
26525                e.matmul(&e4b.proj, &act, t)?
26526            };
26527            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
26528            // ONE launch (glue-fusion lane; last layer emits through output_norm).
26529            let next_norm = if il + 1 < n_layer {
26530                self.layers[il + 1].attn_norm.float_data()
26531            } else {
26532                self.output_norm.float_data()
26533            };
26534            let mut xn = e.uninit(t * n_embd)?;
26535            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
26536                &y,
26537                e4b.post_norm.float_data(),
26538                &resid,
26539                bits.layer_scale,
26540                next_norm,
26541                &mut xn,
26542                n_embd,
26543                t,
26544                eps,
26545            )?;
26546            h_carry = Some(pair);
26547            x = xn;
26548        }
26549        // the head consumes the last layer's fused (output_norm) emit. head_last callers
26550        // (prime, last_only forward) need only the final row's logits — the all-T head is
26551        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
26552        let (oq, odq) = h_carry.take().unwrap();
26553        let h0 = e.zeros(0)?;
26554        let hm = if head_last { 1 } else { t };
26555        let (hq, hd) = if head_last && t > 1 {
26556            let mut q1 = e.uninit_i8(n_embd)?;
26557            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
26558            let nb = n_embd / 32;
26559            let mut d1 = e.uninit(nb)?;
26560            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
26561            (q1, d1)
26562        } else {
26563            (oq, odq)
26564        };
26565        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
26566        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
26567        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
26568        // Logit-returning callers (host logits / spec prime) keep the capped emit.
26569        if cap_logits {
26570            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
26571            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
26572        }
26573        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
26574        Ok((ld, x))
26575    }
26576
26577    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
26578    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
26579    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
26580    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
26581    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
26582    /// covers exactly the layers that appended).
26583    pub fn gemma4_e4b_decode_step_t_am_dev(
26584        &self,
26585        e: &Engine,
26586        tok_d: &CudaSlice<u32>,
26587        t: usize,
26588        pos0: usize,
26589        cache: &mut Cache,
26590    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26591        let n_embd = self.cfg.n_embd as usize;
26592        let eps = self.cfg.rms_eps;
26593        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
26594        let pos_d = e.htod_i32(&pos)?;
26595        let embd_gpu = self
26596            .embd_gpu
26597            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
26598        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
26599        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
26600        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
26601        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
26602        let (ld, xp) =
26603            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
26604        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
26605        // emit is already capped, matching the eager chain bit-for-bit).
26606        let n_vocab = self.output.out_features();
26607        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
26608        for i in 0..t {
26609            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
26610        }
26611        let mut hn = e.uninit(t * n_embd)?;
26612        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
26613        cache.pos += t;
26614        Ok((vam, hn))
26615    }
26616
26617    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
26618    /// prime path — mirror of `gemma4_decode_step_t_h`).
26619    pub(crate) fn gemma4_e4b_decode_step_t_h(
26620        &self,
26621        e: &Engine,
26622        tokens: &[u32],
26623        pos0: usize,
26624        cache: &mut Cache,
26625    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26626        let n_embd = self.cfg.n_embd as usize;
26627        let eps = self.cfg.rms_eps;
26628        let t = tokens.len();
26629        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
26630        let mut hn = e.uninit(t * n_embd)?;
26631        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
26632        cache.pos += t;
26633        Ok((e.dtoh(&ld)?, hn))
26634    }
26635
26636    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
26637    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
26638    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
26639    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
26640    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
26641    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
26642    pub fn gemma4_e4b_decode_step_dcg(
26643        &self,
26644        e: &Engine,
26645        token_d: &mut CudaSlice<u32>,
26646        pos_d: &mut CudaSlice<i32>,
26647        embd_gpu: &CudaSlice<u8>,
26648        embd_qt: i32,
26649        embd_rb: usize,
26650        cache: &mut Cache,
26651        n_vocab: usize,
26652        bucket: usize,
26653    ) -> Result<(), Box<dyn std::error::Error>> {
26654        let n_embd = self.cfg.n_embd as usize;
26655        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
26656        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
26657        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
26658        let (ld, _x) =
26659            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
26660        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
26661        e.inc_seqlen(pos_d)?;
26662        Ok(())
26663    }
26664
26665    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
26666    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
26667    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
26668    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
26669    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
26670    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
26671    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
26672    #[allow(clippy::too_many_arguments)]
26673    pub fn gemma4_e4b_decode_step_dc(
26674        &self,
26675        e: &Engine,
26676        token_d: &CudaSlice<u32>,
26677        pos_d: &mut CudaSlice<i32>,
26678        embd_gpu: &CudaSlice<u8>,
26679        embd_qt: i32,
26680        embd_rb: usize,
26681        cache: &mut Cache,
26682        n_vocab: usize,
26683    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
26684        let n_embd = self.cfg.n_embd as usize;
26685        let eps = self.cfg.rms_eps;
26686        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
26687        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
26688        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
26689        let (ld, _x) =
26690            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
26691        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
26692        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
26693        e.inc_seqlen(pos_d)?;
26694        cache.pos += 1;
26695        let _ = eps;
26696        Ok(tok_out)
26697    }
26698
26699    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
26700    /// pre-output_norm hidden). Advances cache.pos.
26701    pub(crate) fn gemma4_e4b_decode_step_h(
26702        &self,
26703        e: &Engine,
26704        token: u32,
26705        cache: &mut Cache,
26706    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26707        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
26708        let logits = e.dtoh(&ld)?;
26709        cache.pos += 1;
26710        Ok((logits, x))
26711    }
26712
26713    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
26714    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
26715    /// fast; the prefill fa arms come later.
26716    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
26717    pub(crate) fn gemma4_e4b_prime(
26718        &self,
26719        e: &Engine,
26720        tokens: &[u32],
26721        cache: &mut Cache,
26722    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26723        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
26724        // process-kill as gemma4_prime — refuse per-request.
26725        if cache.pos != 0 {
26726            return Err(
26727                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
26728                        call or decode tokenwise"
26729                    .into(),
26730            );
26731        }
26732        let n_embd = self.cfg.n_embd as usize;
26733        let t = tokens.len();
26734        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
26735        cache.pos += t;
26736        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
26737        let xv = e.view(&x, t * n_embd);
26738        let row = xv.slice((t - 1) * n_embd..t * n_embd);
26739        let mut h_seed = e.uninit(n_embd)?;
26740        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
26741        Ok((last, h_seed, x))
26742    }
26743
26744    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
26745    pub(crate) fn gemma4_e4b_forward(
26746        &self,
26747        e: &Engine,
26748        tokens: &[u32],
26749        last_only: bool,
26750    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
26751        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
26752        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
26753        e.dtoh(&ld) // head_last already reduced to the final row when last_only
26754    }
26755}
26756
26757#[cfg(test)]
26758mod prime_chunk_schedule_tests {
26759    use super::{
26760        CUDA_GRID_YZ_MAX, PRIME_CHUNK_LAUNCH_CAP, PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, PrimePpSignal,
26761        PrimePpStageChannels, PrimePpWaveCredits, PrimePpWaveSlot, active_matrix_values,
26762        align_prime_ranges_to_gdn, dynamic_prime_chunk_ranges, explicit_prime_chunk,
26763        fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring, move_prime_cache_layers,
26764        parse_step_ep_grouped_prefill, parse_step_tp_prefill, prime_cache_stage_for_layer,
26765        recv_prime_pp_signal, restore_prime_cache_layers, step_grouped_decode_shape,
26766        step_grouped_prefill_shape, step_tp_prefill_shape, validate_step_prime_batch_modes,
26767    };
26768
26769    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
26770        ranges.iter().map(|(start, end)| end - start).collect()
26771    }
26772
26773    #[allow(clippy::manual_clamp)] // allow: the min/max chain mirrors the reference arithmetic order in pinned sizing/quant math
26774    fn auto_chunk(t: usize) -> usize {
26775        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
26776    }
26777
26778    /// The ornith cold-long 66k defect (darklanes research/ornith-move-20260829 F2,
26779    /// re-hit on prod 2026-09-01): MEMRA_PRIME_CHUNK=0 ("monolithic") must not schedule
26780    /// a prime range wider than the CUDA grid.y limit, and must stay byte-identical
26781    /// (single monolithic range) for every prompt the limit allows.
26782    /// memra#144: the non-hyper prime charge is chunk-bounded under a chunked prime and
26783    /// prompt-scaled under a monolithic one, with the whole-prompt hiddens always charged.
26784    #[test]
26785    fn prime_workspace_shape_is_chunk_bounded_and_prompt_scaled_when_monolithic() {
26786        let shape = crate::hybrid_forward::PrimeWorkspaceShape {
26787            call_row_bytes: 42 * 2048 + 14 * 2048,
26788            prompt_row_bytes: 4 * 2048,
26789            n_layers: 41,
26790        };
26791        let chunked_60k = shape.admission_bytes_with_call_rows(60_000, 4096);
26792        let chunked_257k = shape.admission_bytes_with_call_rows(257_000, 4096);
26793        let mono_60k = shape.admission_bytes_with_call_rows(60_000, 60_000);
26794        let mono_257k = shape.admission_bytes_with_call_rows(257_000, 257_000);
26795        // chunked: the call term is the same 4096 rows at both lengths; only hiddens grow
26796        assert_eq!(
26797            chunked_257k - chunked_60k,
26798            (257_000 - 60_000) * shape.prompt_row_bytes
26799        );
26800        // monolithic: the call term scales with the prompt
26801        assert_eq!(
26802            mono_257k - mono_60k,
26803            (257_000 - 60_000) * (shape.call_row_bytes + shape.prompt_row_bytes)
26804        );
26805        assert!(
26806            mono_257k > chunked_257k * 4,
26807            "monolithic 257k must dwarf the chunked charge"
26808        );
26809        // a prompt shorter than the chunk pays only its own rows
26810        assert_eq!(
26811            shape.admission_bytes_with_call_rows(100, 4096),
26812            100 * (shape.call_row_bytes + shape.prompt_row_bytes)
26813        );
26814    }
26815
26816    #[test]
26817    // assertions_on_constants: the constant relation (cap + fold headroom fits the CUDA
26818    // grid wall) IS the invariant under test; a red here means someone moved a constant.
26819    #[allow(clippy::assertions_on_constants)]
26820    fn monolithic_prime_chunk_caps_at_the_cuda_launch_wall() {
26821        // Ring OFF (ornith serving shape): 0 maps to the launch cap, larger explicit
26822        // values cap there too, workable explicit values pass through untouched.
26823        assert_eq!(explicit_prime_chunk(0, false), PRIME_CHUNK_LAUNCH_CAP);
26824        assert_eq!(explicit_prime_chunk(100_000, false), PRIME_CHUNK_LAUNCH_CAP);
26825        assert_eq!(explicit_prime_chunk(4096, false), 4096);
26826        assert_eq!(
26827            explicit_prime_chunk(PRIME_CHUNK_LAUNCH_CAP, false),
26828            PRIME_CHUNK_LAUNCH_CAP
26829        );
26830        // Ring ON keeps the historical PRIME_CHUNK_MAX_TOKENS clamp exactly.
26831        assert_eq!(
26832            explicit_prime_chunk(0, true),
26833            crate::cache::PRIME_CHUNK_MAX_TOKENS
26834        );
26835        assert_eq!(
26836            explicit_prime_chunk(100_000, true),
26837            crate::cache::PRIME_CHUNK_MAX_TOKENS
26838        );
26839        assert_eq!(explicit_prime_chunk(512, true), 512);
26840        // The fold headroom the cap exists for.
26841        assert!(PRIME_CHUNK_LAUNCH_CAP + PRIME_MIN_T - 1 <= CUDA_GRID_YZ_MAX);
26842    }
26843
26844    #[test]
26845    fn capped_monolithic_ranges_are_identical_below_the_wall_and_legal_above() {
26846        let chunk = explicit_prime_chunk(0, false);
26847        // Every t the CUDA limit allows keeps the exact pre-fix monolithic schedule:
26848        // one range covering the whole prompt (t <= chunk directly, or the < PRIME_MIN_T
26849        // tail folds the split back into a single range).
26850        for t in [
26851            PRIME_MIN_T,
26852            4096,
26853            61_000,
26854            64_984,
26855            PRIME_CHUNK_LAUNCH_CAP,
26856            PRIME_CHUNK_LAUNCH_CAP + 1,
26857            CUDA_GRID_YZ_MAX,
26858        ] {
26859            assert_eq!(
26860                fixed_prime_chunk_ranges_for_ring(t, chunk, false),
26861                vec![(0, t)],
26862                "t={t} must stay a single monolithic range"
26863            );
26864        }
26865        // Above the wall — the sizes the campaign measured failing, the F2 bracket's
26866        // first FAIL, and the boundary — every scheduled range must be launch-legal,
26867        // contiguous, and full-coverage.
26868        for t in [65_536, 65_643, 66_045, 79_717, 82_440, 262_144] {
26869            let ranges = fixed_prime_chunk_ranges_for_ring(t, chunk, false);
26870            assert!(ranges.len() >= 2, "t={t} must chunk");
26871            let mut cursor = 0usize;
26872            for &(start, end) in &ranges {
26873                assert_eq!(start, cursor, "t={t}: ranges must be contiguous");
26874                assert!(
26875                    end - start <= CUDA_GRID_YZ_MAX,
26876                    "t={t}: range width {} exceeds the CUDA grid.y limit",
26877                    end - start
26878                );
26879                assert!(
26880                    end - start >= PRIME_MIN_T,
26881                    "t={t}: range width {} below PRIME_MIN_T",
26882                    end - start
26883                );
26884                cursor = end;
26885            }
26886            assert_eq!(cursor, t, "t={t}: ranges must cover the prompt");
26887        }
26888        // Dense sweep across the boundary band: no width may ever exceed the limit.
26889        for t in (CUDA_GRID_YZ_MAX - 64)..=(CUDA_GRID_YZ_MAX + 2 * PRIME_MIN_T + 64) {
26890            for &(start, end) in &fixed_prime_chunk_ranges_for_ring(t, chunk, false) {
26891                assert!(
26892                    end - start <= CUDA_GRID_YZ_MAX,
26893                    "t={t} width {}",
26894                    end - start
26895                );
26896            }
26897        }
26898    }
26899
26900    #[test]
26901    fn ppn_prime_cache_partition_moves_and_restores_every_layer() {
26902        let round_trip = |fence: &[usize], layers: usize| {
26903            let original: Vec<Option<usize>> = (0..layers).map(Some).collect();
26904            let mut parent = original.clone();
26905            let mut stages: Vec<Vec<Option<usize>>> =
26906                (0..fence.len() - 1).map(|_| vec![None; layers]).collect();
26907
26908            move_prime_cache_layers(&mut parent, &mut stages, fence);
26909            assert!(parent.iter().all(Option::is_none));
26910            for layer in 0..layers {
26911                let owner = prime_cache_stage_for_layer(fence, layer);
26912                for (stage, values) in stages.iter().enumerate() {
26913                    assert_eq!(values[layer], (stage == owner).then_some(layer));
26914                }
26915            }
26916
26917            restore_prime_cache_layers(&mut parent, &mut stages, fence);
26918            assert_eq!(parent, original);
26919            assert!(stages.iter().flatten().all(Option::is_none));
26920        };
26921
26922        // Layers beyond the trunk fence end model MTP/tail state and remain last-stage owned.
26923        round_trip(&[0, 5, 8], 10);
26924        round_trip(&[0, 2, 5, 8], 10);
26925        round_trip(&[0, 1, 3, 6, 8], 10);
26926    }
26927
26928    #[test]
26929    fn ppn_prime_wave_credit_requires_the_exact_oldest_wave_and_slot() {
26930        let mut credits = PrimePpWaveCredits::default();
26931        let wave0 = PrimePpWaveSlot { wave: 0, slot: 1 };
26932        let wave1 = PrimePpWaveSlot { wave: 1, slot: 0 };
26933        credits.record_send(wave0).unwrap();
26934        assert_eq!(credits.release_required(), None);
26935        credits.record_send(wave1).unwrap();
26936        assert_eq!(credits.release_required(), Some(wave0));
26937
26938        assert!(
26939            credits
26940                .record_release(PrimePpWaveSlot { wave: 0, slot: 0 })
26941                .unwrap_err()
26942                .contains("does not match oldest pending")
26943        );
26944        assert_eq!(credits.release_required(), Some(wave0));
26945        credits.record_release(wave0).unwrap();
26946        credits
26947            .record_send(PrimePpWaveSlot { wave: 2, slot: 1 })
26948            .unwrap();
26949        assert!(
26950            credits
26951                .record_send(PrimePpWaveSlot { wave: 4, slot: 0 })
26952                .unwrap_err()
26953                .contains("while wave 3 was next")
26954        );
26955        assert!(
26956            credits
26957                .record_send(PrimePpWaveSlot { wave: 3, slot: 1 })
26958                .unwrap_err()
26959                .contains("reused slot 1")
26960        );
26961    }
26962
26963    #[test]
26964    fn ppn_prime_wave_signal_reports_order_error_injected_error_and_closure() {
26965        let expected = PrimePpWaveSlot { wave: 2, slot: 1 };
26966
26967        let (sender, receiver) = std::sync::mpsc::channel();
26968        sender.send(PrimePpSignal::Slot(expected)).unwrap();
26969        assert_eq!(
26970            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap(),
26971            expected
26972        );
26973
26974        let (sender, receiver) = std::sync::mpsc::channel();
26975        sender
26976            .send(PrimePpSignal::Slot(PrimePpWaveSlot { wave: 3, slot: 1 }))
26977            .unwrap();
26978        assert!(
26979            recv_prime_pp_signal(&receiver, expected, true, "test")
26980                .unwrap_err()
26981                .contains("expected wave/slot")
26982        );
26983
26984        let (sender, receiver) = std::sync::mpsc::channel();
26985        sender
26986            .send(PrimePpSignal::Error("injected stage failure".into()))
26987            .unwrap();
26988        assert_eq!(
26989            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap_err(),
26990            "injected stage failure"
26991        );
26992
26993        let (upstream_sender, upstream_receiver) = std::sync::mpsc::channel();
26994        let (outgoing_sender, outgoing_receiver) = std::sync::mpsc::channel();
26995        let (_release_sender, released_downstream) = std::sync::mpsc::channel();
26996        PrimePpStageChannels {
26997            incoming: None,
26998            release_upstream: Some(upstream_sender),
26999            outgoing: outgoing_sender,
27000            released_downstream,
27001        }
27002        .notify_failure("injected worker error");
27003        assert_eq!(
27004            recv_prime_pp_signal(&upstream_receiver, expected, false, "test").unwrap_err(),
27005            "injected worker error"
27006        );
27007        assert_eq!(
27008            recv_prime_pp_signal(&outgoing_receiver, expected, false, "test").unwrap_err(),
27009            "injected worker error"
27010        );
27011
27012        let (sender, receiver) = std::sync::mpsc::channel::<PrimePpSignal>();
27013        drop(sender);
27014        assert!(
27015            recv_prime_pp_signal(&receiver, expected, true, "test")
27016                .unwrap_err()
27017                .contains("channel closed while waiting for wave 2")
27018        );
27019    }
27020
27021    /// TOOTH for the PP-auto-ranges GDN grid law (lane/hermes-perf-fixes, 2026-08-23;
27022    /// primegrid pattern): the AUTO schedules put internal prime-call boundaries OFF the
27023    /// WY-chunk grid — the broken arm must be demonstrably off-grid, and the aligned twin
27024    /// must land every boundary on it without changing coverage.
27025    #[test]
27026    fn auto_prime_ranges_align_to_the_gdn_grid() {
27027        let c = 32usize; // shipped MEMRA_GDN_CHUNK default/clamp floor
27028        let assert_covers = |ranges: &[(usize, usize)], t: usize| {
27029            assert_eq!(ranges.first().map(|&(s, _)| s), Some(0));
27030            assert_eq!(ranges.last().map(|&(_, e)| e), Some(t));
27031            for w in ranges.windows(2) {
27032                assert_eq!(w[0].1, w[1].0, "ranges must stay contiguous");
27033            }
27034            assert!(ranges.iter().all(|&(s, e)| e > s), "no empty range");
27035        };
27036
27037        // The PP-2 auto geometry at a real agentic length: t=9510 -> fill = 1189 (div_ceil
27038        // by 8), every internal boundary off the 32 grid — the falsified-identity arm.
27039        let t = 9510usize;
27040        let fill = auto_chunk(t);
27041        let fixed = fixed_prime_chunk_ranges(t, fill);
27042        assert!(
27043            fixed[..fixed.len() - 1].iter().any(|&(_, e)| e % c != 0),
27044            "broken arm vanished: fixed auto boundaries all landed on-grid"
27045        );
27046        let dynamic = dynamic_prime_chunk_ranges(t, fill, &fixed);
27047        assert!(
27048            dynamic[..dynamic.len() - 1]
27049                .iter()
27050                .any(|&(_, e)| e % c != 0),
27051            "broken arm vanished: dynamic auto boundaries all landed on-grid"
27052        );
27053
27054        for ranges in [&fixed, &dynamic] {
27055            let aligned = align_prime_ranges_to_gdn(ranges, t, c);
27056            assert_covers(&aligned, t);
27057            for &(_, e) in &aligned[..aligned.len() - 1] {
27058                assert_eq!(e % c, 0, "internal boundary {e} off the {c}-grid");
27059            }
27060            // boundaries only move DOWN, at most c-1 tokens.
27061            for (&(_, a), &(_, b)) in aligned.iter().zip(ranges.iter()) {
27062                assert!(a <= b && b - a < c);
27063            }
27064        }
27065
27066        // Collapse/merge: boundaries inside one grid cell fuse instead of emitting an
27067        // empty range; the schedule survives degenerate short fills.
27068        let tight = vec![(0usize, 33usize), (33, 40), (40, 200)];
27069        let aligned = align_prime_ranges_to_gdn(&tight, 200, c);
27070        assert_covers(&aligned, 200);
27071        assert_eq!(aligned, vec![(0, 32), (32, 200)]);
27072
27073        // No-ops: single range, c=0 (grid off), already-aligned schedules.
27074        assert_eq!(align_prime_ranges_to_gdn(&[(0, 200)], 200, c), [(0, 200)]);
27075        assert_eq!(align_prime_ranges_to_gdn(&tight, 200, 0), tight.as_slice());
27076        let on_grid = vec![(0usize, 128usize), (128, 256), (256, 300)];
27077        assert_eq!(
27078            align_prime_ranges_to_gdn(&on_grid, 300, c),
27079            on_grid.as_slice()
27080        );
27081    }
27082
27083    #[test]
27084    fn active_matrix_prefix_scopes_reused_prime_slabs() {
27085        assert_eq!(
27086            active_matrix_values(40 * 4096, 29, 4096, "activation").unwrap(),
27087            29 * 4096
27088        );
27089        assert_eq!(
27090            active_matrix_values(29 * 4096, 29, 4096, "activation").unwrap(),
27091            29 * 4096
27092        );
27093        assert_eq!(
27094            active_matrix_values(29 * 4096, 24, 4096, "activation").unwrap(),
27095            24 * 4096
27096        );
27097        assert!(active_matrix_values(28 * 4096, 29, 4096, "activation").is_err());
27098        assert!(active_matrix_values(usize::MAX, usize::MAX, 2, "activation").is_err());
27099    }
27100
27101    #[test]
27102    fn step_tp_prefill_batch_refuses_before_scheduler_fallback() {
27103        assert!(validate_step_prime_batch_modes(false, false).is_ok());
27104
27105        let grouped_without_tp = validate_step_prime_batch_modes(false, true).unwrap_err();
27106        assert!(grouped_without_tp.contains("requires MEMRA_STEP_TP_PREFILL=1"));
27107
27108        for grouped in [false, true] {
27109            let err = validate_step_prime_batch_modes(true, grouped).unwrap_err();
27110            assert!(err.contains("did not clear the live-server performance gate"));
27111            assert!(err.contains("per-session grouped prefill"));
27112        }
27113    }
27114
27115    #[test]
27116    fn step_grouped_path_is_eager_single_token_only() {
27117        assert!(step_grouped_decode_shape(false, 1));
27118        assert!(!step_grouped_decode_shape(true, 1));
27119        assert!(!step_grouped_decode_shape(false, 2));
27120        assert!(!step_grouped_decode_shape(true, 2));
27121    }
27122
27123    #[test]
27124    fn step_grouped_prefill_door_is_strict_and_capacity_bounded() {
27125        assert!(!parse_step_ep_grouped_prefill(None).unwrap());
27126        assert!(!parse_step_ep_grouped_prefill(Some("")).unwrap());
27127        assert!(!parse_step_ep_grouped_prefill(Some("0")).unwrap());
27128        assert!(parse_step_ep_grouped_prefill(Some("1")).unwrap());
27129        assert!(parse_step_ep_grouped_prefill(Some("true")).is_err());
27130        assert!(parse_step_ep_grouped_prefill(Some("2")).is_err());
27131
27132        assert!(step_grouped_prefill_shape(true, true, PRIME_MIN_T));
27133        assert!(step_grouped_prefill_shape(
27134            true,
27135            true,
27136            crate::cache::PRIME_CHUNK_MAX_TOKENS,
27137        ));
27138        assert!(!step_grouped_prefill_shape(true, true, PRIME_MIN_T - 1,));
27139        assert!(!step_grouped_prefill_shape(
27140            true,
27141            true,
27142            crate::cache::PRIME_CHUNK_MAX_TOKENS + 1,
27143        ));
27144        assert!(!step_grouped_prefill_shape(false, true, PRIME_MIN_T));
27145        assert!(!step_grouped_prefill_shape(true, false, PRIME_MIN_T));
27146    }
27147
27148    #[test]
27149    fn step_tp_prefill_door_is_strict_and_default_off() {
27150        assert!(!parse_step_tp_prefill(None).unwrap());
27151        assert!(!parse_step_tp_prefill(Some("")).unwrap());
27152        assert!(!parse_step_tp_prefill(Some("0")).unwrap());
27153        assert!(parse_step_tp_prefill(Some("1")).unwrap());
27154        assert!(parse_step_tp_prefill(Some("true")).is_err());
27155        assert!(parse_step_tp_prefill(Some("2")).is_err());
27156    }
27157
27158    #[test]
27159    fn step_tp_prefill_requires_a_qualified_even_rank_shape() {
27160        assert!(step_tp_prefill_shape(
27161            true,
27162            PRIME_MIN_T,
27163            4,
27164            true,
27165            true,
27166            false,
27167        ));
27168        assert!(!step_tp_prefill_shape(
27169            false,
27170            PRIME_MIN_T,
27171            4,
27172            true,
27173            true,
27174            false,
27175        ));
27176        assert!(!step_tp_prefill_shape(
27177            true,
27178            PRIME_MIN_T - 1,
27179            4,
27180            true,
27181            true,
27182            false,
27183        ));
27184        // TP2 admits (2026-08-25); odd/1-card placements still refuse.
27185        assert!(step_tp_prefill_shape(
27186            true,
27187            PRIME_MIN_T,
27188            2,
27189            true,
27190            true,
27191            false
27192        ));
27193        assert!(!step_tp_prefill_shape(
27194            true,
27195            PRIME_MIN_T,
27196            1,
27197            true,
27198            true,
27199            false
27200        ));
27201        assert!(!step_tp_prefill_shape(
27202            true,
27203            PRIME_MIN_T,
27204            3,
27205            true,
27206            true,
27207            false
27208        ));
27209        assert!(!step_tp_prefill_shape(
27210            true,
27211            PRIME_MIN_T,
27212            4,
27213            false,
27214            true,
27215            false,
27216        ));
27217        assert!(!step_tp_prefill_shape(
27218            true,
27219            PRIME_MIN_T,
27220            4,
27221            true,
27222            false,
27223            false,
27224        ));
27225        assert!(!step_tp_prefill_shape(
27226            true,
27227            PRIME_MIN_T,
27228            4,
27229            true,
27230            true,
27231            true,
27232        ));
27233    }
27234
27235    #[test]
27236    fn fixed_schedule_retains_measured_geometry() {
27237        assert_eq!(
27238            sizes(&fixed_prime_chunk_ranges(461, 128)),
27239            vec![128, 128, 128, 77]
27240        );
27241        assert_eq!(
27242            sizes(&fixed_prime_chunk_ranges(1833, 230)),
27243            vec![230, 230, 230, 230, 230, 230, 230, 223]
27244        );
27245        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
27246        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
27247        assert_eq!(capped, vec![4096, 4088, 16]);
27248        assert!(capped.iter().all(|&rows| rows <= 4096));
27249        assert_eq!(
27250            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
27251            vec![4100],
27252            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
27253        );
27254    }
27255
27256    #[test]
27257    fn dynamic_schedule_matches_registered_shapes() {
27258        let cases = [
27259            (461, vec![64, 141, 132, 124]),
27260            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
27261            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
27262        ];
27263        for (t, expected) in cases {
27264            let chunk = auto_chunk(t);
27265            let fixed = fixed_prime_chunk_ranges(t, chunk);
27266            assert_eq!(
27267                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
27268                expected
27269            );
27270        }
27271    }
27272
27273    #[test]
27274    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
27275        for t in 256..=8192 {
27276            let chunk = auto_chunk(t);
27277            let fixed = fixed_prime_chunk_ranges(t, chunk);
27278            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
27279            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
27280            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
27281            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
27282            for pair in dynamic.windows(2) {
27283                assert_eq!(pair[0].1, pair[1].0, "T={t}");
27284            }
27285            assert!(
27286                dynamic
27287                    .iter()
27288                    .all(|(start, end)| end - start >= PRIME_MIN_T),
27289                "T={t} sizes={:?}",
27290                sizes(&dynamic)
27291            );
27292            if dynamic.len() >= 3 {
27293                let chunk_sizes = sizes(&dynamic);
27294                assert!(
27295                    chunk_sizes[0] < chunk_sizes[1],
27296                    "T={t} sizes={chunk_sizes:?}"
27297                );
27298                assert!(
27299                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
27300                    "T={t} sizes={chunk_sizes:?}"
27301                );
27302            }
27303        }
27304    }
27305}
27306
27307#[cfg(test)]
27308mod page_prefetch_tests {
27309    use super::{
27310        grouped_worker_prefetch_position, page_prefetch_positions,
27311        page_prefetch_window_from_values, worker_prefetch_positions,
27312    };
27313
27314    #[test]
27315    fn page_prefetch_window_keeps_existing_opt_in_default() {
27316        assert_eq!(page_prefetch_window_from_values(false, None), 0);
27317        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
27318        assert_eq!(page_prefetch_window_from_values(true, None), 1);
27319        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
27320        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
27321        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
27322    }
27323
27324    #[test]
27325    fn rolling_page_prefetch_advises_each_future_expert_once() {
27326        let advised: Vec<_> = (0..7)
27327            .flat_map(|position| page_prefetch_positions(position, 7, 3))
27328            .collect();
27329        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
27330
27331        let one_ahead: Vec<_> = (0..4)
27332            .flat_map(|position| page_prefetch_positions(position, 4, 1))
27333            .collect();
27334        assert_eq!(one_ahead, vec![1, 2, 3]);
27335        assert!(page_prefetch_positions(0, 4, 0).is_empty());
27336    }
27337
27338    #[test]
27339    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
27340        assert_eq!(grouped_worker_prefetch_position(0, None), None);
27341        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
27342            .chain(
27343                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
27344            )
27345            .collect();
27346        assert_eq!(positions, vec![0, 1, 2, 3]);
27347        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
27348    }
27349
27350    #[test]
27351    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
27352        let queued: Vec<_> = (0..8)
27353            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
27354            .collect();
27355        assert_eq!(queued, (0..8).collect::<Vec<_>>());
27356
27357        let one_at_a_time: Vec<_> = (0..4)
27358            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
27359            .collect();
27360        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
27361        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
27362    }
27363}
27364
27365pub struct G4DcSlots {
27366    x: CudaSlice<f32>,
27367    xn: CudaSlice<f32>,
27368    cur: CudaSlice<f32>,
27369    hq: CudaSlice<i8>,
27370    hd_: CudaSlice<f32>,
27371    q0: CudaSlice<f32>,
27372    k0: CudaSlice<f32>,
27373    v0: CudaSlice<f32>,
27374    q: CudaSlice<f32>,
27375    k: CudaSlice<f32>,
27376    v: CudaSlice<f32>,
27377    attn: CudaSlice<f32>,
27378    o: CudaSlice<f32>,
27379    attn_out: CudaSlice<f32>,
27380    zsh: CudaSlice<f32>,
27381    zq: CudaSlice<i8>,
27382    zd: CudaSlice<f32>,
27383    gate: CudaSlice<f32>,
27384    up: CudaSlice<f32>,
27385    act: CudaSlice<f32>,
27386    actq: CudaSlice<i8>,
27387    actd: CudaSlice<f32>,
27388    f0: CudaSlice<f32>,
27389    sn: CudaSlice<f32>,
27390    hn: CudaSlice<f32>,
27391    logits: CudaSlice<f32>,
27392}
27393
27394/// Whole-token decode graph state (step TP graph increment B). One stitched multi-device
27395/// parent per fa bucket, plus the persistent host->graph plumbing: the device token id the
27396/// in-graph embed gathers, the device position the ropes read (advanced in-graph), and the
27397/// fixed logits stage the head writes.
27398pub struct Step35TokenGraphState {
27399    /// (bucket_max, graph) — bucket keyed by the fa split geometry (fa_geom_eager).
27400    pub graphs: Vec<(usize, crate::tp::TokenGraph)>,
27401    pub token_d: cudarc::driver::CudaSlice<u32>,
27402    pub pos_d: cudarc::driver::CudaSlice<i32>,
27403    pub logits_stage: cudarc::driver::CudaSlice<f32>,
27404    /// Cross-child intermediates MUST live at fixed addresses (graph mem nodes remap at
27405    /// launch, so an alloc made inside one captured child is not referable from another):
27406    /// the running residual, the post-attention pair, the shared-expert row, and the
27407    /// e-context mirrors of the root-produced attention output and K/V shadow rows.
27408    pub x: cudarc::driver::CudaSlice<f32>,
27409    pub x1: cudarc::driver::CudaSlice<f32>,
27410    pub mixed_stage: cudarc::driver::CudaSlice<f32>,
27411    pub sh_stage: cudarc::driver::CudaSlice<f32>,
27412    pub k_shadow_stage: cudarc::driver::CudaSlice<f32>,
27413    pub v_shadow_stage: cudarc::driver::CudaSlice<f32>,
27414    /// Alloc-free e-section scratch (child graphs cannot contain mem nodes): router logits,
27415    /// shared-expert gate/up/act rows + sigmoid scalar, dense-FFN z/gate/up/act, head hidden.
27416    pub router_logits: cudarc::driver::CudaSlice<f32>,
27417    pub shexp_gate: cudarc::driver::CudaSlice<f32>,
27418    pub shexp_up: cudarc::driver::CudaSlice<f32>,
27419    pub shexp_act: cudarc::driver::CudaSlice<f32>,
27420    pub gate_sig: cudarc::driver::CudaSlice<f32>,
27421    pub dense_z: cudarc::driver::CudaSlice<f32>,
27422    pub dense_gate: cudarc::driver::CudaSlice<f32>,
27423    pub dense_up: cudarc::driver::CudaSlice<f32>,
27424    pub dense_act: cudarc::driver::CudaSlice<f32>,
27425    pub hn: cudarc::driver::CudaSlice<f32>,
27426    /// MEMRA_TG_PROBE_LAYER diagnostics: capture-time copies of layer K's attention output
27427    /// and post-FFN residual, dumped after replay for graph-vs-eager layer bisection.
27428    pub probe_mixed: cudarc::driver::CudaSlice<f32>,
27429    pub probe_x: cudarc::driver::CudaSlice<f32>,
27430    /// Chunk-loop (F-lite): device token history ring + its device write index, filled by
27431    /// the in-graph tail argmax chain; host reads the ring once per chunk.
27432    pub token_hist: cudarc::driver::CudaSlice<u32>,
27433    pub hist_idx: cudarc::driver::CudaSlice<i32>,
27434}
27435
27436impl HybridModel {
27437    /// Whole-token decode graph (step TP graph increment B): ONE stitched multi-device parent
27438    /// replays the entire 45-layer token — the launch-collapse the per-layer minis could not
27439    /// reach. Returns Some(logits) when the graph handled the token, None for eager fallback
27440    /// (door off, ineligible class, any layer's ring would rebase, or a bucket boundary that
27441    /// needs a rebuild this token).
27442    ///
27443    /// v1 SCOPE (diagnostic door, default OFF): the LOCAL shadow caches advance their lengths
27444    /// but not their contents under this door (the TP rank caches are fully maintained
27445    /// in-graph via the dcw counters). Sessions relying on shadow CONTENT (save/rollback)
27446    /// must not run with the door on until the local-dcw twin lands.
27447    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
27448    pub(crate) fn step35_token_graph_step(
27449        &self,
27450        e: &Engine,
27451        token: u32,
27452        cache: &mut Cache,
27453    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
27454        if !self.uses_sliding_gated_moe_program()
27455            || !crate::tp::step_tp_graph_enabled()?
27456            || !crate::tp::step_tp_dcw_enabled()?
27457            || !crate::tp::step_tp_qkv_fused_enabled()?
27458            || !crate::tp::step_tp_dev_router_enabled()?
27459            || !crate::tp::step_nvfp4_dev_routes_enabled()?
27460        {
27461            return Ok(None);
27462        }
27463        // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): the eager
27464        // token step is this route's byte-identical twin — warmup and rebase tokens
27465        // already ride it — so below the driver-free floor the token goes eager
27466        // (`Ok(None)` = the caller's eager fallback) instead of feeding cuGraphLaunch
27467        // an exhausted card (lane/graph-launch-guard-sweep-20260831).
27468        if !crate::spec::graph_launch_headroom_ok(e) {
27469            static NOTED: std::sync::Once = std::sync::Once::new();
27470            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
27471            return Ok(None);
27472        }
27473        let n_embd = self.cfg.n_embd as usize;
27474        let n_vocab = self.cfg.n_vocab as usize;
27475        let n_layers = self.layers.len();
27476        let pos = cache.pos;
27477        let staged_next = pos + 1;
27478        if staged_next < 96 {
27479            return Ok(None); // sub-vec-floor contexts keep eager (fa kernel-class boundary)
27480        }
27481
27482        // Per-layer eligibility: every TP layer contiguous-appends this token (any rebase ->
27483        // eager fallback for the whole token; the host path also updates base_d there).
27484        for il in 0..n_layers {
27485            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
27486                return Ok(None); // caches not hydrated yet — eager warms them
27487            };
27488            if tp_kv.peek_append_ring(1)?.1 {
27489                return Ok(None);
27490            }
27491        }
27492
27493        // Bucket key: the global layers' fa split geometry at this depth (SWA layers cap at
27494        // their window and share one bucket forever after ctx > window).
27495        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
27496        if !fa_vec {
27497            return Ok(None);
27498        }
27499        let sp = crate::fa_split_keys(staged_next, 8);
27500        let bucket_max = (n_splits * sp).max(staged_next);
27501
27502        let mut state_guard = self
27503            .step35_token_graph
27504            .lock()
27505            .map_err(|_| "step35 token graph lock is poisoned")?;
27506        if state_guard.is_none() {
27507            let _main = e.gpu.enter_main()?;
27508            let n_expert = self
27509                .cfg
27510                .moe
27511                .as_ref()
27512                .map(|m| m.expert_count as usize)
27513                .unwrap_or(0);
27514            let n_ff_sh = self
27515                .layers
27516                .iter()
27517                .find_map(|l| match &l.ffn {
27518                    crate::hybrid::Ffn::Moe(m) => m.gate_shexp.as_ref().map(|g| g.out_features()),
27519                    _ => None,
27520                })
27521                .unwrap_or(0);
27522            let n_ff_dense = self
27523                .layers
27524                .iter()
27525                .find_map(|l| match &l.ffn {
27526                    crate::hybrid::Ffn::Dense { ffn_gate, .. } => Some(ffn_gate.out_features()),
27527                    _ => None,
27528                })
27529                .unwrap_or(0);
27530            *state_guard = Some(Step35TokenGraphState {
27531                graphs: Vec::new(),
27532                token_d: e.stream().clone_htod(&[0u32])?,
27533                pos_d: e.htod_i32(&[pos as i32])?,
27534                logits_stage: e.htod(&vec![0.0f32; n_vocab])?,
27535                x: e.htod(&vec![0.0f32; n_embd])?,
27536                x1: e.htod(&vec![0.0f32; n_embd])?,
27537                mixed_stage: e.htod(&vec![0.0f32; n_embd])?,
27538                sh_stage: e.htod(&vec![0.0f32; n_embd])?,
27539                k_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
27540                v_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
27541                router_logits: e.htod(&vec![0.0f32; n_expert.max(1)])?,
27542                shexp_gate: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
27543                shexp_up: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
27544                shexp_act: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
27545                gate_sig: e.htod(&[1.0f32; 1])?,
27546                dense_z: e.htod(&vec![0.0f32; n_embd])?,
27547                dense_gate: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
27548                dense_up: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
27549                dense_act: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
27550                hn: e.htod(&vec![0.0f32; n_embd])?,
27551                probe_mixed: e.htod(&vec![0.0f32; n_embd])?,
27552                probe_x: e.htod(&vec![0.0f32; n_embd])?,
27553                token_hist: e.stream().clone_htod(&[0u32; 16])?,
27554                hist_idx: e.htod_i32(&[0])?,
27555            });
27556        }
27557        let state = state_guard.as_mut().expect("state armed above");
27558        // Pre-arm the argmax partials OUTSIDE any capture: the launcher allocates them on
27559        // first use, and an alloc inside a captured section is a mem node (child graphs
27560        // reject those — the tail argmax chain needs them already resident).
27561        {
27562            let _main = e.gpu.enter_main()?;
27563            let Step35TokenGraphState {
27564                logits_stage,
27565                token_d,
27566                ..
27567            } = &mut *state;
27568            e.argmax_token_device_into(logits_stage, token_d, n_vocab)?;
27569        }
27570
27571        // ONE graph, retargeted per bucket (increment C): the per-16-token whole rebuild was
27572        // ~55ms (3.4ms/token persistent); the M1 exec update path moves nsp/ski/gridDimY and
27573        // the partial-pool memset widths in ~1ms. The partial pool is pre-grown to the run
27574        // ceiling at build so the baked pointers never move.
27575        if state.graphs.is_empty() {
27576            // Build the parent at this bucket. Capture executes nothing; correctness is
27577            // pinned at replay by the token-identity gate.
27578            self.step35_token_graph_build(e, cache, state, bucket_max)?;
27579        }
27580        {
27581            let (b, g) = state.graphs.first_mut().expect("graph built above");
27582            if *b != bucket_max {
27583                g.retarget_bucket(bucket_max)?;
27584                *b = bucket_max;
27585            }
27586        }
27587        let graph = state
27588            .graphs
27589            .first()
27590            .map(|(_, g)| g)
27591            .expect("graph built above");
27592
27593        let tg_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
27594        let t_fence = tg_timing.then(std::time::Instant::now);
27595        // Rank-stream fence: an eager token (warmup, rebase fallback) leaves len-mirror sets
27596        // queued on the rank streams, and graph children carry no ordering edge to those
27597        // streams. Graph tokens themselves enqueue nothing there (external commit), so this
27598        // sync is a no-op between consecutive replays.
27599        {
27600            let fa0 = match &self.layers[0].mixer {
27601                Mixer::Full(fa) => fa,
27602                _ => return Err("step35 token graph expects full-attention layers".into()),
27603            };
27604            let tp0 = fa0
27605                .step_tp_qkv
27606                .as_ref()
27607                .ok_or("step35 token graph lost its TP state")?;
27608            for rank in 0..tp0.runtime.devices().len() {
27609                let engine = tp0
27610                    .runtime
27611                    .rank_engine(rank)
27612                    .ok_or("step35 token graph lost a rank engine")?;
27613                let _main = engine.gpu.enter_main()?;
27614                engine.stream().synchronize()?;
27615            }
27616        }
27617
27618        // Replay: feed the token, launch, read the logits, mirror the host bookkeeping.
27619        {
27620            let _main = e.gpu.enter_main()?;
27621            e.set_u32_one(&mut state.token_d, token)?;
27622            e.set_i32_one(&mut state.pos_d, pos as i32)?;
27623        }
27624        let t_launch = tg_timing.then(std::time::Instant::now);
27625        graph.launch(e)?;
27626        let t_book = tg_timing.then(std::time::Instant::now);
27627        // Host bookkeeping OVERLAPS the replay (nsys 2026-08-21: the 45-layer txn loop was
27628        // 5.4ms/token of inter-token gap when it ran after the sync). Host-only work except
27629        // the local len_d set, which is stream-ordered AFTER the graph on e's stream. On a
27630        // replay error the counters are already advanced — acceptable: the decode aborts.
27631        for il in 0..n_layers {
27632            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
27633            let transaction = tp_kv.begin_transaction()?;
27634            let fa = match &self.layers[il].mixer {
27635                Mixer::Full(fa) => fa,
27636                _ => return Err("step35 token graph expects full-attention layers".into()),
27637            };
27638            let tp = fa
27639                .step_tp_qkv
27640                .as_ref()
27641                .ok_or("step35 token graph lost its TP state")?;
27642            // Bookkeeping-only txn (external appends; mirror sets skipped — the in-graph
27643            // incs own the counters). Shards unused.
27644            let empty: [CudaSlice<f32>; 0] = [];
27645            tp.runtime.append_tp_kv_transaction_inner(
27646                tp_kv,
27647                transaction,
27648                &empty,
27649                &empty,
27650                1,
27651                true,
27652            )?;
27653            tp.runtime
27654                .commit_tp_kv_transaction_external(tp_kv, transaction, 1)?;
27655            // Local shadow: lengths advance (v1 keeps contents stale under the door).
27656            if let Some(local) = cache.kv[il].as_mut() {
27657                local.len = pos + 1;
27658                let _main = e.gpu.enter_main()?;
27659                e.set_i32_one(&mut local.len_d, (pos + 1) as i32)?;
27660            }
27661        }
27662        cache.pos = pos + 1;
27663        let t_sync = tg_timing.then(std::time::Instant::now);
27664        let (logits, h_seed) = {
27665            let _main = e.gpu.enter_main()?;
27666            e.stream().synchronize()?;
27667            (e.dtoh(&state.logits_stage)?, e.clone_dtod(&state.x)?)
27668        };
27669        if let (Some(f), Some(l), Some(b), Some(sy)) = (t_fence, t_launch, t_book, t_sync) {
27670            use std::sync::atomic::{AtomicU64, Ordering};
27671            static NS: [AtomicU64; 5] = [
27672                AtomicU64::new(0),
27673                AtomicU64::new(0),
27674                AtomicU64::new(0),
27675                AtomicU64::new(0),
27676                AtomicU64::new(0),
27677            ];
27678            static CALLS: AtomicU64 = AtomicU64::new(0);
27679            let now = std::time::Instant::now();
27680            NS[0].fetch_add((l - f).as_nanos() as u64, Ordering::Relaxed); // fence+set
27681            NS[1].fetch_add((b - l).as_nanos() as u64, Ordering::Relaxed); // launch call
27682            NS[2].fetch_add((sy - b).as_nanos() as u64, Ordering::Relaxed); // bookkeeping
27683            NS[3].fetch_add((now - sy).as_nanos() as u64, Ordering::Relaxed); // sync+dtoh
27684            NS[4].fetch_add((now - f).as_nanos() as u64, Ordering::Relaxed); // total
27685            let calls = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
27686            if calls.is_multiple_of(100) {
27687                let avg = |i: usize| NS[i].load(Ordering::Relaxed) as f64 / calls as f64 / 1e3;
27688                eprintln!(
27689                    "[tg-timing] calls={calls} fence_us={:.0} launch_us={:.0} book_us={:.0} \
27690                     syncdtoh_us={:.0} total_us={:.0}",
27691                    avg(0),
27692                    avg(1),
27693                    avg(2),
27694                    avg(3),
27695                    avg(4)
27696                );
27697            }
27698        }
27699        // MEMRA_TG_PROBE_LAYER diagnostics: append the captured layer-K probes.
27700        if std::env::var("MEMRA_TG_PROBE_LAYER").is_ok() {
27701            use std::io::Write;
27702            let (pm, px) = {
27703                let _main = e.gpu.enter_main()?;
27704                (e.dtoh(&state.probe_mixed)?, e.dtoh(&state.probe_x)?)
27705            };
27706            for (path, data) in [
27707                ("/root/tg-probe-mixed.bin", &pm),
27708                ("/root/tg-probe-x.bin", &px),
27709            ] {
27710                let mut fo = std::fs::OpenOptions::new()
27711                    .create(true)
27712                    .append(true)
27713                    .open(path)?;
27714                for v in data {
27715                    fo.write_all(&v.to_le_bytes())?;
27716                }
27717            }
27718        }
27719        // MEMRA_DUMP_HN twin of the eager tail's dump (same format: appended raw LE f32 rows)
27720        // so a graph arm and an eager arm produce position-aligned pre-head hidden streams.
27721        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
27722            let hh = {
27723                let _main = e.gpu.enter_main()?;
27724                e.dtoh(&state.hn)?
27725            };
27726            use std::io::Write;
27727            let mut fo = std::fs::OpenOptions::new()
27728                .create(true)
27729                .append(true)
27730                .open(path)?;
27731            for v in &hh {
27732                fo.write_all(&v.to_le_bytes())?;
27733            }
27734        }
27735        // MEMRA_STEP_TP_GRAPH_DEBUG=1: per-token device-counter dump (drift hunts). One dtoh
27736        // per rank per token; diagnostics only.
27737        if std::env::var("MEMRA_STEP_TP_GRAPH_DEBUG").as_deref() == Ok("1") {
27738            for il in [0usize, 1, 44] {
27739                let tp_kv = cache.tp_kv[il].as_ref().expect("eligibility checked above");
27740                let host_len = tp_kv.staged_len();
27741                let fa = match &self.layers[il].mixer {
27742                    Mixer::Full(fa) => fa,
27743                    _ => continue,
27744                };
27745                let tp = fa
27746                    .step_tp_qkv
27747                    .as_ref()
27748                    .ok_or("step35 token graph lost its TP state")?;
27749                for rank in 0..tp.runtime.devices().len() {
27750                    let engine = tp
27751                        .runtime
27752                        .rank_engine(rank)
27753                        .ok_or("step35 token graph lost a rank engine")?;
27754                    let rank_cache = tp_kv.rank(rank).ok_or("debug rank cache missing")?;
27755                    let _main = engine.gpu.enter_main()?;
27756                    engine.stream().synchronize()?;
27757                    let len_d = engine.dtoh_i32_one(rank_cache.len_d())?;
27758                    let base_d = match rank_cache.base_d() {
27759                        Some(b) => engine.dtoh_i32_one(b)?,
27760                        None => -1,
27761                    };
27762                    eprintln!(
27763                        "[graph-debug] pos={pos} il={il} rank={rank} host_len={host_len} \
27764                         len_d={len_d} base_d={base_d}"
27765                    );
27766                }
27767            }
27768        }
27769        Ok(Some((logits, h_seed)))
27770    }
27771
27772    /// MEMRA_HEAD_SPLIT worker: dev0 computes lm-head rows [0, half), rank1 computes
27773    /// [half, n_vocab) concurrently on its otherwise-idle tail, and the halves concatenate
27774    /// on e — bit-identical per logit to the single-device matvec. Process-static workspace
27775    /// (SHEXP_WS pattern) pinned by the head tensor pointer; rank1 holds a one-time 0.5GB
27776    /// replica of its row half. Returns None when ineligible (no bf16 head / no rank1).
27777    pub(crate) fn head_split_matvec(
27778        &self,
27779        e: &Engine,
27780        hn: &CudaSlice<f32>,
27781    ) -> Result<Option<Vec<f32>>, Box<dyn std::error::Error>> {
27782        if self.head_split_fill_device(e, hn)?.is_none() {
27783            return Ok(None);
27784        }
27785        let guard = HEAD_SPLIT_WS
27786            .lock()
27787            .map_err(|_| "head split lock is poisoned")?;
27788        let ws = guard.as_ref().expect("filled above");
27789        let _main = e.gpu.enter_main()?;
27790        Ok(Some(e.dtoh(&ws.logits_e)?))
27791    }
27792
27793    /// Compute body of the split head: arms the replica + staging on first use, then fills
27794    /// the persistent full-logits row (e's half by view matvec, rank1's half by raw P2P
27795    /// push) and orders e's stream behind it. None = ineligible.
27796    fn head_split_fill_device(
27797        &self,
27798        e: &Engine,
27799        hn: &CudaSlice<f32>,
27800    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
27801        use cudarc::driver::DevicePtr;
27802        let crate::model::GpuTensor::FloatBf16 { data: head, .. } = &self.output else {
27803            return Ok(None);
27804        };
27805        let Some(rank1) = self.layers.first().and_then(|l| match &l.mixer {
27806            Mixer::Full(fa) => fa
27807                .step_tp_qkv
27808                .as_ref()
27809                .and_then(|tp| tp.runtime.rank_engine(1)),
27810            _ => None,
27811        }) else {
27812            return Ok(None);
27813        };
27814        let n_embd = self.cfg.n_embd as usize;
27815        let n_vocab = self.cfg.n_vocab as usize;
27816        let half = n_vocab / 2;
27817        let mut guard = HEAD_SPLIT_WS
27818            .lock()
27819            .map_err(|_| "head split lock is poisoned")?;
27820        let pin = {
27821            let _main = e.gpu.enter_main()?;
27822            let stream = e.stream();
27823            let (ptr, _g) = head.device_ptr(&stream);
27824            ptr
27825        };
27826        if guard.as_ref().is_none_or(|ws| ws.pin != pin) {
27827            // One-time: upload rank1's row half + persistent staging.
27828            let hi_rows = n_vocab - half;
27829            let (w1, hn1, y1, ev_done) = {
27830                let _r1 = rank1.gpu.enter_main()?;
27831                (
27832                    rank1.alloc_u8_uninit(hi_rows * n_embd * 2)?,
27833                    rank1.htod(&vec![0.0f32; n_embd])?,
27834                    rank1.htod(&vec![0.0f32; hi_rows])?,
27835                    rank1.ctx().new_event(None)?,
27836                )
27837            };
27838            {
27839                use cudarc::driver::sys;
27840                let src = pin + (half * n_embd * 2) as u64;
27841                let dst = {
27842                    let _r1 = rank1.gpu.enter_main()?;
27843                    let rstream = rank1.stream();
27844                    let (d, _g) = w1.device_ptr(&rstream);
27845                    d
27846                };
27847                let _r1 = rank1.gpu.enter_main()?;
27848                let r = unsafe {
27849                    sys::cuMemcpyAsync(
27850                        dst as sys::CUdeviceptr,
27851                        src as sys::CUdeviceptr,
27852                        hi_rows * n_embd * 2,
27853                        rank1.stream().cu_stream() as sys::CUstream,
27854                    )
27855                };
27856                if r != sys::CUresult::CUDA_SUCCESS {
27857                    return Err(format!("head split replica upload: {r:?}").into());
27858                }
27859                rank1.stream().synchronize()?;
27860            }
27861            let (logits_e, ev_hn) = {
27862                let _main = e.gpu.enter_main()?;
27863                (e.htod(&vec![0.0f32; n_vocab])?, e.ctx().new_event(None)?)
27864            };
27865            let (raw_hn1, raw_y1) = {
27866                let _r1 = rank1.gpu.enter_main()?;
27867                let rstream = rank1.stream();
27868                let (a, _g0) = hn1.device_ptr(&rstream);
27869                let (b, _g1) = y1.device_ptr(&rstream);
27870                (a, b)
27871            };
27872            let raw_logits_hi = {
27873                let _main = e.gpu.enter_main()?;
27874                let stream = e.stream();
27875                let (l, _g) = logits_e.device_ptr(&stream);
27876                l + (half * 4) as u64
27877            };
27878            *guard = Some(HeadSplit {
27879                pin,
27880                w1,
27881                hn1,
27882                y1,
27883                logits_e,
27884                ev_hn,
27885                ev_done,
27886                raw_hn1,
27887                raw_y1,
27888                raw_logits_hi,
27889                samp: None,
27890            });
27891        }
27892        let ws = guard.as_mut().expect("armed above");
27893        let hi_rows = n_vocab - half;
27894        // e: signal hn ready; rank1: pull hn, matvec its half, push the logits half back.
27895        let raw_hn = {
27896            let _main = e.gpu.enter_main()?;
27897            let stream = e.stream();
27898            let (h, _g) = hn.device_ptr(&stream);
27899            ws.ev_hn.record(&stream)?;
27900            h
27901        };
27902        {
27903            let _r1 = rank1.gpu.enter_main()?;
27904            rank1.stream().wait(&ws.ev_hn)?;
27905            crate::tp::raw_copy_bytes(ws.raw_hn1, raw_hn, n_embd * 4, rank1)?;
27906            let HeadSplit { w1, hn1, y1, .. } = &mut *ws;
27907            rank1.matvec_bf16_into(w1, hn1, y1, n_embd, hi_rows)?;
27908            crate::tp::raw_copy_bytes(ws.raw_logits_hi, ws.raw_y1, hi_rows * 4, rank1)?;
27909            ws.ev_done.record(&rank1.stream())?;
27910        }
27911        {
27912            let _main = e.gpu.enter_main()?;
27913            let head_lo = head.slice(0..half * n_embd * 2);
27914            let HeadSplit { logits_e, .. } = &mut *ws;
27915            // Writes rows [0, half) of logits_e; rank1's raw push fills [half, n_vocab).
27916            e.matvec_bf16_view_into(&head_lo, hn, logits_e, n_embd, half)?;
27917            e.stream().wait(&ws.ev_done)?;
27918            Ok(Some(()))
27919        }
27920    }
27921
27922    /// Device twin of `head_split_matvec` for the chain: fills the persistent full-logits
27923    /// row exactly like the host variant (identical halves, identical concat) and runs the
27924    /// device argmax into `token_d` — NO host readback. Returns false when the split is
27925    /// ineligible (caller falls back to the plain matmul head).
27926    pub(crate) fn head_split_argmax_device(
27927        &self,
27928        e: &Engine,
27929        hn: &CudaSlice<f32>,
27930        token_d: &mut CudaSlice<u32>,
27931    ) -> Result<bool, Box<dyn std::error::Error>> {
27932        if self.head_split_fill_device(e, hn)?.is_none() {
27933            return Ok(false);
27934        }
27935        let n_vocab = self.cfg.n_vocab as usize;
27936        let guard = HEAD_SPLIT_WS
27937            .lock()
27938            .map_err(|_| "head split lock is poisoned")?;
27939        let ws = guard.as_ref().expect("filled above");
27940        let _main = e.gpu.enter_main()?;
27941        e.argmax_token_device_into(&ws.logits_e, token_d, n_vocab)?;
27942        Ok(true)
27943    }
27944
27945    /// SAMPLED twin of `head_split_argmax_device`. The split head already materializes the
27946    /// full concatenated row in `ws.logits_e`, so sampling does NOT have to give up HEAD_SPLIT
27947    /// — it draws from that row on device (filter thresholds, Gumbel perturbation, argmax)
27948    /// exactly as the serve tick does. Worth ~0.2 ms/token: the post-W8 census had the
27949    /// unsplit q8 head at ~364 us against ~82 us per half.
27950    pub(crate) fn head_split_sample_device(
27951        &self,
27952        e: &Engine,
27953        hn: &CudaSlice<f32>,
27954        token_d: &mut CudaSlice<u32>,
27955        samp: &crate::decode_batch::DevSamp,
27956        ctr: u32,
27957    ) -> Result<bool, Box<dyn std::error::Error>> {
27958        if self.head_split_fill_device(e, hn)?.is_none() {
27959            return Ok(false);
27960        }
27961        let n_vocab = self.cfg.n_vocab as usize;
27962        let guard = HEAD_SPLIT_WS
27963            .lock()
27964            .map_err(|_| "head split lock is poisoned")?;
27965        let mut guard = guard;
27966        let ws = guard.as_mut().expect("filled above");
27967        let _main = e.gpu.enter_main()?;
27968        if ws.samp.is_none() {
27969            ws.samp = Some(SampScratch {
27970                pb: e.zeros(n_vocab)?,
27971                th: e.zeros(1)?,
27972                z: e.zeros(1)?,
27973                mx: e.zeros(1)?,
27974                rows: e.htod_i32(&[0i32])?,
27975            });
27976        }
27977        let filtered = samp.top_k > 0 || samp.top_p < 1.0 || samp.min_p > 0.0;
27978        let HeadSplit {
27979            logits_e,
27980            samp: scratch,
27981            ..
27982        } = &mut *ws;
27983        let sc = scratch.as_mut().expect("armed above");
27984        if filtered {
27985            e.filter_stats(
27986                logits_e, n_vocab, &sc.rows, &mut sc.th, &mut sc.z, &mut sc.mx, n_vocab, 1,
27987                samp.temp, samp.top_k, samp.top_p, samp.min_p,
27988            )?;
27989            let SampScratch { pb, th, mx, .. } = sc;
27990            e.gumbel_perturb_filtered_col(
27991                logits_e, 0, pb, n_vocab, samp.seed, ctr, samp.temp, mx, th, 0,
27992            )?;
27993        } else {
27994            e.gumbel_perturb_col(logits_e, 0, &mut sc.pb, n_vocab, samp.seed, ctr, samp.temp)?;
27995        }
27996        e.argmax_token_device_col(&sc.pb, 0, n_vocab, token_d, 0)?;
27997        Ok(true)
27998    }
27999
28000    /// One-per-chunk host readback of the persistent full-logits row (the LAST chain
28001    /// token's row).
28002    pub(crate) fn head_split_logits_dtoh(
28003        &self,
28004        e: &Engine,
28005    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
28006        let guard = HEAD_SPLIT_WS
28007            .lock()
28008            .map_err(|_| "head split lock is poisoned")?;
28009        let ws = guard.as_ref().ok_or("head split logits not armed")?;
28010        let _main = e.gpu.enter_main()?;
28011        e.dtoh(&ws.logits_e)
28012    }
28013
28014    /// Chunk-loop replay (F-lite, MEMRA_STEP_TP_GRAPH_LOOP): run up to `k_target` greedy
28015    /// tokens as back-to-back graph launches chained through the in-graph tail argmax —
28016    /// ONE host sync, ONE history readback, and ONE bulk KV transaction per layer per
28017    /// chunk. Returns None when ineligible (caller falls back to the per-token path).
28018    /// The chunk consumes `token` (already emitted by the caller) as launch 0's input and
28019    /// returns the ids the chain argmax'd (hist[0..k]) plus the LAST launch's logits row —
28020    /// hist[k-1] is exactly argmax(logits), so the caller emits hist[..k-1] and lets its
28021    /// own loop re-derive hist[k-1] from the returned row.
28022    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
28023    pub fn step35_token_graph_chunk(
28024        &self,
28025        e: &Engine,
28026        token: u32,
28027        k_target: usize,
28028        cache: &mut Cache,
28029    ) -> Result<Option<(Vec<u32>, Vec<f32>)>, Box<dyn std::error::Error>> {
28030        if !self.uses_sliding_gated_moe_program()
28031            || !crate::tp::step_tp_graph_enabled()?
28032            || !crate::tp::step_tp_dcw_enabled()?
28033            || !crate::tp::step_tp_qkv_fused_enabled()?
28034            || !crate::tp::step_tp_dev_router_enabled()?
28035            || !crate::tp::step_nvfp4_dev_routes_enabled()?
28036        {
28037            return Ok(None);
28038        }
28039        // GRAPH-LAUNCH HEADROOM GUARD: same guard, same eager twin as
28040        // `step35_token_graph_step` (the chunk is that step replayed k times).
28041        if !crate::spec::graph_launch_headroom_ok(e) {
28042            static NOTED: std::sync::Once = std::sync::Once::new();
28043            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
28044            return Ok(None);
28045        }
28046        let n_layers = self.layers.len();
28047        let pos = cache.pos;
28048        let staged_next = pos + 1;
28049        if staged_next < 96 {
28050            return Ok(None);
28051        }
28052        // Bucket for the FIRST token; the chunk must not cross the bucket boundary (the
28053        // exec's n_splits ladder must match eager per depth).
28054        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
28055        if !fa_vec {
28056            return Ok(None);
28057        }
28058        let sp = crate::fa_split_keys(staged_next, 8);
28059        let bucket_max = (n_splits * sp).max(staged_next);
28060        let to_boundary = bucket_max.saturating_sub(staged_next) + 1;
28061        let mut k = k_target.min(to_boundary).min(16);
28062        if k < 2 {
28063            return Ok(None);
28064        }
28065        // Every layer must contiguous-append all k rows (no rebase inside the chunk).
28066        for il in 0..n_layers {
28067            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
28068                return Ok(None);
28069            };
28070            while k >= 2 && tp_kv.peek_append_ring(k)?.1 {
28071                k -= 1;
28072            }
28073            if k < 2 {
28074                return Ok(None);
28075            }
28076        }
28077
28078        let mut state_guard = self
28079            .step35_token_graph
28080            .lock()
28081            .map_err(|_| "step35 token graph lock is poisoned")?;
28082        let Some(state) = state_guard.as_mut() else {
28083            return Ok(None); // per-token path arms the state + stages first
28084        };
28085        if state.graphs.is_empty() {
28086            return Ok(None);
28087        }
28088        {
28089            let (b, g) = state.graphs.first_mut().expect("checked above");
28090            if *b != bucket_max {
28091                g.retarget_bucket(bucket_max)?;
28092                *b = bucket_max;
28093            }
28094        }
28095        let graph = state.graphs.first().map(|(_, g)| g).expect("checked above");
28096
28097        // Rank-stream fence (eager stragglers; see the per-token path).
28098        {
28099            let fa0 = match &self.layers[0].mixer {
28100                Mixer::Full(fa) => fa,
28101                _ => return Err("step35 token graph expects full-attention layers".into()),
28102            };
28103            let tp0 = fa0
28104                .step_tp_qkv
28105                .as_ref()
28106                .ok_or("step35 token graph lost its TP state")?;
28107            for rank in 0..tp0.runtime.devices().len() {
28108                let engine = tp0
28109                    .runtime
28110                    .rank_engine(rank)
28111                    .ok_or("step35 token graph lost a rank engine")?;
28112                let _main = engine.gpu.enter_main()?;
28113                engine.stream().synchronize()?;
28114            }
28115        }
28116
28117        // Seed the chain and fire k launches back-to-back: launch i embeds the token the
28118        // PREVIOUS launch's tail argmax wrote (launch 0 embeds the host-seeded `token`).
28119        {
28120            let _main = e.gpu.enter_main()?;
28121            e.set_u32_one(&mut state.token_d, token)?;
28122            e.set_i32_one(&mut state.pos_d, pos as i32)?;
28123            e.set_i32_one(&mut state.hist_idx, 0)?;
28124        }
28125        for _ in 0..k {
28126            graph.launch(e)?;
28127        }
28128        // Bulk host bookkeeping overlaps the replays: one k-row txn per layer.
28129        for il in 0..n_layers {
28130            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
28131            let transaction = tp_kv.begin_transaction()?;
28132            let fa = match &self.layers[il].mixer {
28133                Mixer::Full(fa) => fa,
28134                _ => return Err("step35 token graph expects full-attention layers".into()),
28135            };
28136            let tp = fa
28137                .step_tp_qkv
28138                .as_ref()
28139                .ok_or("step35 token graph lost its TP state")?;
28140            let empty: [CudaSlice<f32>; 0] = [];
28141            tp.runtime.append_tp_kv_transaction_inner(
28142                tp_kv,
28143                transaction,
28144                &empty,
28145                &empty,
28146                k,
28147                true,
28148            )?;
28149            tp.runtime
28150                .commit_tp_kv_transaction_external(tp_kv, transaction, k)?;
28151            if let Some(local) = cache.kv[il].as_mut() {
28152                local.len = pos + k;
28153                let _main = e.gpu.enter_main()?;
28154                e.set_i32_one(&mut local.len_d, (pos + k) as i32)?;
28155            }
28156        }
28157        cache.pos = pos + k;
28158        let (hist, logits) = {
28159            let _main = e.gpu.enter_main()?;
28160            e.stream().synchronize()?;
28161            (e.dtoh_u32(&state.token_hist)?, e.dtoh(&state.logits_stage)?)
28162        };
28163        Ok(Some((hist[..k].to_vec(), logits)))
28164    }
28165}
28166
28167impl HybridModel {
28168    /// Capture the whole-token parent for one fa bucket. Capture executes nothing; the
28169    /// section closures issue the SAME calls the eager dcw path runs (bit-proven), with the
28170    /// per-token operands living at the fixed stage addresses. Groups: the two rank sections
28171    /// of each phase fork in parallel and merge into the following root section.
28172    #[allow(clippy::too_many_arguments)]
28173    fn step35_token_graph_build(
28174        &self,
28175        e: &Engine,
28176        cache: &mut Cache,
28177        state: &mut Step35TokenGraphState,
28178        bucket_max: usize,
28179    ) -> Result<(), Box<dyn std::error::Error>> {
28180        use cudarc::driver::DevicePtr;
28181        let n_embd = self.cfg.n_embd as usize;
28182        let eps = self.cfg.rms_eps;
28183        let n_layers = self.layers.len();
28184        let started = std::time::Instant::now();
28185        if !crate::router_kernel_on() {
28186            return Err(
28187                "step35 token graph requires the router kernel (MEMRA_ROUTER_KERNEL=0)".into(),
28188            );
28189        }
28190        if !Engine::bf16_mmv_on() || !n_embd.is_multiple_of(8) {
28191            return Err("step35 token graph requires MEMRA_BF16_MMV bf16-resident matvecs".into());
28192        }
28193
28194        // Device embed table (the spec/graph lanes' lazily-uploaded copy).
28195        let embd_gpu = self
28196            .embd_gpu_try(e)
28197            .ok_or("step35 token graph could not upload the device embed table")?;
28198        let embd_qtype = match self.embd.ggml_type {
28199            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
28200            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
28201            other => return Err(format!("token graph embed dtype {other:?} unhandled").into()),
28202        };
28203        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
28204
28205        // Fixed-stage pointers the sections reference.
28206        let (p_mixed, p_kshadow, p_vshadow) = {
28207            let _main = e.gpu.enter_main()?;
28208            let stream = e.stream();
28209            let (a, _g) = state.mixed_stage.device_ptr(&stream);
28210            let (b, _g) = state.k_shadow_stage.device_ptr(&stream);
28211            let (c, _g) = state.v_shadow_stage.device_ptr(&stream);
28212            (a, b, c)
28213        };
28214
28215        crate::tp::token_graph_build_begin()?;
28216        let mut group_id: u32 = 0;
28217        for il in 0..n_layers {
28218            let layer = &self.layers[il];
28219            let fa = match &layer.mixer {
28220                Mixer::Full(fa) => fa,
28221                _ => return Err("step35 token graph expects full-attention layers".into()),
28222            };
28223            let tp = fa
28224                .step_tp_qkv
28225                .as_ref()
28226                .ok_or("step35 token graph lost its TP state")?;
28227            let attention = tp
28228                .attention
28229                .as_ref()
28230                .ok_or("step35 token graph lost its attention aux")?;
28231            let geometry = self.step35_geom(il);
28232            let window = geometry.window.map(|w| w as usize);
28233            let head_dim = geometry.head_dim_k as usize;
28234            let heads = geometry.n_head as usize;
28235            let kv_heads = geometry.n_head_kv as usize;
28236            let ranks = tp.runtime.devices().len();
28237            let local_heads = heads / ranks;
28238            let local_kv_heads = kv_heads / ranks;
28239            let layer_bucket = window.map(|w| bucket_max.min(w)).unwrap_or(bucket_max);
28240            let use_gate_shards =
28241                attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some();
28242            if !use_gate_shards {
28243                return Err("step35 token graph requires the fused gate shards".into());
28244            }
28245
28246            let ws_index = tp
28247                .runtime
28248                .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
28249            let ws_mutex = tp.runtime.decode_v2_workspace();
28250            let mut ws_guard = ws_mutex
28251                .lock()
28252                .map_err(|_| "step TP decode v2 workspace lock is poisoned")?;
28253            let ws = ws_guard
28254                .get_mut(ws_index)
28255                .ok_or("step TP decode v2 workspace missing after ensure")?;
28256            tp.runtime
28257                .decode_v2_arm_token_mirrors(ws, p_mixed, (p_kshadow, p_vshadow))?;
28258            let mut rope_freqs = Vec::with_capacity(ranks);
28259            for rank in 0..ranks {
28260                let engine = tp
28261                    .runtime
28262                    .rank_engine(rank)
28263                    .ok_or("step35 token graph lost a rank engine")?;
28264                rope_freqs.push(if geometry.rope_factors {
28265                    self.step35_aux
28266                        .as_ref()
28267                        .and_then(|aux| aux.rope_freqs(engine))
28268                } else {
28269                    None
28270                });
28271            }
28272            let gate_shards_arg = if let Some(shards) = attention.gate_shards.as_deref() {
28273                Some(crate::tp::StepTpGateShards::F32(shards))
28274            } else {
28275                attention
28276                    .gate_shards_bf16
28277                    .as_deref()
28278                    .map(crate::tp::StepTpGateShards::Bf16)
28279            };
28280
28281            // ---- E1: embed (layer 0) / attn norm into h_stage + pos copy ----
28282            let decode_input = attention
28283                .decode_input
28284                .as_ref()
28285                .ok_or("step35 token graph requires the replicated decode input")?;
28286            let mut decode_input = decode_input
28287                .lock()
28288                .map_err(|_| "replicated decode input lock is poisoned")?;
28289            // Stage arming happens through the eager stage flow once; require it here.
28290            if ws.h_stage.is_none() {
28291                return Err(
28292                    "step35 token graph requires the stage flow armed (run eager dcw first)".into(),
28293                );
28294            }
28295            {
28296                let state_x = &mut state.x;
28297                let token_d = &state.token_d;
28298                let pos_d = &state.pos_d;
28299                crate::tp::graph_section(e, None, || {
28300                    let _main = e.gpu.enter_main()?;
28301                    if il == 0 {
28302                        e.embed_gather_device_into(
28303                            embd_gpu,
28304                            token_d,
28305                            state_x,
28306                            n_embd,
28307                            embd_qtype,
28308                            embd_row_bytes,
28309                        )?;
28310                    }
28311                    {
28312                        let h_stage = ws.h_stage.as_mut().expect("stage armed checked above");
28313                        e.rms_norm(
28314                            state_x,
28315                            layer.attn_norm.float_data(),
28316                            h_stage,
28317                            n_embd,
28318                            1,
28319                            eps,
28320                        )?;
28321                    }
28322                    {
28323                        let pos_stage = ws.pos_stage.as_mut().expect("stage armed above");
28324                        let mut dst = pos_stage.slice_mut(0..1);
28325                        e.stream().memcpy_dtod(&pos_d.slice(0..1), &mut dst)?;
28326                    }
28327                    Ok(())
28328                })?;
28329            }
28330
28331            // ---- R0/R1 (parallel): projections + dcw attention interior ----
28332            group_id += 1;
28333            for rank in 0..ranks {
28334                let engine = tp
28335                    .runtime
28336                    .rank_engine(rank)
28337                    .ok_or("step35 token graph lost a rank engine")?;
28338                {
28339                    // fa partial pool must reach the RUN CEILING before capture — an
28340                    // in-capture grow is a mem node (child graphs reject those), and the
28341                    // retarget path (increment C) widens the baked memsets up to the ceiling
28342                    // without moving the pool pointers. Two ensures cover both sp rungs.
28343                    let ceiling = window
28344                        .map(|w| cache.max_ctx.min(w))
28345                        .unwrap_or(cache.max_ctx);
28346                    let _main = engine.gpu.enter_main()?;
28347                    engine.fa_dcw_pool_ensure(
28348                        head_dim,
28349                        local_heads,
28350                        local_kv_heads,
28351                        ceiling.min(2048),
28352                    )?;
28353                    engine.fa_dcw_pool_ensure(head_dim, local_heads, local_kv_heads, ceiling)?;
28354                    engine.fa_dcw_pool_ensure(
28355                        head_dim,
28356                        local_heads,
28357                        local_kv_heads,
28358                        layer_bucket,
28359                    )?;
28360                }
28361                let runtime = &tp.runtime;
28362                let q_norm = &attention.q_norm;
28363                let k_norm = &attention.k_norm;
28364                let gate_ref = gate_shards_arg.as_ref();
28365                crate::tp::graph_section(engine, Some(group_id), || {
28366                    runtime.decode_v2_input_qkv_rank(
28367                        ws,
28368                        &state.pos_d,
28369                        &mut decode_input,
28370                        &tp.q,
28371                        &tp.k,
28372                        &tp.v,
28373                        q_norm,
28374                        k_norm,
28375                        head_dim,
28376                        geometry.n_rot as usize,
28377                        geometry.rope_base,
28378                        &rope_freqs,
28379                        eps,
28380                        gate_ref,
28381                        true,
28382                        true,
28383                        false,
28384                        rank,
28385                        None,
28386                    )?;
28387                    // Merged dcw interior at the BUCKET geometry (one-partition law makes the
28388                    // replayed values track the live counters).
28389                    let distributed = cache.tp_kv[il]
28390                        .as_mut()
28391                        .ok_or("step35 token graph lost a TP cache")?;
28392                    let (kv_dim_k, kv_dim_v) = (distributed.kv_dim_k(), distributed.kv_dim_v());
28393                    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
28394                    let capacity = distributed.physical_capacity();
28395                    {
28396                        let rank_cache = distributed
28397                            .rank_mut(rank)
28398                            .ok_or("step35 token graph lost a rank cache")?;
28399                        let (k_plane, v_plane, len_d, base_d) =
28400                            rank_cache.planes_and_counters_mut();
28401                        engine.append_kv_quantized_dcw(
28402                            &ws.k[rank],
28403                            &ws.v_raw[rank],
28404                            k_plane,
28405                            v_plane,
28406                            len_d,
28407                            base_d,
28408                            kv_dim_k,
28409                            kv_dim_v,
28410                            ktb,
28411                            vtb,
28412                        )?;
28413                    }
28414                    {
28415                        let rank_cache = distributed
28416                            .rank_mut(rank)
28417                            .ok_or("step35 token graph lost a rank cache")?;
28418                        engine.inc_i32(rank_cache.len_d_mut())?;
28419                    }
28420                    let rank_cache = distributed
28421                        .rank(rank)
28422                        .ok_or("step35 token graph lost a rank cache")?;
28423                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * ktb);
28424                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * vtb);
28425                    // Graph build keeps the UNFUSED combine + gate pair: the bucket
28426                    // retarget addresses combine's nsp at arg slot 6, and the fused
28427                    // combine_gate kernel shifts it to 7 — bit-identical either way, so
28428                    // only the eager arm takes FUSION #2d.
28429                    engine.fa_decode_dcw(
28430                        &ws.q[rank],
28431                        &k_ring,
28432                        &v_ring,
28433                        &mut ws.attn_out[rank],
28434                        head_dim,
28435                        local_heads,
28436                        local_kv_heads,
28437                        rank_cache.len_d(),
28438                        rank_cache.base_d(),
28439                        window.unwrap_or(0),
28440                        layer_bucket,
28441                        geometry.attention_scale(),
28442                        ktb,
28443                        vtb,
28444                        None,
28445                    )?;
28446                    engine.attn_head_gate(
28447                        &ws.attn_out[rank],
28448                        &ws.gate[rank],
28449                        &mut ws.gated[rank],
28450                        None,
28451                        head_dim,
28452                        local_heads,
28453                        1,
28454                    )?;
28455                    runtime.decode_v2_finish_rank_partial(ws, &tp.o, true, rank)?;
28456                    Ok(())
28457                })?;
28458            }
28459
28460            // ---- ROOT: combine + shadows + e-mirrors ----
28461            {
28462                let root = tp
28463                    .runtime
28464                    .rank_engine(0)
28465                    .ok_or("step35 token graph lost the root engine")?;
28466                let runtime = &tp.runtime;
28467                crate::tp::graph_section(root, None, || runtime.decode_v2_finish_root_fused(ws))?;
28468            }
28469            drop(ws_guard);
28470            drop(decode_input);
28471
28472            let probe_layer: Option<usize> = std::env::var("MEMRA_TG_PROBE_LAYER")
28473                .ok()
28474                .and_then(|v| v.parse().ok());
28475            if probe_layer == Some(il) {
28476                let Step35TokenGraphState {
28477                    mixed_stage,
28478                    probe_mixed,
28479                    ..
28480                } = &mut *state;
28481                crate::tp::graph_section(e, None, || {
28482                    let _main = e.gpu.enter_main()?;
28483                    let mut dst = probe_mixed.slice_mut(0..n_embd);
28484                    e.stream()
28485                        .memcpy_dtod(&mixed_stage.slice(0..n_embd), &mut dst)?;
28486                    Ok(())
28487                })?;
28488            }
28489
28490            // ---- FFN half ----
28491            match &layer.ffn {
28492                crate::hybrid::Ffn::Dense {
28493                    ffn_gate,
28494                    ffn_up,
28495                    ffn_down,
28496                } => {
28497                    let n_ff = ffn_gate.out_features();
28498                    let lim = self.cfg.clamp_shexp_at(il as u32);
28499                    // Alloc-free inline of ffn_swiglu_decode's bf16 tail: dual gate/up matvec is
28500                    // bit-identical per row to the two matmul-dispatched matvec_bf16 launches.
28501                    // A clamped dense layer would take eager's q8_1 branch instead -- refuse.
28502                    if lim.is_some() {
28503                        return Err("step35 token graph dense FFN with clamp unsupported".into());
28504                    }
28505                    let (wg_d, wu_d, wd_d) = match (ffn_gate, ffn_up, ffn_down) {
28506                        (
28507                            crate::model::GpuTensor::FloatBf16 { data: wg, .. },
28508                            crate::model::GpuTensor::FloatBf16 { data: wu, .. },
28509                            crate::model::GpuTensor::FloatBf16 { data: wd, .. },
28510                        ) => (wg, wu, wd),
28511                        _ => {
28512                            return Err(
28513                                "step35 token graph dense FFN requires bf16-resident weights"
28514                                    .into(),
28515                            );
28516                        }
28517                    };
28518                    crate::tp::graph_section(e, None, || {
28519                        let _main = e.gpu.enter_main()?;
28520                        let Step35TokenGraphState {
28521                            x,
28522                            x1,
28523                            mixed_stage,
28524                            dense_z,
28525                            dense_gate,
28526                            dense_up,
28527                            dense_act,
28528                            sh_stage,
28529                            ..
28530                        } = &mut *state;
28531                        e.add_rms_norm(
28532                            x,
28533                            mixed_stage,
28534                            layer.post_attn_norm.float_data(),
28535                            x1,
28536                            dense_z,
28537                            n_embd,
28538                            1,
28539                            eps,
28540                        )?;
28541                        // TWO SINGLE matvecs, not the dual: eager dense rides two
28542                        // matmul-dispatched matvec_bf16 launches; the dual twin measured a
28543                        // ~2e-9 residual difference here (token-graph bisection, 2026-08-21).
28544                        e.matvec_bf16_into(wg_d, dense_z, dense_gate, n_embd, n_ff)?;
28545                        e.matvec_bf16_into(wu_d, dense_z, dense_up, n_embd, n_ff)?;
28546                        Self::ffn_act_lim(
28547                            e, &self.cfg, dense_gate, dense_up, 1.0, 1.0, lim, dense_act, n_ff,
28548                        )?;
28549                        e.matvec_bf16_into(wd_d, dense_act, sh_stage, n_ff, n_embd)?;
28550                        e.add(x1, sh_stage, x, n_embd)?;
28551                        Ok(())
28552                    })?;
28553                }
28554                crate::hybrid::Ffn::Moe(m) => {
28555                    let moe = self
28556                        .cfg
28557                        .moe
28558                        .as_ref()
28559                        .ok_or("step35 token graph needs moe cfg")?;
28560                    let n_expert = moe.expert_count as usize;
28561                    let n_used = moe.expert_used_count as usize;
28562                    let sigmoid = self
28563                        .cfg
28564                        .sigmoid_router()
28565                        .ok_or("step35 token graph needs the sigmoid router")?;
28566                    let step_tp = m
28567                        .step_tp
28568                        .as_ref()
28569                        .ok_or("step35 token graph needs TP experts")?;
28570                    let bank = match &step_tp.experts {
28571                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => bank,
28572                        _ => return Err("step35 token graph needs the NVFP4 bank".into()),
28573                    };
28574                    let routes_ws_mutex = bank.device_workspace_handle();
28575                    let mut routes_guard = routes_ws_mutex
28576                        .lock()
28577                        .map_err(|_| "routes workspace lock is poisoned")?;
28578                    let routes_ws = routes_guard
28579                        .as_mut()
28580                        .ok_or("step35 token graph requires the routes workspace warmed")?;
28581                    routes_ws.arm_stages(e, bank.input_width, n_used)?;
28582                    step_tp.runtime.routes_arm_raw(bank, routes_ws)?;
28583                    let p_z = {
28584                        let root = step_tp
28585                            .runtime
28586                            .rank_engine(0)
28587                            .ok_or("routes root engine missing")?;
28588                        let _main = root.gpu.enter_main()?;
28589                        let stream = root.stream();
28590                        let in_stage = routes_ws
28591                            .in_stage_handle()
28592                            .ok_or("routes in stage not armed")?;
28593                        let (a, _g) = in_stage.device_ptr(&stream);
28594                        a
28595                    };
28596                    let local_out = bank.expert_width / ranks;
28597
28598                    // ---- E2: post-attn norm into the routes in-stage + router + staging ----
28599                    crate::tp::graph_section(e, None, || {
28600                        let _main = e.gpu.enter_main()?;
28601                        {
28602                            let in_stage = routes_ws
28603                                .in_stage_mut()
28604                                .ok_or("routes in stage not armed")?;
28605                            let Step35TokenGraphState {
28606                                x, x1, mixed_stage, ..
28607                            } = &mut *state;
28608                            e.add_rms_norm(
28609                                x,
28610                                mixed_stage,
28611                                layer.post_attn_norm.float_data(),
28612                                x1,
28613                                in_stage,
28614                                n_embd,
28615                                1,
28616                                eps,
28617                            )?;
28618                        }
28619                        {
28620                            let z_ref = routes_ws
28621                                .in_stage_handle()
28622                                .ok_or("routes in stage not armed")?;
28623                            e.router_gemv_into(
28624                                m.gate_inp.float_data(),
28625                                z_ref,
28626                                &mut state.router_logits,
28627                                n_embd,
28628                                n_expert,
28629                                1,
28630                            )?;
28631                        }
28632                        let (sel_e, w_e) = routes_ws
28633                            .dev_route_e_mut()
28634                            .ok_or("routes staging not armed")?;
28635                        e.moe_router_sigmoid_topk_into(
28636                            &state.router_logits,
28637                            1,
28638                            n_expert,
28639                            n_used,
28640                            m.active_count(),
28641                            &m.exp_probs_b_dev,
28642                            &m.active_experts_dev,
28643                            sigmoid.0,
28644                            sigmoid.1,
28645                            sel_e,
28646                            w_e,
28647                        )?;
28648                        Ok(())
28649                    })?;
28650
28651                    // ---- R0r/R1r (parallel): routes sweeps ----
28652                    group_id += 1;
28653                    for rank in 0..ranks {
28654                        let engine = step_tp
28655                            .runtime
28656                            .rank_engine(rank)
28657                            .ok_or("routes rank engine missing")?;
28658                        let runtime = &step_tp.runtime;
28659                        crate::tp::graph_section(engine, Some(group_id), || {
28660                            runtime.routes_rank_section(
28661                                bank,
28662                                routes_ws,
28663                                p_z,
28664                                local_out,
28665                                n_used,
28666                                step_tp.activation_limit,
28667                                rank,
28668                            )
28669                        })?;
28670                    }
28671
28672                    // ---- ROOTr: combine into the out stage ----
28673                    {
28674                        let root = step_tp
28675                            .runtime
28676                            .rank_engine(0)
28677                            .ok_or("routes root engine missing")?;
28678                        let runtime = &step_tp.runtime;
28679                        crate::tp::graph_section(root, None, || {
28680                            runtime.routes_root_section(bank, routes_ws)
28681                        })?;
28682                    }
28683
28684                    // ---- E3: shexp + add_shared onto the out stage + residual ----
28685                    // Alloc-free inline of moe_ffn_grouped_add_shared's bf16_dual arm (the arm
28686                    // eager takes under MEMRA_BF16_MMV, guarded at fn entry).
28687                    let lim_sh = self.cfg.clamp_shexp_at(il as u32);
28688                    let (wg_sh, wu_sh, wd_sh) = match (&m.gate_shexp, &m.up_shexp, &m.down_shexp) {
28689                        (
28690                            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
28691                            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
28692                            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
28693                        ) => (wg, wu, wd),
28694                        _ => {
28695                            return Err(
28696                                "step35 token graph shexp requires bf16-resident weights".into()
28697                            );
28698                        }
28699                    };
28700                    let n_ff_sh = m
28701                        .gate_shexp
28702                        .as_ref()
28703                        .expect("matched Some above")
28704                        .out_features();
28705                    // No sigmoid gate on this model's shexp -> gate_sig stays at its 1.0
28706                    // init, reproducing eager's ones vector without a launch.
28707                    let gate_inp_shexp = m.gate_inp_shexp.as_ref();
28708                    crate::tp::graph_section(e, None, || {
28709                        let _main = e.gpu.enter_main()?;
28710                        let (z_ref, out_stage) = routes_ws
28711                            .in_and_out_stages_mut()
28712                            .ok_or("routes stages not armed")?;
28713                        let Step35TokenGraphState {
28714                            x,
28715                            x1,
28716                            sh_stage,
28717                            shexp_gate,
28718                            shexp_up,
28719                            shexp_act,
28720                            gate_sig,
28721                            ..
28722                        } = &mut *state;
28723                        e.matvec_bf16_dual_into(
28724                            wg_sh, wu_sh, z_ref, shexp_gate, shexp_up, n_embd, n_ff_sh,
28725                        )?;
28726                        Self::ffn_act_lim(
28727                            e, &self.cfg, shexp_gate, shexp_up, 1.0, 1.0, lim_sh, shexp_act,
28728                            n_ff_sh,
28729                        )?;
28730                        e.matvec_bf16_into(wd_sh, shexp_act, sh_stage, n_ff_sh, n_embd)?;
28731                        if let Some(gate_w) = gate_inp_shexp {
28732                            e.sigmoid_dot_rows_into(
28733                                z_ref,
28734                                gate_w.float_data(),
28735                                gate_sig,
28736                                n_embd,
28737                                1,
28738                            )?;
28739                        }
28740                        e.add_scaled_rows(sh_stage, gate_sig, out_stage, n_embd, 1)?;
28741                        e.add(x1, out_stage, x, n_embd)?;
28742                        Ok(())
28743                    })?;
28744                }
28745            }
28746            if probe_layer == Some(il) {
28747                let Step35TokenGraphState { x, probe_x, .. } = &mut *state;
28748                crate::tp::graph_section(e, None, || {
28749                    let _main = e.gpu.enter_main()?;
28750                    let mut dst = probe_x.slice_mut(0..n_embd);
28751                    e.stream().memcpy_dtod(&x.slice(0..n_embd), &mut dst)?;
28752                    Ok(())
28753                })?;
28754            }
28755        }
28756
28757        // ---- Tail: output norm + head into the logits stage ----
28758        let head = match &self.output {
28759            crate::model::GpuTensor::FloatBf16 { data, .. } => data,
28760            _ => return Err("step35 token graph head requires the bf16-resident output".into()),
28761        };
28762        crate::tp::graph_section(e, None, || {
28763            let _main = e.gpu.enter_main()?;
28764            let Step35TokenGraphState {
28765                x,
28766                hn,
28767                logits_stage,
28768                token_d,
28769                pos_d,
28770                token_hist,
28771                hist_idx,
28772                ..
28773            } = &mut *state;
28774            e.rms_norm(x, self.output_norm.float_data(), hn, n_embd, 1, eps)?;
28775            e.matvec_bf16_into(head, hn, logits_stage, n_embd, self.cfg.n_vocab as usize)?;
28776            // Chunk-loop tail: greedy argmax feeds token_d (host-identical tie-break,
28777            // argmax_gate-validated), the id lands in the history ring, and pos advances on
28778            // device — consecutive launches chain with NO host sync. Single-token mode
28779            // overwrites token_d/pos_d from the host before each launch, so these nodes are
28780            // harmless there.
28781            e.argmax_token_device_into(logits_stage, token_d, self.cfg.n_vocab as usize)?;
28782            e.u32_hist_append(token_d, token_hist, hist_idx)?;
28783            e.inc_i32(pos_d)?;
28784            Ok(())
28785        })?;
28786
28787        let graph = crate::tp::token_graph_build_finish()?;
28788        state.graphs.push((bucket_max, graph));
28789        eprintln!(
28790            "[step35-token-graph] built bucket={bucket_max} layers={n_layers} \
28791             build_ms={:.0} performance_claim=false",
28792            started.elapsed().as_secs_f64() * 1e3
28793        );
28794        Ok(())
28795    }
28796}