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                })
281            })
282            .collect();
283        Self {
284            parent,
285            fence: fence.to_vec(),
286            stages,
287            committed: false,
288        }
289    }
290
291    fn pp2_parts(&mut self) -> (&mut Cache, &mut Cache) {
292        assert_eq!(self.stages.len(), 2);
293        let (stage0, stage1) = self.stages.split_at_mut(1);
294        (
295            stage0[0]
296                .get_mut()
297                .unwrap_or_else(|poisoned| poisoned.into_inner()),
298            stage1[0]
299                .get_mut()
300                .unwrap_or_else(|poisoned| poisoned.into_inner()),
301        )
302    }
303
304    fn stages(&self) -> &[std::sync::Mutex<Cache>] {
305        &self.stages
306    }
307
308    fn commit(&mut self) {
309        self.committed = true;
310    }
311}
312
313impl Drop for PrimeCacheStages<'_> {
314    fn drop(&mut self) {
315        let n = self.parent.kv.len();
316        for i in 0..n {
317            let stage = prime_cache_stage_for_layer(&self.fence, i);
318            let source = self.stages[stage]
319                .get_mut()
320                .unwrap_or_else(|poisoned| poisoned.into_inner());
321            debug_assert!(self.parent.kv[i].is_none());
322            debug_assert!(self.parent.tp_kv[i].is_none());
323            debug_assert!(self.parent.recur[i].is_none());
324            debug_assert!(self.parent.latent[i].is_none());
325            self.parent.kv[i] = source.kv[i].take();
326            self.parent.tp_kv[i] = source.tp_kv[i].take();
327            self.parent.recur[i] = source.recur[i].take();
328            self.parent.latent[i] = source.latent[i].take();
329            self.parent.glm5_tp_recur[i] = source.glm5_tp_recur[i].take();
330            self.parent.glm5_tp_latent_peer[i] = source.glm5_tp_latent_peer[i].take();
331        }
332        self.parent.pos = self
333            .stages
334            .iter_mut()
335            .map(|stage| {
336                stage
337                    .get_mut()
338                    .unwrap_or_else(|poisoned| poisoned.into_inner())
339                    .pos
340            })
341            .min()
342            .unwrap_or(self.parent.pos);
343        if !self.committed {
344            self.parent.mark_tainted();
345        }
346    }
347}
348
349/// Fail-stop transaction marker for concat-prime paths. These paths mutate several independent
350/// caches before their final epilogue can fail; an error must make every member permanently
351/// ineligible for retry/reuse rather than replaying a queue over partially advanced state.
352struct CacheTaintGuard {
353    caches: Vec<*mut Cache>,
354    committed: bool,
355}
356
357impl CacheTaintGuard {
358    fn arm(caches: &mut [&mut Cache]) -> Self {
359        Self {
360            caches: caches
361                .iter_mut()
362                .map(|cache| *cache as *mut Cache)
363                .collect(),
364            committed: false,
365        }
366    }
367
368    fn commit(&mut self) {
369        self.committed = true;
370    }
371}
372
373impl Drop for CacheTaintGuard {
374    fn drop(&mut self) {
375        if self.committed {
376            return;
377        }
378        for cache in &self.caches {
379            // SAFETY: `arm` receives the function's unique cache references. The guard never
380            // escapes that call or dereferences them until unwind/return after active borrows end.
381            unsafe { (&mut **cache).mark_tainted() };
382        }
383    }
384}
385
386#[derive(Debug, Clone, Copy, PartialEq, Eq)]
387struct PrimePpWaveSlot {
388    wave: usize,
389    slot: usize,
390}
391
392#[derive(Debug)]
393enum PrimePpSignal {
394    Slot(PrimePpWaveSlot),
395    Error(String),
396}
397
398#[derive(Default)]
399struct PrimePpWaveCredits {
400    next_wave: usize,
401    pending: std::collections::VecDeque<PrimePpWaveSlot>,
402}
403
404impl PrimePpWaveCredits {
405    fn release_required(&self) -> Option<PrimePpWaveSlot> {
406        (self.pending.len() == 2).then(|| self.pending[0])
407    }
408
409    fn record_release(&mut self, released: PrimePpWaveSlot) -> Result<(), String> {
410        let expected =
411            self.pending.front().copied().ok_or_else(|| {
412                "prime PP received a slot release with no pending wave".to_string()
413            })?;
414        if released != expected {
415            return Err(format!(
416                "prime PP slot release {:?} does not match oldest pending {:?}",
417                released, expected
418            ));
419        }
420        self.pending.pop_front();
421        Ok(())
422    }
423
424    fn record_send(&mut self, sent: PrimePpWaveSlot) -> Result<(), String> {
425        if sent.wave != self.next_wave {
426            return Err(format!(
427                "prime PP sent wave {} while wave {} was next",
428                sent.wave, self.next_wave
429            ));
430        }
431        if sent.slot >= 2 {
432            return Err(format!(
433                "prime PP boundary returned invalid slot {}",
434                sent.slot
435            ));
436        }
437        if self.pending.iter().any(|pending| pending.slot == sent.slot) {
438            return Err(format!(
439                "prime PP reused slot {} before its exact-wave release",
440                sent.slot
441            ));
442        }
443        self.pending.push_back(sent);
444        self.next_wave += 1;
445        Ok(())
446    }
447}
448
449fn recv_prime_pp_signal(
450    receiver: &std::sync::mpsc::Receiver<PrimePpSignal>,
451    expected: PrimePpWaveSlot,
452    exact_slot: bool,
453    label: &str,
454) -> Result<PrimePpWaveSlot, String> {
455    match receiver.recv() {
456        Ok(PrimePpSignal::Error(error)) => Err(error),
457        Ok(PrimePpSignal::Slot(received))
458            if received.wave == expected.wave
459                && (!exact_slot || received.slot == expected.slot) =>
460        {
461            if received.slot >= 2 {
462                Err(format!(
463                    "{label}: wave {} carried invalid slot {}",
464                    received.wave, received.slot
465                ))
466            } else {
467                Ok(received)
468            }
469        }
470        Ok(PrimePpSignal::Slot(received)) => Err(format!(
471            "{label}: expected wave/slot {:?}, received {:?}",
472            expected, received
473        )),
474        Err(_) => Err(format!(
475            "{label}: channel closed while waiting for wave {}",
476            expected.wave
477        )),
478    }
479}
480
481fn send_prime_pp_signal(
482    sender: &std::sync::mpsc::Sender<PrimePpSignal>,
483    signal: PrimePpSignal,
484    label: &str,
485) -> Result<(), String> {
486    sender
487        .send(signal)
488        .map_err(|_| format!("{label}: channel closed"))
489}
490
491struct PrimePpWave<'a> {
492    start: usize,
493    end: usize,
494    tokens: &'a [u32],
495}
496
497struct PrimePpStageChannels {
498    incoming: Option<std::sync::mpsc::Receiver<PrimePpSignal>>,
499    release_upstream: Option<std::sync::mpsc::Sender<PrimePpSignal>>,
500    outgoing: std::sync::mpsc::Sender<PrimePpSignal>,
501    released_downstream: std::sync::mpsc::Receiver<PrimePpSignal>,
502}
503
504impl PrimePpStageChannels {
505    fn notify_failure(&self, error: &str) {
506        if let Some(upstream) = &self.release_upstream {
507            let _ = upstream.send(PrimePpSignal::Error(error.to_string()));
508        }
509        let _ = self.outgoing.send(PrimePpSignal::Error(error.to_string()));
510    }
511}
512
513/// The DSA k-pool indexer's resident state, borrowed for one `mla_attn_core` call.
514///
515/// TWO PLANES, DIFFERENT LIFETIMES. `state` is the packed `[k_norm | gate]` row per cached token
516/// (`LatentKvLayer::index_rows`); it is append-only and grows with the cache. `pool_keys` is the
517/// collapsed key per COMPLETE pool of `pool` such rows, and it is the residency win: a pool's key
518/// is final the moment its last row lands, so pools `[0, *ready)` are never recomputed and each
519/// call builds only the pools its own tokens completed. `ready` is written back through the
520/// borrow, so the caller must persist it alongside the buffers.
521///
522/// `pool_keys` is `Option` because its size needs the indexer's `pool`, which the state plan does
523/// not carry — `mla_kpool_indices` allocates it on first use and leaves it resident thereafter.
524/// A caller that hands over a fresh `None` every call (the stateless arm) gets the old
525/// rebuild-everything behaviour, which is exactly right when the state itself is per-call.
526pub struct IndexerPlanes<'a> {
527    pub state: &'a mut CudaSlice<f32>,
528    pub pool_keys: &'a mut Option<CudaSlice<f32>>,
529    pub ready: &'a mut usize,
530    /// PHYSICAL rows of `state` when it is a TAIL RING; 0 when the plane is flat (one row per
531    /// cached token, absolute addressing). `mla_kpool_indices` rounds this DOWN to a multiple of
532    /// the indexer's `pool` — the state plan does not carry `pool`, so the allocator cannot — and
533    /// proves the liveness bound against the rounded value before it appends.
534    pub state_ring_rows: usize,
535    /// Token capacity of the session, which sizes `pool_keys`. It is NOT derivable from
536    /// `state.len()` once `state` is a ring: the ring holds one call's tail, the pool-key plane
537    /// holds the whole context collapsed `pool`-to-one.
538    pub capacity_tokens: usize,
539}
540
541/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
542pub(crate) struct AttnPre {
543    pub q: cudarc::driver::CudaSlice<f32>,
544    pub k: cudarc::driver::CudaSlice<f32>,
545    pub v: cudarc::driver::CudaSlice<f32>,
546    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
547}
548
549/// task #18: one sequence's GDN prep outputs (the scan inputs).
550pub(crate) struct GdnPrep {
551    pub hk: usize,
552    pub q_l2: cudarc::driver::CudaSlice<f32>,
553    pub k_l2: cudarc::driver::CudaSlice<f32>,
554    pub v_g: cudarc::driver::CudaSlice<f32>,
555    pub beta: cudarc::driver::CudaSlice<f32>,
556    pub g_log: cudarc::driver::CudaSlice<f32>,
557    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
558    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
559}
560
561/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
562pub(crate) struct VerifyStreamScratch {
563    pub pos_d: CudaSlice<i32>,
564    pub row_ctrs: Vec<CudaSlice<i32>>,
565}
566use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
567
568struct MoeInputTraceWriter {
569    dir: std::path::PathBuf,
570    index: std::fs::File,
571    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
572}
573
574static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
575    std::sync::OnceLock::new();
576
577/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
578/// per-expert launch chain). See `moe_gdec_token`.
579fn gdec_enabled() -> bool {
580    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
581    *E.get_or_init(|| {
582        std::env::var("MEMRA_MOE_GDEC")
583            .map(|v| v != "0")
584            .unwrap_or(true)
585    })
586}
587
588/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
589/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
590/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
591/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
592/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
593/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
594/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
595/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
596/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
597/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
598fn moe_slab_enabled() -> bool {
599    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
600}
601
602/// `MEMRA_MOE_FUSED_EPI` — the glm5_next fused MoE epilogue (sigmoid-routed, PRE-clamped SwiGLU,
603/// per-expert macro fold) collapsed into one launch pair per token-layer.
604///
605/// DEFAULT OFF, deliberately (docs/FLAGS.md carries the row and the reasons). The arm is proven
606/// EXACT against `memra_reference` by `tests/glm5_moe_epilogue_gpu.rs`, but it has no throughput
607/// receipt: the rig is correctness-only by law and the 190.7 GB artifact has never been on it, so
608/// the launch-count claim is arithmetic from source and nothing has been measured on serving
609/// hardware. Unmeasured behavior does not default ON.
610///
611/// Read PER CALL, not latched in a `OnceLock`: the acceptance gate flips both arms inside one
612/// test process (the interleave unit is a model load, not a boot), and a latched flag would make
613/// the second arm silently a copy of the first.
614fn moe_fused_epi_enabled() -> bool {
615    std::env::var("MEMRA_MOE_FUSED_EPI")
616        .map(|v| v != "0")
617        .unwrap_or(false)
618}
619
620/// `MEMRA_HC_DECODE_WS` — the persistent hc-glue decode workspace (lane/glm5-decode-diet
621/// lever 2): the T=1 hc decode walk lands its glue transients (mixes, gates, comb, collapse
622/// y, both norm scratches, the per-site post output) in one per-engine `HyperDecodeWs`
623/// instead of ~12 fresh `cuMemAllocAsync`+free pairs per layer per token (the launch-diet
624/// census's 2,358-calls/token class). Same kernels, same call order, same operand bytes —
625/// byte identity ON/OFF gated by `tests/hc_decode_ws_gpu.rs`.
626///
627/// DEFAULT OFF, deliberately (docs/FLAGS.md row): the alloc-call reduction is proven on the
628/// rig by counter receipt, but the ms/token value is arithmetic against the box's measured
629/// launch/alloc constants — nothing has been measured on serving hardware yet. Unmeasured
630/// behavior does not default ON.
631///
632/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent).
633fn hyper_decode_ws_on() -> bool {
634    std::env::var("MEMRA_HC_DECODE_WS").as_deref() == Ok("1")
635}
636
637/// Engagement counter for the workspace walk — the receipt the gate and any box A/B arm
638/// must show (engagement lines are receipts, never inferred).
639pub static HC_DECODE_WS_DISPATCHES: std::sync::atomic::AtomicU64 =
640    std::sync::atomic::AtomicU64::new(0);
641
642/// `MEMRA_MLA_TC_PREFILL` — the glm5_next tensor-core MLA prefill chain
643/// (lane/glm5-mla-tc-prefill, 2026-08-30): at prefill widths the three per-position f32
644/// kernels the launch-diet census named (`memra_mla_attn_gathered_kernel` 139 ms +
645/// `memra_mla_absorb_q_kernel` 44.5 ms + `memra_mla_decompress_v_kernel` 43.6 ms per
646/// layer-chunk, 75.8% of a 98%-GPU-busy cold prime) are replaced by two strided-batched
647/// bf16 tensor-core GEMMs (absorb / decompress, one launch each) and one gathered
648/// flash-attention MMA kernel (`fa_mla_gathered_bf16`). Selection, the latent cache, the
649/// q/kv projections, and decode are UNTOUCHED.
650///
651/// DEFAULT ON (owner acceptance 2026-08-30, "why not? i dont see why not", on the two-box
652/// A/B receipts): interleaved x5 fresh boots per arm on BOTH the Server-Edition and
653/// Workstation-Edition 4-card boxes, zero violations, zero argmax flips across 20 boots,
654/// TTFD -62%..-69% (7.45->2.83 s @4.6k / 6.58->2.51 s), prefill 619-724 -> 1629-2255 tok/s,
655/// vendor-default sampled twin -66/-67%, decode untouched, engagement receipted in every ON
656/// boot with no cuBLASLt declines (docs/FLAGS.md row carries the pointers). The numeric
657/// config remains band-gated (bf16 operands, f32 accumulate — the fa_prefill/MEMRA_PP_BF16
658/// class, `tests/mla_tc_prefill_gpu.rs`, never bit). `MEMRA_MLA_TC_PREFILL=0` is the
659/// rollback seam.
660///
661/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent): the gate
662/// flips both arms inside one test process, and a latched flag would make the second arm
663/// silently a copy of the first.
664fn mla_tc_prefill_enabled() -> bool {
665    std::env::var("MEMRA_MLA_TC_PREFILL")
666        .map(|v| v != "0")
667        .unwrap_or(true)
668}
669
670/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
671/// default flip. `=0` selects the established path, while any other explicit value enables the
672/// grouped research arm for the current call.
673fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
674    std::env::var("MEMRA_MOE_GROUPED")
675        .map(|value| value != "0")
676        .unwrap_or(false)
677}
678
679/// `MEMRA_MOE_GROUPED_PREFILL`: the glm5_next expert-grouped MoE PREFILL arm, token-sort by
680/// expert (host CSR, the `moe_align_block_size` shape), then ONE grouped tensor-core GEMM per
681/// projection per layer-chunk over the resident NVFP4 bank, with the sigmoid `noaux_tc` routing,
682/// the PRE-clamped SwiGLU epilogue and the per-expert `weight_scale_2` macro fold the fused
683/// epilogue lane qualified for this family.
684///
685/// DEFAULT ON since 2026-08-29 (owner acceptance; `=0` is the rollback seam). The flip carries
686/// its receipts, per the flag-default law: reference-band + routing-exactness gate green
687/// (`tests/glm5_moe_grouped_prefill_gpu.rs`; grouped GEMM is measured non-bit-stable, so byte
688/// identity is not the honest bar; routing sel/w stay bit-identical by construction, the same
689/// `moe_router_logits` + `moe_route_sigmoid_cfg` invocation as the sequential arm), plus the
690/// interleaved x5 box A/B on the serving card class: TTFD 54.2 -> 7.5 s / 65.5 -> 8.9 /
691/// 75.9 -> 10.3 at 4.6/5.5/6.5k-token real prompts (85 -> 616-639 tok/s prefill, decode
692/// unchanged, sampled vendor-default twin green, engagement 42/42). The one greedy first-token
693/// flip (B5550) sits at a position the 8-draw vendor-default census measured as SOFT in both
694/// arms (the OFF arm itself draws the ON arm's token there) and was accepted by the OWNER on
695/// 2026-08-29, the MEMRA_BF16_MMV acceptance class. Receipts:
696/// `research/glm53-flash-bringup-20260827/moe-grouped-prefill-receipts/` (`box-ab-20260829/`).
697///
698/// Read PER CALL, not latched: the acceptance gate flips both arms inside one test process.
699fn moe_grouped_prefill_enabled() -> bool {
700    std::env::var("MEMRA_MOE_GROUPED_PREFILL")
701        .map(|v| v != "0")
702        .unwrap_or(true)
703}
704
705/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
706/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
707fn moe_prefetch_enabled() -> bool {
708    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
709    *E.get_or_init(|| {
710        std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
711            || crate::spill_pread::worker_enabled()
712    })
713}
714
715/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
716/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
717/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
718/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
719fn moe_page_prefetch_window() -> usize {
720    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
721    *W.get_or_init(|| {
722        page_prefetch_window_from_values(
723            std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
724            std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
725                .ok()
726                .as_deref(),
727        )
728    })
729}
730
731fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
732    if !enabled {
733        return 0;
734    }
735    raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
736}
737
738/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
739/// full window; each later position adds one expert at the far edge. Thus widening the window does
740/// not repeatedly issue `MADV_WILLNEED` for the same range.
741fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
742    if window == 0 || position >= len {
743        return len..len;
744    }
745    let (start, count) = if position == 0 {
746        (1, window)
747    } else {
748        (position.saturating_add(window), 1)
749    };
750    let start = start.min(len);
751    start..start.saturating_add(count).min(len)
752}
753
754/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
755/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
756fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
757    let position = current.map_or(0, |position| position.saturating_add(1));
758    (position < order_len).then_some(position)
759}
760
761/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
762/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
763/// window. Position zero primes the current expert too: its three independent reads can run in
764/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
765fn worker_prefetch_window() -> usize {
766    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
767    *WINDOW.get_or_init(|| {
768        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
769        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
770            .ok()
771            .and_then(|value| value.parse::<usize>().ok())
772            .unwrap_or(automatic.max(1))
773    })
774}
775
776/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
777/// this includes the current expert when the window is seeded so all three current projections
778/// enter the CPU pool together.
779fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
780    if window == 0 || position >= len {
781        return len..len;
782    }
783    let (start, count) = if position == 0 {
784        (0, window)
785    } else {
786        (position.saturating_add(window).saturating_sub(1), 1)
787    };
788    let start = start.min(len);
789    start..start.saturating_add(count).min(len)
790}
791
792/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
793/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
794/// expert weight pointers come from the per-layer device table. Requires the fused router (the
795/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
796fn moe_dev_enabled() -> bool {
797    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
798    *E.get_or_init(|| {
799        std::env::var("MEMRA_MOE_DEV")
800            .map(|v| v != "0")
801            .unwrap_or(true)
802            && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
803    })
804}
805
806/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
807/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
808/// Where the verify-rows MoE pair's routed selection lives for one layer-call
809/// (lane/glm5-moe-loc door D). `Host` is the shipped arm: the router's pinned readback gave the
810/// host `sel`/`w`, and the host builds the pointer/scale tables. `Dev` is door D's arm: the
811/// router's own device `sel_idx`/`sel_w` are still live, so the tables are built where they are
812/// and the readback (2 DtoH + a full `cuStreamSynchronize` per MoE layer-call) never happens.
813/// ONE launch path consumes both — only the table build differs, term-for-term identically.
814enum VrowsSel<'a> {
815    Host(&'a [u32], &'a [f32]),
816    Dev(&'a CudaSlice<i32>, &'a CudaSlice<f32>),
817}
818
819fn sigmoid_router_enabled() -> bool {
820    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
821    *E.get_or_init(|| {
822        std::env::var("MEMRA_SIG_ROUTER")
823            .map(|v| v != "0")
824            .unwrap_or(true)
825    })
826}
827
828/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
829/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
830/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
831/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
832fn moe_q8_enabled() -> bool {
833    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
834    *E.get_or_init(|| {
835        std::env::var("MEMRA_MOE_Q8")
836            .map(|v| v != "0")
837            .unwrap_or(true)
838    })
839}
840
841/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
842/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
843fn expert_dp4a_supported(qt: i32) -> bool {
844    qt == crate::QT_Q4_0
845        || qt == crate::QT_IQ3_S
846        || qt == crate::QT_IQ4_XS
847        || qt == crate::QT_Q3_K
848        || qt == crate::QT_Q4_K
849        || qt == crate::QT_Q6_K
850}
851
852fn q8_expert_supported(qt: i32) -> bool {
853    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
854    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
855    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
856    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
857    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
858    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
859    let kq = *KQ.get_or_init(|| {
860        std::env::var("MEMRA_MOE_Q8_KQ")
861            .map(|v| v != "0")
862            .unwrap_or(true)
863    });
864    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
865    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
866    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
867    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
868    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
869    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
870    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
871        .map(|v| v != "0")
872        .unwrap_or(true);
873    qt == crate::QT_IQ3_S
874        || qt == crate::QT_IQ4_XS
875        || (nvfp4_q8 && qt == crate::QT_NVFP4)
876        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
877}
878
879/// ModelOpt `W4A16_NVFP4` is weight-only: feeding its expert weights a q8_1 activation changes
880/// the declared numeric program to W4A8. Keep the established q8 path for every other artifact,
881/// but force Hy3 W4A16 experts through the BF16-activation `qmatvec_view` oracle.
882fn q8_expert_supported_for_model(cfg: &ModelConfig, qt: i32) -> bool {
883    let weight_only_nvfp4 = cfg.hy3.as_ref().is_some_and(|hy3| hy3.weight_only_nvfp4);
884    q8_expert_supported(qt) && !(weight_only_nvfp4 && qt == crate::QT_NVFP4)
885}
886
887fn moe_q8_enabled_for_model(cfg: &ModelConfig, m: &MoeWeights) -> bool {
888    m.has_uniform_expert_layout()
889        && moe_q8_enabled()
890        && q8_expert_supported_for_model(cfg, m.gate_exps.qtype)
891        && q8_expert_supported_for_model(cfg, m.up_exps.qtype)
892        && q8_expert_supported_for_model(cfg, m.down_exps.qtype)
893}
894
895#[cfg(test)]
896mod w4a16_dispatch_tests {
897    use super::q8_expert_supported_for_model;
898    use memra_gguf::config::{HfConfig, ModelConfig};
899
900    #[test]
901    fn hy3_w4a16_never_admits_q8_activations() {
902        let hf = HfConfig::parse(
903            r#"{"model_type":"hy_v3","num_hidden_layers":2,"hidden_size":8,
904            "num_attention_heads":2,"intermediate_size":16,"vocab_size":32,
905            "max_position_embeddings":32,
906            "quantization_config":{"quant_method":"modelopt","quant_algo":"W4A16_NVFP4"}}"#,
907        );
908        let cfg = ModelConfig::from_hf(&hf);
909        assert!(!q8_expert_supported_for_model(&cfg, crate::QT_NVFP4));
910        assert!(q8_expert_supported_for_model(&cfg, crate::QT_IQ4_XS));
911    }
912}
913
914/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
915/// k-quant tensors must fall to the _em dot path instead.
916fn q8_expert_dec_supported(qt: i32) -> bool {
917    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
918}
919
920/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
921/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
922/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
923/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
924/// q35 layers, which is why that cell measured FLAT.
925fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
926    match qt {
927        crate::QT_Q4_0 => in_f.is_multiple_of(32),
928        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
929            in_f.is_multiple_of(256)
930        }
931        // NVFP4 (block 64) added lane/moebatch-q35moe 2026-08-21: the ornith15 expert bank is
932        // uniform NVFP4, which passed the pairs q8 gate but missed BOTH batched doors
933        // (use_mma's dec set and this table), so 14.7k-token prefill rode the per-pair _em
934        // fallback — 88.6% of the prime wall (prime-anatomy receipt).
935        crate::QT_NVFP4 => in_f.is_multiple_of(64),
936        _ => false,
937    }
938}
939
940/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
941/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
942fn moe_prewarm_enabled() -> bool {
943    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
944    *E.get_or_init(|| {
945        std::env::var("MEMRA_MOE_PREWARM")
946            .map(|v| v != "0")
947            .unwrap_or(true)
948    })
949}
950
951/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
952/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
953/// can then vote for and exercise those experts on GPU before the cache is frozen.
954fn cpu_expert_profile_admit_enabled() -> bool {
955    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
956    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
957}
958
959/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
960/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
961/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
962pub const PRIME_MIN_T: usize = 16;
963
964/// CUDA grid dimensions y and z cap at 65,535 (an architecture constant on every compute
965/// capability we target; only grid.x is 2^31-1). Several prime-path kernels launch with the
966/// call's token count in grid.y — the fused GDN conv (`ssm_conv1d_gdn_state_f32`,
967/// lib.rs `ssm_conv1d_gdn_state_pad`), the dp4a matvec family taken at out_f < 128
968/// (`qmatvec*_fast`/`qmatvec_dp4a_named`, which is where the GDN ssm_beta/ssm_alpha
969/// projections land), and the MoE `router_gemv` — so ONE prime call must never carry more
970/// tokens than this. A monolithic prime above it is a guaranteed
971/// `DriverError(CUDA_ERROR_INVALID_VALUE)` at launch: measured on ornith-1.5 serving with
972/// MEMRA_PRIME_CHUNK=0 (cold 64,984 PASS / 65,643 FAIL, darklanes
973/// research/ornith-move-20260829 F2; re-hit on prod by the 2026-09-01 stress campaign at
974/// 66,045/79,717/82,440), and the same wall was hit and chunk-walked away by the glm5 1M
975/// lane's ppN prime.
976pub const CUDA_GRID_YZ_MAX: usize = 65_535;
977
978/// Widest prime range that stays launch-legal through the ring-off tail fold of
979/// `fixed_prime_chunk_ranges_for_ring`: a trailing remainder shorter than `PRIME_MIN_T`
980/// folds INTO the previous range, widening it by up to `PRIME_MIN_T - 1` tokens, so the
981/// cap keeps `chunk + PRIME_MIN_T - 1 <= CUDA_GRID_YZ_MAX`. With this value every
982/// t <= 65,535 still schedules as the identical single monolithic range (t <= chunk, or
983/// the fold collapses the split), so behavior below the CUDA wall is byte-for-byte
984/// unchanged — only prompts that today CANNOT launch get chunked.
985pub const PRIME_CHUNK_LAUNCH_CAP: usize = CUDA_GRID_YZ_MAX - (PRIME_MIN_T - 1);
986
987/// The explicit-`MEMRA_PRIME_CHUNK` chunk width, ring-aware. Extracted pure for tests.
988/// Ring ON keeps the historical clamp to `PRIME_CHUNK_MAX_TOKENS`. Ring OFF preserves the
989/// operator's value except at the CUDA launch wall: `0` ("monolithic") now means
990/// "monolithic up to `PRIME_CHUNK_LAUNCH_CAP`", and any larger explicit value is capped
991/// there too — an uncapped value above the wall never produced output, only
992/// CUDA_ERROR_INVALID_VALUE (see `CUDA_GRID_YZ_MAX`).
993fn explicit_prime_chunk(parsed: usize, ring_on: bool) -> usize {
994    if ring_on {
995        if parsed == 0 {
996            crate::cache::PRIME_CHUNK_MAX_TOKENS
997        } else {
998            parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
999        }
1000    } else if parsed == 0 {
1001        PRIME_CHUNK_LAUNCH_CAP
1002    } else {
1003        parsed.min(PRIME_CHUNK_LAUNCH_CAP)
1004    }
1005}
1006
1007/// MEMRA_STEP_GEMM_PRIME_SUFFIX: does a CONTINUATION prime (`cache.pos > 0` — a rewound
1008/// session's suffix, or a prompt remainder split across scheduler ticks) ride the batched
1009/// GEMM prime, like a fresh prompt does?
1010///
1011/// DEFAULT ON since 2026-08-29, by decision, under the flip bar the OFF-era FLAGS row
1012/// wrote down (never byte identity — a prime-decomposition m-dependence that EVERY
1013/// measured prime path shares, walk included, bars that gate for all of them):
1014///  1. vendor-default sampled rows: the blind, rubric-pre-registered 8-turn quality A/B
1015///     (research/step37-sampled-quality-20260828, 72/72 valid rows, engagement receipts
1016///     per row) — WARM-GEMM sits inside COLD's own self-spread at t4 and t8 (t8 carried
1017///     at n=16; the round-1 walk-over-gemm signal collapsed at p~0.91).
1018///  2. the 8-turn cache-on twin: warm TTFT 0.58 s (door) vs 7.15 s (walk) on the real
1019///     warm serving shape, zero faults.
1020///  3. the batched prime's own standard: acceptance 0.80-0.86 across all arms with the
1021///     door arm highest at t8; interleaved arms; zero ILLEGAL/#87/panics in 19 boots.
1022///
1023/// Precondition shipped first: the SWA-ring checkpoint restore fix (c9a617ca99) — real
1024/// session reuse crosses the grow path before any door question matters.
1025/// Why it is worth it, measured: the walk continuation costs 5.5978 ms/suffix-token
1026/// against this path's 0.99 ms/token (five-point sweep, R^2 0.9976), a 7.97x suffix
1027/// slope collapse.
1028///
1029/// The `seq_end` fix beneath is NOT gated on this door — it is unconditional, because
1030/// the chunk-local `seq_end` it replaced is wrong for a fresh prompt of 4096+k tokens
1031/// (k in [PRIME_MIN_T, 512)) with no continuation anywhere in sight.
1032///
1033/// `=0` is the kill switch (continuations back on the walk, fresh primes keep the fast
1034/// path); `=1` forces; `MEMRA_STEP_GEMM_PRIME=0` remains the whole-path seam. Read per
1035/// call, not cached — probes flip it in process.
1036fn step_gemm_prime_suffix_on() -> bool {
1037    std::env::var("MEMRA_STEP_GEMM_PRIME_SUFFIX").as_deref() != Ok("0")
1038}
1039
1040/// Widest tick the MoE DEV per-token program serves (lane/orndecode-20260822). PRIME_MIN_T
1041/// doubled as the dev-arm's upper bound on the assumption that t==16 only ever meant real
1042/// prefill; the exact-16 decode tier broke that assumption — at B=16 the MoE stage crossed
1043/// onto the t>=MMA_T grouped/kq GEMM program (m_e ~1.6 rows/expert: 52.6% of the tick at
1044/// ~104 us/launch) or the `_em` per-pair fallback (67.7 us), both catastrophically slower
1045/// than the dev q8 kernels that serve B<=8 (8.8 us gate_up covering a token's whole expert
1046/// set). Decode widths 2..=16 now ride dev; the grouped/pairs prefill programs start at 17.
1047/// gate2/gate3 byte batteries at B=12/16 are the qualification (bit-checked vs isolated).
1048const MOE_DEV_MAX_T: usize = 16;
1049const PRIME_PIPE_MICROBATCHES: usize = 8;
1050const PRIME_PIPE_MIN_CHUNK: usize = 128;
1051const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
1052const PRIME_PIPE_LINEAR_WORK: usize = 8;
1053
1054fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
1055    crate::pp::prime_pp_on()
1056        && !crate::pp::pp2_streams_off()
1057        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
1058}
1059
1060fn prime_ppn_wave_auto_geometry(n_layers: usize) -> bool {
1061    crate::pp::prime_pp_on()
1062        && !crate::pp::pp2_streams_off()
1063        && crate::pp::pp_wave_on() == Ok(true)
1064        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| matches!(cuts.len(), 4 | 5))
1065}
1066
1067fn prime_pipeline_auto_geometry(n_layers: usize) -> bool {
1068    prime_pp2_auto_geometry(n_layers) || prime_ppn_wave_auto_geometry(n_layers)
1069}
1070
1071/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative up to the
1072/// CUDA launch wall (`PRIME_CHUNK_LAUNCH_CAP`; 0 = monolithic up to that wall — see
1073/// `explicit_prime_chunk`).
1074/// Pipelined PP primes use the measured PP-2 geometry: up to eight microchunks, never below
1075/// 128 tokens, while the legacy 4096-token cap remains the long-context bound. PP-3/4 inherit
1076/// only the geometry when their separate MEMRA_PP_WAVE door is explicitly open.
1077pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
1078    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
1079        let parsed = value
1080            .parse::<usize>()
1081            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
1082        return explicit_prime_chunk(parsed, crate::cache::swa_ring_on());
1083    }
1084    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
1085    if prime_pipeline_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
1086        chunk.min(
1087            t.div_ceil(PRIME_PIPE_MICROBATCHES)
1088                .max(PRIME_PIPE_MIN_CHUNK),
1089        )
1090    } else {
1091        chunk
1092    }
1093}
1094
1095fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
1096    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
1097}
1098
1099fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
1100    if chunk == 0 || t <= chunk {
1101        return vec![(0, t)];
1102    }
1103    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
1104    let mut start = 0usize;
1105    while start < t {
1106        let mut end = (start + chunk).min(t);
1107        if t - end > 0 && t - end < PRIME_MIN_T {
1108            if ring_on {
1109                let shifted = t - PRIME_MIN_T;
1110                end = if shifted > start { shifted } else { t };
1111            } else {
1112                end = t;
1113            }
1114        }
1115        ranges.push((start, end));
1116        start = end;
1117    }
1118    ranges
1119}
1120
1121fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
1122    let prefix = prefix as u128;
1123    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
1124}
1125
1126fn dynamic_prime_chunk_ranges(
1127    t: usize,
1128    fixed_chunk: usize,
1129    fixed: &[(usize, usize)],
1130) -> Vec<(usize, usize)> {
1131    let n = fixed.len();
1132    if n < 3 {
1133        return fixed.to_vec();
1134    }
1135
1136    let max_first = t - (n - 1) * PRIME_MIN_T;
1137    let first = fixed_chunk
1138        .div_ceil(2)
1139        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
1140        .min(max_first);
1141    let mut ranges = Vec::with_capacity(n);
1142    ranges.push((0, first));
1143
1144    let first_work = prime_chunk_work(first, t);
1145    let work_span = prime_chunk_work(t, t) - first_work;
1146    let denominator = (n - 1) as u128;
1147    let mut previous = first;
1148    for boundary in 1..n - 1 {
1149        let target = first_work * denominator + work_span * (boundary as u128);
1150        let remaining = n - 1 - boundary;
1151        let mut low = previous + PRIME_MIN_T;
1152        let mut high = t - remaining * PRIME_MIN_T;
1153        while low < high {
1154            let mid = low + (high - low) / 2;
1155            if prime_chunk_work(mid, t) * denominator >= target {
1156                high = mid;
1157            } else {
1158                low = mid + 1;
1159            }
1160        }
1161        ranges.push((previous, low));
1162        previous = low;
1163    }
1164    ranges.push((previous, t));
1165    ranges
1166}
1167
1168/// Internal prime ranges. A pipelined PP prime defaults to a short-fill, equal-modeled-time
1169/// schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
1170/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
1171///
1172/// `gdn_grid`: the model runs the chunked GDN WY scan (`HybridModel::gdn_prime_grid_on`) —
1173/// AUTO-scheduled internal boundaries are then snapped down to the WY-chunk grid
1174/// (`align_prime_ranges_to_gdn`; the spec-longctx grid law, extended from serve splits to
1175/// the PP prime microchunks). Explicit MEMRA_PRIME_CHUNK keeps its operator-authoritative
1176/// (fixed, unaligned) semantics — the FLAGS caveat documents that identity contract.
1177pub fn prime_chunk_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1178    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
1179    let chunk = prime_chunk_tokens(t, n_layers);
1180    let fixed = fixed_prime_chunk_ranges(t, chunk);
1181    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
1182        Ok(value) => value == "dynamic",
1183        Err(_) => true,
1184    };
1185    if explicit_chunk {
1186        return fixed;
1187    }
1188    let ranges = if !dynamic || !prime_pipeline_auto_geometry(n_layers) {
1189        fixed
1190    } else {
1191        dynamic_prime_chunk_ranges(t, chunk, &fixed)
1192    };
1193    // MEMRA_PRIME_GRID_ALIGN=0 is the shared rollback seam of the grid law (same env the
1194    // worker's serve-boundary alignment honors, read per call so gates can flip it
1195    // in-process): the legacy off-grid auto schedule — the toothed cell's broken arm.
1196    if gdn_grid && std::env::var("MEMRA_PRIME_GRID_ALIGN").as_deref() != Ok("0") {
1197        align_prime_ranges_to_gdn(&ranges, t, Engine::gdn_chunk_size())
1198    } else {
1199        ranges
1200    }
1201}
1202
1203/// The mHC prime schedule: how `prime_cache_hyper` SPLITS one prompt into calls.
1204///
1205/// THIS IS THE RULE THE CAPACITY GATE ASSERTS ON, and it is separate from
1206/// [`prime_chunk_ranges`] so the hyper walk's split can be read, gated and changed without
1207/// touching the serial trunk's. It DELEGATES to the serial schedule rather than deriving a
1208/// second one: the transient pressure it answers is the same pressure, and two schedules would
1209/// be two things to keep aligned with `MEMRA_PRIME_CHUNK`.
1210///
1211/// WHY THE SPLIT IS SEMANTICALLY INERT, term by term. Read this precisely: it says the split
1212/// computes the SAME PROGRAM, not that it computes the same BITS. The bit claim is FALSE on this
1213/// trunk and measured so — see the near-tie note at the end.
1214///
1215///   * **The mHC residual is strictly PER TOKEN.** `crate::hyper`'s contract is `mixes[t,:]`, an
1216///     RMS rescale over that token's own `streams*hidden` slab, a Sinkhorn per token per site,
1217///     a per-token collapse and a per-token post. The stream state is expanded at the start of a
1218///     call and collapsed at its end; it carries NOTHING between tokens. Splitting the token
1219///     axis cannot move a value.
1220///   * **KDA prefill is a SEQUENTIAL scan** (`kda.rs`: `memra_kda_scan_s128` runs prefill and
1221///     decode alike, the chunked UT transform is not the shipped path). A sequential recurrence
1222///     has no fold grid, so the GDN WY grid law has NO KDA analogue to violate and
1223///     `align_prime_ranges_to_gdn` has nothing to align. The conv ring carries across calls
1224///     already — it is the seam every decode step uses. **DEBT, named:** if the chunked KDA twin
1225///     ever becomes the prefill path, it acquires a fold grid and this schedule's internal
1226///     boundaries must be snapped to it exactly as the GDN ones are, or chunked prime stops
1227///     being bit-identical. The `gdn_grid` argument is the seam that change lands on.
1228///   * **The latent KV plane is f32** (`LatentKvLayer::rows`), so a later call reads earlier
1229///     calls' rows in the SAME numeric class it would have computed them in. There is no
1230///     analogue of the serial trunk's f32-vs-quantized-KV class edge — the thing that made
1231///     `MEMRA_PRIME_CHUNK` steer arithmetic until the 2026-08-05 grain-free fix.
1232///   * **The DSA pool keys are already incremental** (`index_pools_ready`): a pool key is a pure
1233///     function of its own `pool` state rows and the constant `kpool_ape`, final the instant the
1234///     pool's last row lands, so no boundary can move one. Selection is per query over resident
1235///     keys, with visibility keyed on the query's ABSOLUTE cache row.
1236///
1237/// NOT BIT-STABLE ACROSS CHUNK SIZES, and the cause is NOT the split. Measured on the rig
1238/// (`research/glm53-flash-bringup-20260827/1m-context-20260828/02`): the arms diverge at ROW 0,
1239/// which no cross-token state can reach, and `Engine::linear` — the cuBLASLt f32 `mixes` GEMM in
1240/// `hyper::pre` — is itself not m-invariant (m=32 vs m=200 moves 9601/12288 output bits at worst
1241/// 3.815e-6, the same worst the chunked prime reports; m=128 and m=199 vs m=200 are identical).
1242/// cuBLASLt reselects its algorithm by shape and the reduction order goes with it, which
1243/// `hyper.rs`'s header already concedes ("a serving trunk, not a byte-parity oracle"). So this
1244/// is a documented near-tie class the split EXPOSES, not one it creates. It is written into the
1245/// `MEMRA_PRIME_CHUNK` FLAGS row, and `glm5_chunked_prime_gpu` holds the split to a calibrated
1246/// band anchored on `memra_reference::execute` rather than on the monolithic sibling.
1247///
1248/// One thing the split provably cannot break here: POSITIONS. glm5_next is NoPE end to end
1249/// (`qk_rope_head_dim = 0`, `mla_use_nope`), and KDA is positionless, so `pos_d` reaches no
1250/// kernel on this path — a mutation that made it call-local instead of session-absolute moved
1251/// nothing at all.
1252///
1253/// A prompt at or under one chunk takes the monolithic body unchanged, and `MEMRA_PRIME_CHUNK=0`
1254/// restores the monolithic walk up to `PRIME_CHUNK_LAUNCH_CAP` (65,520 tokens; the CUDA
1255/// grid.y wall is 65,535 and the old walk always died at launch above it, 08-29 ppN receipts)
1256/// — the rollback seam, and the oracle arm the correctness gate compares against, both now
1257/// bounded by that cap; longer prompts split at the cap even under `0`.
1258pub fn hyper_prime_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1259    prime_chunk_ranges(t, n_layers, gdn_grid)
1260}
1261
1262/// The largest number of token rows any ONE mHC prime call carries under
1263/// [`hyper_prime_ranges`]. Every per-call transient in the walk — the `t*streams*hidden` stream
1264/// state, the MLA query planes, and the DSA indexer's `t * n_pools` score plane — is
1265/// proportional to this, so it is the single number a capacity assertion needs.
1266pub fn hyper_prime_call_rows(t: usize, n_layers: usize, gdn_grid: bool) -> usize {
1267    hyper_prime_ranges(t, n_layers, gdn_grid)
1268        .iter()
1269        .map(|&(start, end)| end - start)
1270        .max()
1271        .unwrap_or(0)
1272}
1273
1274/// Per-request prefill WORKSPACE coefficients for a HyperConnections trunk, published to
1275/// admission (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper model: their
1276/// admission arithmetic is byte-identical to the pre-lane behavior.
1277///
1278/// These are the FORMULA behind the 262k 2-card cell's measured ~0.8 MiB/token/card prefill
1279/// wall (`research/glm53-flash-bringup-20260827/262k-2card-20260830/LANE.md`), not the slope
1280/// itself: each term is the size of a named allocation in the walk, summed per token of ONE
1281/// prime call. On GLM-5.3-Flash geometry (H=4096, S=4, F=2048, U=8, heads=64, qk=256, v=256,
1282/// topk=2048, P=4) `chunk_token_bytes` evaluates to ~0.86 MiB — the receipt's slope with the
1283/// conservative side up. The attribution table naming every term lives in
1284/// `research/glm53-flash-bringup-20260827/gpf-workspace-20260830/LANE.md` §1.
1285#[derive(Debug, Clone, Copy)]
1286pub struct HyperPrimeWorkspaceShape {
1287    /// Bytes of per-call prefill transients PER TOKEN OF ONE PRIME CALL: the double-buffered
1288    /// `[t, streams, hidden]` stream state + ppN boundary slots, the pre/norm transients, the
1289    /// grouped-MoE staging (CSR activations + three f32 partial planes + f16 mirrors + scatter
1290    /// planes), the MLA query/attention planes, the k-pool idx plane, and the prime-tail
1291    /// hidden/norm pair. Multiplied by [`hyper_prime_call_rows`] this bounds the workspace of
1292    /// the CHUNKED prime; on the monolithic rollback (`MEMRA_PRIME_CHUNK=0`) the call rows are
1293    /// the whole prompt up to `PRIME_CHUNK_LAUNCH_CAP` (65,535 launch-legal max; admission
1294    /// re-derives `hyper_prime_call_rows` so the arithmetic stays consistent either way) and
1295    /// the same product stays honest.
1296    pub chunk_token_bytes: usize,
1297    /// Bytes per PROMPT token that live for the WHOLE prime on the last stage: the returned
1298    /// pre-output_norm `hiddens` stack (`n_embd` f32), consumed by the MTP-spec `prompt_h`
1299    /// and the embed capture.
1300    pub prompt_bytes_per_token: usize,
1301    /// DSA k-pool group size `P`, or 0 when the model runs no k-pool indexer. The selection
1302    /// score plane of ONE call is `call_rows * (ctx / P)` f32 — the one prefill transient that
1303    /// stays COUPLED TO CONTEXT DEPTH after chunking (it is the allocation the 3-card 1M prime
1304    /// died on at 97.2 GiB).
1305    pub kpool_score_pool: usize,
1306    /// Trunk layer count, for re-deriving [`hyper_prime_call_rows`] at admission time with the
1307    /// same env-sensitive schedule the prime itself will walk.
1308    pub n_layers: usize,
1309    /// The model's own GDN grid-alignment input to the schedule.
1310    pub gdn_grid: bool,
1311}
1312
1313impl HyperPrimeWorkspaceShape {
1314    /// The admission charge for one request: workspace of the LARGEST prime call this request
1315    /// can produce, plus the ctx-coupled score plane at that call width, plus the prompt-long
1316    /// hiddens stack.
1317    ///
1318    /// Keyed on PROMPT rows, deliberately not on `ctx_cap`: every term here is a function of
1319    /// what the PRIME walks, and a `max_tokens`-omitted request carries a `ctx_cap` of the
1320    /// whole server window — charging the window would refuse every vendor-default short
1321    /// prompt on a deep-window box for workspace it never allocates. A continuation request's
1322    /// `prompt` is the full rendered conversation (the suffix optimization is internal
1323    /// reuse), so the score plane's `t_kv` is covered too.
1324    pub fn admission_bytes(&self, prompt_rows: usize) -> usize {
1325        let rows = hyper_prime_call_rows(prompt_rows, self.n_layers, self.gdn_grid);
1326        let chunk = self.chunk_token_bytes.saturating_mul(rows);
1327        let score = prompt_rows
1328            .checked_div(self.kpool_score_pool)
1329            .map(|pools| rows.saturating_mul(pools).saturating_mul(size_of::<f32>()))
1330            .unwrap_or(0);
1331        chunk
1332            .saturating_add(score)
1333            .saturating_add(self.prompt_bytes_per_token.saturating_mul(prompt_rows))
1334    }
1335}
1336
1337/// Snap AUTO prime-range internal boundaries DOWN to the GDN WY-chunk grid (lane/
1338/// hermes-perf-fixes, 2026-08-23 — the missing helper the PP-auto-ranges finding names).
1339///
1340/// THE LAW THIS EXTENDS (measured, research/multiturn-cache-20260821/
1341/// LONGCTX-EXACTNESS-20260821.md; the serve-split half already ships as the worker's
1342/// `grid_align_boundary`): under the chunked WY scan a prompt primed as two calls split at
1343/// L is bit-identical to the monolithic prime iff `L % gdn_chunk_size() == 0` — an off-grid
1344/// call start shifts the fold grid and materializes recurrent state at a point the
1345/// monolithic program never computes. The prime loop walks these ranges as separate
1346/// `prime_layers` calls, so INTERNAL microchunk boundaries are the same seam: the PP-2
1347/// auto geometry (`t.div_ceil(8).max(128)` fills, and every dynamic short-fill boundary)
1348/// lands off the 32-token grid for most prompt lengths, which is exactly the
1349/// chunk-value bit-identity the GDN lane falsified (FLAGS PRIME_CHUNK/SCHED caveat).
1350///
1351/// Boundaries only move DOWN (earlier is always semantically safe — same argument as the
1352/// worker's alignment); a boundary that collapses onto its predecessor is dropped (ranges
1353/// merge). The final range always ends at `t`. Aligning down only GROWS the tail
1354/// remainder, so the fixed-schedule tail-merge rule is never re-violated. Cost bound: at
1355/// most `c-1` tokens shift per boundary.
1356pub fn align_prime_ranges_to_gdn(
1357    ranges: &[(usize, usize)],
1358    t: usize,
1359    c: usize,
1360) -> Vec<(usize, usize)> {
1361    if c == 0 || ranges.len() < 2 {
1362        return ranges.to_vec();
1363    }
1364    let mut out: Vec<(usize, usize)> = Vec::with_capacity(ranges.len());
1365    let mut start = 0usize;
1366    for (i, &(_, end)) in ranges.iter().enumerate() {
1367        let e = if i + 1 == ranges.len() {
1368            t
1369        } else {
1370            end / c * c
1371        };
1372        if e > start {
1373            out.push((start, e));
1374            start = e;
1375        } // else: boundary collapsed onto its predecessor — merge into the next range
1376    }
1377    debug_assert_eq!(out.last().map(|&(_, e)| e), Some(t));
1378    out
1379}
1380
1381struct HeadSplit {
1382    pin: u64,
1383    w1: CudaSlice<u8>,
1384    hn1: CudaSlice<f32>,
1385    y1: CudaSlice<f32>,
1386    logits_e: CudaSlice<f32>,
1387    ev_hn: cudarc::driver::CudaEvent,
1388    ev_done: cudarc::driver::CudaEvent,
1389    raw_hn1: u64,
1390    raw_y1: u64,
1391    raw_logits_hi: u64,
1392    /// SAMPLED-TAIL scratch (perturbed row + the filter's threshold/z/max slots + the row
1393    /// index). Allocating these per token cost more than the split head saved: the first
1394    /// sampled-split measurement came in at 78.25 tok/s against 78.96 for the unsplit head,
1395    /// which is five allocations per token, not arithmetic.
1396    samp: Option<SampScratch>,
1397}
1398
1399struct SampScratch {
1400    pb: CudaSlice<f32>,
1401    th: CudaSlice<f32>,
1402    z: CudaSlice<f32>,
1403    mx: CudaSlice<f32>,
1404    rows: CudaSlice<i32>,
1405}
1406/// HEAD-SPLIT workspace (host + device twins share it).
1407static HEAD_SPLIT_WS: std::sync::Mutex<Option<HeadSplit>> = std::sync::Mutex::new(None);
1408
1409/// DEV1-LOCAL ROUTER replicas (MEMRA_DEV1_ROUTER): per-layer (gate_inp_f32, exp_probs_b,
1410/// active_experts) on rank1 + a shared logits scratch. Deterministic kernels on identical
1411/// input bits — rank1's local selection is bit-equal to the root's.
1412#[allow(clippy::type_complexity)]
1413static DEV1_ROUTER_REPS: std::sync::Mutex<
1414    Option<(
1415        std::collections::HashMap<u16, (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<u8>)>,
1416        Option<CudaSlice<f32>>,
1417    )>,
1418> = std::sync::Mutex::new(None);
1419
1420/// SHEXP-ON-DEV1 workspace (MEMRA_SHEXP_DEV1): replica weights + scratch on rank1, the
1421/// down row lands ROOT-resident over P2P (single store pass), and apply adds it on e
1422/// behind ev_done. (pins, wg1, wu1, wd1, act1, sh_root, ev_z, ev_done).
1423#[allow(clippy::type_complexity)]
1424static SHEXP_D1_REPS: std::sync::Mutex<
1425    Option<std::collections::HashMap<u16, (CudaSlice<u8>, CudaSlice<u8>, CudaSlice<u8>)>>,
1426> = std::sync::Mutex::new(None);
1427#[allow(clippy::type_complexity)]
1428static SHEXP_D1_WS: std::sync::Mutex<
1429    Option<(
1430        (usize, usize),
1431        CudaSlice<f32>,
1432        CudaSlice<f32>,
1433        CudaSlice<f32>,
1434        cudarc::driver::CudaEvent,
1435        cudarc::driver::CudaEvent,
1436    )>,
1437> = std::sync::Mutex::new(None);
1438
1439/// SHEXP OVERLAP workspace (issue writes, apply reads): (device, n_embd, n_ff_sh, act, sh).
1440#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1441static SHEXP_OV_WS: std::sync::Mutex<
1442    Option<(usize, usize, usize, CudaSlice<f32>, CudaSlice<f32>)>,
1443> = std::sync::Mutex::new(None);
1444
1445impl HybridModel {
1446    /// Does this model's prime schedule live under the GDN WY-chunk grid law? True when the
1447    /// trunk has GDN (linear-attention) layers AND the chunked scan is on — the regime where
1448    /// an off-grid prime-call boundary shifts the WY fold grid (see
1449    /// `align_prime_ranges_to_gdn`). Attention-only models and the sequential scan
1450    /// (`MEMRA_GDN_CHUNKED=0`) are split-invariant, so the grid is a no-op contract there.
1451    pub fn gdn_prime_grid_on(&self) -> bool {
1452        Engine::gdn_chunked_enabled()
1453            && self
1454                .layers
1455                .iter()
1456                .any(|l| matches!(l.mixer, crate::hybrid::Mixer::Linear(_)))
1457    }
1458
1459    /// Can the step TP runtime run the DEVICE-RESIDENT activation path from this serving
1460    /// engine? Native P2P (peer copies replace the host staging) AND a shared root context
1461    /// (the device buffers must be addressable on both sides — the TP registry builds its
1462    /// own Engine per rank, so this is a real seam, not a formality).
1463    fn full_attn_tp_device_resident(e: &Engine, tp: &crate::hybrid::StepTpQkv) -> bool {
1464        tp.runtime.native_p2p() && tp.runtime.root_shares_ctx(e)
1465    }
1466
1467    pub(crate) fn full_attn_tp_qkv(
1468        &self,
1469        e: &Engine,
1470        fa: &FullAttnLayer,
1471        h: &CudaSlice<f32>,
1472        t: usize,
1473    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
1474        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1475            return Ok(None);
1476        };
1477        let values = active_matrix_values(
1478            h.len(),
1479            t,
1480            self.cfg.n_embd as usize,
1481            "Step TP QKV activation",
1482        )?;
1483        // DEVICE-RESIDENT NATIVE PATH (lane/hermes-perf-fixes, 2026-08-23 — the host-bounce
1484        // finding): the native-P2P arm used to dtoh the FULL hidden state per layer, run
1485        // from a host copy, gather q/k/v to host vectors, and htod all three back — a host
1486        // round-trip on every execute that the peer transport exists to remove. The
1487        // device twins are byte-identical by construction (the same bytes travel dtod
1488        // instead of dtoh+htod; kernels, peer copies, and gather order are shared code).
1489        // The host arm below remains the transport for !native_p2p (host staging IS that
1490        // transport) and for a root context this engine cannot address.
1491        if Self::full_attn_tp_device_resident(e, tp) {
1492            // Producer fence: h was written on THIS engine's stream; the TP ranks read it
1493            // on theirs (same context, different streams).
1494            e.stream().synchronize()?;
1495            let q = tp
1496                .runtime
1497                .bf16_column_parallel_resident_native_device(&tp.q, h, t)?;
1498            let k = tp
1499                .runtime
1500                .bf16_column_parallel_resident_native_device(&tp.k, h, t)?;
1501            let v = tp
1502                .runtime
1503                .bf16_column_parallel_resident_native_device(&tp.v, h, t)?;
1504            Self::full_attn_tp_log_once(tp, "qkv", "device-resident");
1505            return Ok(Some(vec![q, k, v]));
1506        }
1507        let host = e.dtoh_view(&h.slice(0..values))?;
1508        let q = if tp.runtime.native_p2p() {
1509            tp.runtime
1510                .bf16_column_parallel_resident_native(&tp.q, &host, t)?
1511        } else {
1512            tp.runtime
1513                .bf16_column_parallel_resident(&tp.q, &host, t)?
1514                .gathered
1515        };
1516        let k = if tp.runtime.native_p2p() {
1517            tp.runtime
1518                .bf16_column_parallel_resident_native(&tp.k, &host, t)?
1519        } else {
1520            tp.runtime
1521                .bf16_column_parallel_resident(&tp.k, &host, t)?
1522                .gathered
1523        };
1524        let v = if tp.runtime.native_p2p() {
1525            tp.runtime
1526                .bf16_column_parallel_resident_native(&tp.v, &host, t)?
1527        } else {
1528            tp.runtime
1529                .bf16_column_parallel_resident(&tp.v, &host, t)?
1530                .gathered
1531        };
1532        Self::full_attn_tp_log_once(tp, "qkv", "host-canonical");
1533        Ok(Some(vec![e.htod(&q)?, e.htod(&k)?, e.htod(&v)?]))
1534    }
1535
1536    /// One transport banner per (projection, transport) — the old per-call eprintln fired
1537    /// on EVERY layer of EVERY step, itself a decode-rate cost on the path this lane is
1538    /// unbouncing (the sibling grouped-EP path already learned this).
1539    fn full_attn_tp_log_once(tp: &crate::hybrid::StepTpQkv, proj: &str, activation: &'static str) {
1540        use std::sync::atomic::{AtomicBool, Ordering};
1541        static LOGGED: [AtomicBool; 4] = [
1542            AtomicBool::new(false),
1543            AtomicBool::new(false),
1544            AtomicBool::new(false),
1545            AtomicBool::new(false),
1546        ];
1547        let idx = 2 * usize::from(proj == "o") + usize::from(activation == "device-resident");
1548        if LOGGED[idx].swap(true, Ordering::Relaxed) {
1549            return;
1550        }
1551        eprintln!(
1552            "[step-tp-{proj}] execute layer={} devices={:?} projections={proj} \
1553             tensor_parallel=true attention_local=true kv_local=true transport={} \
1554             native_p2p={} bulk_p2p={} activation={activation} \
1555             output={} performance_claim=false (logged once per transport)",
1556            tp.layer,
1557            tp.devices,
1558            tp.runtime.transport_label(),
1559            tp.runtime.native_p2p(),
1560            tp.runtime.bulk_p2p(),
1561            if activation == "device-resident" {
1562                "root-resident"
1563            } else {
1564                "root-readback"
1565            },
1566        );
1567    }
1568
1569    pub(crate) fn full_attn_tp_o(
1570        &self,
1571        e: &Engine,
1572        fa: &FullAttnLayer,
1573        activation: &CudaSlice<f32>,
1574        tokens: usize,
1575    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1576        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1577            return Ok(None);
1578        };
1579        // DEVICE-RESIDENT NATIVE PATH — the O-projection half of the same finding: no DtoH
1580        // of the attention output, no host O staging, root-resident reduction consumed in
1581        // place (byte-identical shared core: `step_bf16_row_native_reduce_from_root`).
1582        if Self::full_attn_tp_device_resident(e, tp) {
1583            e.stream().synchronize()?; // producer fence, as the QKV half
1584            let output = tp
1585                .runtime
1586                .step_bf16_row_parallel_resident_native_device(&tp.o, activation, tokens)?;
1587            Self::full_attn_tp_log_once(tp, "o", "device-resident");
1588            return Ok(Some(output));
1589        }
1590        let host = e.dtoh(activation)?;
1591        let output = if tp.runtime.native_p2p() {
1592            tp.runtime
1593                .step_bf16_row_parallel_resident_native(&tp.o, &host, tokens)?
1594        } else {
1595            tp.runtime
1596                .step_bf16_row_parallel_resident(&tp.o, &host, tokens)?
1597        };
1598        Self::full_attn_tp_log_once(tp, "o", "host-canonical");
1599        Ok(Some(e.htod(&output)?))
1600    }
1601
1602    fn full_attn_o(
1603        &self,
1604        e: &Engine,
1605        fa: &FullAttnLayer,
1606        activation: &CudaSlice<f32>,
1607        tokens: usize,
1608    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1609        match self.full_attn_tp_o(e, fa, activation, tokens)? {
1610            Some(output) => Ok(output),
1611            None => e.matmul(&fa.wo, activation, tokens),
1612        }
1613    }
1614
1615    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
1616    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
1617    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
1618    /// (it forces a dtoh + host hash per layer).
1619    fn prime_trace_path() -> Option<&'static str> {
1620        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
1621        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
1622            .as_deref()
1623    }
1624
1625    /// PRIME ANATOMY (diagnostic): `MEMRA_PRIME_ANATOMY=1` synchronizes the stream around
1626    /// each prime_layers stage and accumulates wall time per stage class, printed after
1627    /// every prime_layers call (cumulative across chunks/reps). The per-stage syncs
1628    /// serialize launch/execute overlap, so the summed total exceeds the naked prime wall —
1629    /// attribution ratios only, never a measured default run. Non-seg serial arm only.
1630    fn prime_anatomy_on() -> bool {
1631        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1632        *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
1633    }
1634
1635    fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
1636        static S: [std::sync::atomic::AtomicU64; 5] = [
1637            std::sync::atomic::AtomicU64::new(0), // 0 mixer full-attn
1638            std::sync::atomic::AtomicU64::new(0), // 1 mixer linear-attn (GDN)
1639            std::sync::atomic::AtomicU64::new(0), // 2 ffn MoE (router + experts + shexp)
1640            std::sync::atomic::AtomicU64::new(0), // 3 ffn dense
1641            std::sync::atomic::AtomicU64::new(0), // 4 norms/adds/glue
1642        ];
1643        &S
1644    }
1645
1646    /// Fail closed on a path that has not been taught the mHC residual program.
1647    ///
1648    /// A serial residual on an hc model is not a degraded answer, it is a DIFFERENT function
1649    /// computed at full speed and full confidence — the exact failure `crate::hyper` exists to
1650    /// prevent. Every trunk entry point that has not been converted calls this first, so the
1651    /// unconverted set is a list of named refusals rather than a list of silent wrong answers.
1652    pub(crate) fn refuse_hyper(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
1653        if let Some(topology) = self.hyper.as_ref() {
1654            return Err(format!(
1655                "{path} runs a serial residual, but this model's ModelPlan declares \
1656                 ResidualTopology::HyperConnections{{ streams: {}, collapse: {:?} }}. Refusing: \
1657                 that path would compute a different model. Converted paths: forward, \
1658                 forward_last, prime_cache, decode_step, and the batched serving chain \
1659                 decode_step_batch / _sampled / _lean / _masked.",
1660                topology.streams, topology.collapse
1661            )
1662            .into());
1663        }
1664        Ok(())
1665    }
1666
1667    /// The FFN branch of one hc site, from an already-normed `[t, hidden]` input.
1668    ///
1669    /// Split out because under hyper-connections the FFN's input is `rms_norm(hc_pre(x))`, not
1670    /// `rms_norm(x + attn)` — the fused add+norm+quantize forms the serial paths use have no
1671    /// residual to fold, so this is the unfused dispatch by construction.
1672    fn hyper_ffn_branch(
1673        &self,
1674        e: &Engine,
1675        layer: &crate::hybrid::HybridLayer,
1676        z: &CudaSlice<f32>,
1677        t: usize,
1678        il: usize,
1679        prefill: bool,
1680    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1681        match &layer.ffn {
1682            crate::hybrid::Ffn::Dense {
1683                ffn_gate,
1684                ffn_up,
1685                ffn_down,
1686            } => {
1687                let n_ff = ffn_gate.out_features();
1688                let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], z, t)?;
1689                let up = g2.pop().unwrap();
1690                let gate = g2.pop().unwrap();
1691                let mut act = e.uninit(t * n_ff)?;
1692                // A dense FFN reads the SHEXP clamp array — see forward()'s note.
1693                Self::ffn_act_lim(
1694                    e,
1695                    &self.cfg,
1696                    &gate,
1697                    &up,
1698                    1.0,
1699                    1.0,
1700                    self.cfg.clamp_shexp_at(il as u32),
1701                    &mut act,
1702                    t * n_ff,
1703                )?;
1704                e.matmul(ffn_down, &act, t)
1705            }
1706            crate::hybrid::Ffn::Moe(m) => {
1707                if prefill {
1708                    self.moe_ffn_il_prefill(e, m, z, t, il as u16)
1709                } else {
1710                    self.moe_ffn_il_zq8(e, m, z, None, t, il as u16)
1711                }
1712            }
1713        }
1714    }
1715
1716    /// Stateless prefill under the mHC residual (`crate::hyper`), the hc twin of `forward` /
1717    /// `forward_last`. The mixers, the FFNs and the norms are the SAME calls the serial paths
1718    /// make; only the residual program around them changes, which is the whole point — a mixer
1719    /// never sees the stream dimension.
1720    fn forward_hyper(
1721        &self,
1722        e: &Engine,
1723        tokens: &[u32],
1724        last_only: bool,
1725    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
1726        let topology = *self
1727            .hyper
1728            .as_ref()
1729            .ok_or("forward_hyper on a model with no HyperConnections topology")?;
1730        // M2 ppN door for the mHC trunk. The generic arm's door lives in `decode_step_h`;
1731        // this walk is reached BEFORE it (decode.rs routes `hyper.is_some()` first), so the
1732        // hc walks own their own door. Loud refusal, never silent fallback: an unqualified
1733        // pipeline rewrite errors here rather than running a single-engine walk over weights
1734        // the loader has already sharded across devices.
1735        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1736            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1737                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1738            }
1739            return self.forward_hyper_ppn(e, tokens, last_only, &topology, &fence);
1740        }
1741        let n_embd = self.cfg.n_embd as usize;
1742        let t = tokens.len();
1743        let eps = self.cfg.rms_eps;
1744        let pos: Vec<i32> = (0..t as i32).collect();
1745        let pos_d = e.htod_i32(&pos)?;
1746
1747        let embedded = self.embed(e, tokens)?;
1748        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
1749        let trace = memra_reference::hidden_trace::enabled();
1750        if trace {
1751            memra_reference::hidden_trace::emit_tokens(tokens);
1752            let streams = x.len() / (t * n_embd);
1753            memra_reference::hidden_trace::emit_last_row(
1754                "expand",
1755                -1,
1756                t,
1757                streams * n_embd,
1758                &e.dtoh(&x)?,
1759            );
1760        }
1761
1762        x = self.hyper_range_forward(e, &topology, x, 0, self.layers.len(), &pos_d, t, trace)?;
1763
1764        // SHARED EXIT with the ppN twin: one trunk exit, so the split and unsplit arms cannot
1765        // drift apart in the head. That is what makes `glm5-hyper-ppn-gate`'s bit-identity bar
1766        // a structural property rather than a coincidence of two maintained copies.
1767        self.hyper_head_logits(e, &topology, &x, t, n_embd, eps, last_only)
1768    }
1769
1770    /// Stateful prefill under the mHC residual: `prime_cache_overlaid`'s contract (leave a
1771    /// decode-ready cache behind, return last-row logits + the pre-output_norm hidden seed and
1772    /// stack) over the hc layer program.
1773    ///
1774    /// Deliberately UNCHUNKED and UNCAPTURED. The serial prime's chunking, prime slabs, S-mid
1775    /// graph capture and core-split arms are all keyed to the serial residual's transient set;
1776    /// re-deriving them for a stream state is a tuning lane, not a correctness one, and this
1777    /// path is the one the reference gate pins. Long prompts therefore hold `T*streams*hidden`
1778    /// f32 of stream state — 4x the serial trunk's — and that ceiling is the named cost of the
1779    /// simple form.
1780    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1781    fn prime_cache_hyper(
1782        &self,
1783        e: &Engine,
1784        tokens: &[u32],
1785        cache: &mut Cache,
1786        queued_after: usize,
1787        overlay: Option<&crate::vision::EmbedOverlay>,
1788    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1789        let topology = *self
1790            .hyper
1791            .as_ref()
1792            .ok_or("prime_cache_hyper on a model with no HyperConnections topology")?;
1793        // M2 ppN door for the mHC prime (see `forward_hyper`'s note). glm5_next has no
1794        // batched prime arm: the split twin is a straight layer-range split with the
1795        // [t, streams, hidden] state on the wire.
1796        //
1797        // CHUNKED, like the single-engine walk below (lane/glm53-1m-demo, 2026-08-29 — the
1798        // follow-up the previous note named). The monolithic ppN prime carried the WHOLE
1799        // prompt as one call, which capped it three independent ways on the 4x96 GB box:
1800        // per-call transients proportional to t OOM'd from ~32k tokens, and every launch
1801        // that places t in grid.y (kda_conv_silu, kda_gate, rms_norm over rows, the router)
1802        // hits the CUDA 65,535 grid.y ceiling from t=65,536 (measured: instant
1803        // CUDA_ERROR_INVALID_VALUE at a 128,566-token prime, receipts in
1804        // research/glm53-flash-bringup-20260827/1m-demo-20260829/). The chunk loop reuses
1805        // the SAME schedule as the single-engine walk (`hyper_prime_ranges`), so per-chunk
1806        // t is bounded and — because the per-chunk staged walk is bit-identical to the
1807        // per-chunk unsplit walk (glm5_hyper_ppn_gate arm 2) — the chunked ppN prime
1808        // composes to bit-identity with the chunked single-engine prime over the same
1809        // schedule. `queued_after + (t - end)` keeps the REQUEST-level `seq_end` invariant
1810        // across chunks (each call recomputes pos0+start + (end-start) + rest = pos0 + t +
1811        // queued_after). A prompt at or under one chunk takes the monolithic ppN body
1812        // unchanged, and `MEMRA_PRIME_CHUNK=0` restores the monolithic walk up to
1813        // PRIME_CHUNK_LAUNCH_CAP (grid.y-legal ceiling; above it the old walk always died
1814        // at launch) — the same rollback seam the single-engine chunk walk documents.
1815        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1816            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1817                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1818            }
1819            // The mixed-embedding overlay rides the ppN twin too (lane/glm5-vision-default-on,
1820            // 2026-08-30): the splice is an EMBEDDING-INTAKE transform, and the ppN walk embeds
1821            // on stage 0 only — every later stage receives the already-expanded stream state.
1822            // Under the chunked ppN prime each chunk takes the overlay WINDOWED to its own
1823            // call-relative range (`EmbedOverlay::window`, the same rebase seam the serve
1824            // prefill tick uses), so splice placement is chunk-schedule-invariant. Gated by
1825            // glm5-hyper-ppn-gate's overlay arm (bit-identity vs the substituted-token truth,
1826            // red arm = shifted spans).
1827            let n_embd = self.cfg.n_embd as usize;
1828            let t = tokens.len();
1829            if cache.pos + t > cache.max_ctx {
1830                return Err("prime_cache: prompt exceeds cache max_ctx".into());
1831            }
1832            let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1833            if ranges.len() == 1 {
1834                return self.prime_cache_hyper_ppn(
1835                    e,
1836                    tokens,
1837                    cache,
1838                    queued_after,
1839                    &topology,
1840                    &fence,
1841                    overlay,
1842                );
1843            }
1844            let mut hiddens = e.uninit(t * n_embd)?;
1845            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1846            for &(start, end) in &ranges {
1847                let ov = overlay.and_then(|o| o.window(start, end - start));
1848                let (l, hs, x) = self.prime_cache_hyper_ppn(
1849                    e,
1850                    &tokens[start..end],
1851                    cache,
1852                    queued_after + (t - end),
1853                    &topology,
1854                    &fence,
1855                    ov.as_ref(),
1856                )?;
1857                e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1858                last = Some((l, hs));
1859            }
1860            let (logits, h_seed) =
1861                last.expect("hyper_prime_ranges never returns an empty schedule");
1862            return Ok((logits, h_seed, hiddens));
1863        }
1864        let n_embd = self.cfg.n_embd as usize;
1865        let t = tokens.len();
1866        if cache.pos + t > cache.max_ctx {
1867            return Err("prime_cache: prompt exceeds cache max_ctx".into());
1868        }
1869        // The REQUEST's absolute end position, computed ONCE before the walk: every chunk sees
1870        // the same value whatever the chunk size, which is the tick-seg law the serial loop
1871        // above carries verbatim (`+ queued_after` closes the serve-split axis).
1872        let seq_end = cache.pos + t + queued_after;
1873        let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1874        if ranges.len() == 1 {
1875            return self.prime_chunk_hyper(e, tokens, cache, seq_end, 0, overlay);
1876        }
1877        let mut hiddens = e.uninit(t * n_embd)?;
1878        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1879        for &(start, end) in &ranges {
1880            let (l, hs, x) =
1881                self.prime_chunk_hyper(e, &tokens[start..end], cache, seq_end, start, overlay)?;
1882            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1883            last = Some((l, hs));
1884        }
1885        let (logits, h_seed) = last.expect("hyper_prime_ranges never returns an empty schedule");
1886        Ok((logits, h_seed, hiddens))
1887    }
1888
1889    /// Publish this model's prefill-workspace coefficients to admission
1890    /// (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper trunk — the server's
1891    /// admission arithmetic is then byte-identical to the pre-lane behavior for that family.
1892    ///
1893    /// Term-by-term, each anchored on a named allocation of the hyper prime walk (glm5 numbers
1894    /// in parentheses; the full attribution table is the lane doc's §1):
1895    ///   * stream state, double-buffered at `hyper::post`, PLUS the ppN boundary tx/rx pair of
1896    ///     the same `[t, streams, hidden]` payload: `4 * S * H * 4` (256 KiB/token);
1897    ///   * pre/norm transients (`hyper::pre` y + `rms_norm` h/z + ffn_out): `4 * H * 4`
1898    ///     (64 KiB/token);
1899    ///   * prime-tail hidden/norm pair (`collapse` + output-norm stack): `2 * H * 4`
1900    ///     (32 KiB/token);
1901    ///   * grouped-MoE prefill staging (`moe_ffn_grouped_prefill_sigmoid`): the f16 CSR
1902    ///     activations `U*2H`, three f32 partial planes `3*U*4F` (gate/up/act), the f16 down
1903    ///     mirror `U*2F`, the CSR-order down output + pair-order permute `2*U*4H`, and the
1904    ///     scatter target `4H` — `U*(10H + 14F) + 4H` (560 KiB/token);
1905    ///   * MLA query/attention planes: `heads * (qk_head_dim + v_head_dim) * 4`
1906    ///     (128 KiB/token) and the k-pool idx plane `(topk/P + 1) * 4` (~2 KiB/token).
1907    ///
1908    /// Validation against truth: at GLM-5.3-Flash geometry the sum is ~0.92 MiB per call
1909    /// token, against the 262k cell's MEASURED retained slope of ~0.8 MiB/token/card
1910    /// (vramwatch.csv: +6.3 GiB across the 8,072-token prime) — the formula sits above the
1911    /// measurement, never below it. The ctx-coupled score plane and the prompt-long hiddens
1912    /// stack are separate coefficients on the shape; see [`HyperPrimeWorkspaceShape`].
1913    pub fn hyper_prime_workspace_shape(&self) -> Option<HyperPrimeWorkspaceShape> {
1914        let topology = self.hyper.as_ref()?;
1915        let h = self.cfg.n_embd as usize;
1916        let s = topology.streams;
1917        let f32b = std::mem::size_of::<f32>();
1918        // Stream state (x2) + ppN boundary slots (x2), pre/norm transients, prime tail.
1919        let mut chunk_token_bytes = 4 * s * h * f32b + 4 * h * f32b + 2 * h * f32b;
1920        if let Some(moe) = self.cfg.moe.as_ref() {
1921            let u = moe.expert_used_count as usize;
1922            let f = moe.expert_ff_length as usize;
1923            chunk_token_bytes += u * (10 * h + 14 * f) + 4 * h;
1924        }
1925        let mut kpool_score_pool = 0;
1926        if let Some(glm5) = self.cfg.glm5.as_ref() {
1927            let heads = self.cfg.n_head as usize;
1928            chunk_token_bytes +=
1929                heads * (glm5.qk_head_dim as usize + glm5.v_head_dim as usize) * f32b;
1930            if glm5.index_kpool > 0 {
1931                chunk_token_bytes += (glm5.index_topk as usize / glm5.index_kpool as usize + 1)
1932                    * std::mem::size_of::<i32>();
1933                kpool_score_pool = glm5.index_kpool as usize;
1934            }
1935        }
1936        Some(HyperPrimeWorkspaceShape {
1937            chunk_token_bytes,
1938            prompt_bytes_per_token: h * f32b,
1939            kpool_score_pool,
1940            n_layers: self.layers.len(),
1941            gdn_grid: self.gdn_prime_grid_on(),
1942        })
1943    }
1944
1945    /// One T=1 decode step under the mHC residual: `decode_step_h`'s contract over the hc layer
1946    /// program. The stream state is INTRA-STEP — expanded from the embedded row, collapsed for
1947    /// the logits — so no cache format changes and the mixers keep their own state exactly as
1948    /// they do on the serial path.
1949    pub(crate) fn decode_step_hyper(
1950        &self,
1951        e: &Engine,
1952        token: u32,
1953        cache: &mut Cache,
1954    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1955        let topology = *self
1956            .hyper
1957            .as_ref()
1958            .ok_or("decode_step_hyper on a model with no HyperConnections topology")?;
1959        // M2 ppN door for the mHC decode step (see `forward_hyper`'s note). This is the door
1960        // the GLM-5.3-Flash residency arc turns on: with it shut, every routed expert has to
1961        // fit beside card 0's trunk, which 171.2 GB of experts cannot do.
1962        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1963            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1964                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1965            }
1966            return self.decode_step_hyper_ppn(e, token, cache, &topology, &fence);
1967        }
1968        let n_embd = self.cfg.n_embd as usize;
1969        let eps = self.cfg.rms_eps;
1970        let pos = cache.pos;
1971        let pos_d = e.htod_i32(&[pos as i32])?;
1972
1973        let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
1974        let mut x = crate::hyper::expand(e, &topology, &embedded, 1, n_embd)?;
1975
1976        x = self.hyper_range_decode(e, &topology, x, 0, self.layers.len(), &pos_d, pos, cache)?;
1977
1978        // SHARED EXIT with the ppN twin (see `forward_hyper`'s note).
1979        self.hyper_decode_tail(e, &topology, &x, n_embd, eps, cache)
1980    }
1981
1982    /// One hc layer RANGE `[lo, hi)` of the STATELESS prefill walk, driven by engine `e`.
1983    ///
1984    /// Extracted so the unsplit walk and every pipeline stage run the SAME code over their own
1985    /// range: the ppN arm's bit-identity claim is then structural, not a coincidence of two
1986    /// hand-kept-in-sync copies. `x` enters and leaves as the `[t, streams, hidden]` stream
1987    /// state, which is exactly the payload a stage boundary transports.
1988    #[allow(clippy::too_many_arguments)]
1989    fn hyper_range_forward(
1990        &self,
1991        e: &Engine,
1992        topology: &crate::hyper::HyperTopology,
1993        mut x: CudaSlice<f32>,
1994        lo: usize,
1995        hi: usize,
1996        pos_d: &CudaSlice<i32>,
1997        t: usize,
1998        trace: bool,
1999    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2000        let n_embd = self.cfg.n_embd as usize;
2001        let eps = self.cfg.rms_eps;
2002        for il in lo..hi {
2003            let layer = &self.layers[il];
2004            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2005                format!("layer {il} carries no hyper-connection weights under an hc plan")
2006            })?;
2007
2008            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2009            let mut h = e.uninit(t * n_embd)?;
2010            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2011            let mixed = match &layer.mixer {
2012                Mixer::Full(fa) => self.full_attn(e, fa, &h, pos_d, t, il)?,
2013                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
2014                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, pos_d, t, il)?,
2015                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
2016            };
2017            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2018            if trace {
2019                let index = il as i64;
2020                memra_reference::hidden_trace::emit_last_row(
2021                    "mixer",
2022                    index,
2023                    t,
2024                    n_embd,
2025                    &e.dtoh(&mixed)?,
2026                );
2027                let streams = x.len() / (t * n_embd);
2028                memra_reference::hidden_trace::emit_last_row(
2029                    "attn",
2030                    index,
2031                    t,
2032                    streams * n_embd,
2033                    &e.dtoh(&x)?,
2034                );
2035            }
2036
2037            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2038            let mut z = e.uninit(t * n_embd)?;
2039            e.rms_norm(
2040                &y,
2041                layer.post_attn_norm.float_data(),
2042                &mut z,
2043                n_embd,
2044                t,
2045                eps,
2046            )?;
2047            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2048            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2049            if trace {
2050                let index = il as i64;
2051                memra_reference::hidden_trace::emit_last_row(
2052                    "ffn",
2053                    index,
2054                    t,
2055                    n_embd,
2056                    &e.dtoh(&ffn_out)?,
2057                );
2058                let streams = x.len() / (t * n_embd);
2059                memra_reference::hidden_trace::emit_last_row(
2060                    "layer",
2061                    index,
2062                    t,
2063                    streams * n_embd,
2064                    &e.dtoh(&x)?,
2065                );
2066            }
2067        }
2068        Ok(x)
2069    }
2070
2071    /// One hc layer RANGE `[lo, hi)` of the STATEFUL prime walk (see `hyper_range_forward`).
2072    /// Every mixer writes its own layer's cache state through `e`, so under the ppN door a
2073    /// stage's KDA conv ring / delta-rule state and its MLA latent rows + kpool indexer plane
2074    /// are written by the SAME engine `pp::new_cache` allocated them on.
2075    #[allow(clippy::too_many_arguments)]
2076    fn hyper_range_prime(
2077        &self,
2078        e: &Engine,
2079        topology: &crate::hyper::HyperTopology,
2080        mut x: CudaSlice<f32>,
2081        lo: usize,
2082        hi: usize,
2083        pos_d: &CudaSlice<i32>,
2084        t: usize,
2085        cache: &mut Cache,
2086        seq_end: usize,
2087    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2088        let n_embd = self.cfg.n_embd as usize;
2089        let eps = self.cfg.rms_eps;
2090        for il in lo..hi {
2091            let layer = &self.layers[il];
2092            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2093                format!("layer {il} carries no hyper-connection weights under an hc plan")
2094            })?;
2095
2096            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2097            let mut h = e.uninit(t * n_embd)?;
2098            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2099            let mixed = match &layer.mixer {
2100                Mixer::Full(fa) => {
2101                    self.full_attn_prime(e, fa, &h, None, pos_d, t, cache, il, seq_end)?
2102                }
2103                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2104                Mixer::Mla(mla) if mla.tp.is_some() => {
2105                    self.mla_tp_attn_cached(e, mla, &h, pos_d, t, il, cache, false)?
2106                }
2107                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, t, il, cache)?,
2108                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2109                    e,
2110                    la,
2111                    &h,
2112                    t,
2113                    eps,
2114                    cache,
2115                    il,
2116                    crate::kda::ConvArm::Prefill,
2117                )?,
2118                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2119            };
2120            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2121
2122            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2123            let mut z = e.uninit(t * n_embd)?;
2124            e.rms_norm(
2125                &y,
2126                layer.post_attn_norm.float_data(),
2127                &mut z,
2128                n_embd,
2129                t,
2130                eps,
2131            )?;
2132            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2133            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2134            // glm5 DFlash2 feature tap (lane/glm5-dflash-draft-src): the CONTRACTED
2135            // completed layer output, host-staged — one Option check when unarmed.
2136            self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
2137        }
2138        Ok(x)
2139    }
2140
2141    /// One hc layer RANGE `[lo, hi)` of the T=1 decode step (see `hyper_range_forward`).
2142    #[allow(clippy::too_many_arguments)]
2143    fn hyper_range_decode(
2144        &self,
2145        e: &Engine,
2146        topology: &crate::hyper::HyperTopology,
2147        mut x: CudaSlice<f32>,
2148        lo: usize,
2149        hi: usize,
2150        pos_d: &CudaSlice<i32>,
2151        pos: usize,
2152        cache: &mut Cache,
2153    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2154        // MEMRA_HC_DECODE_WS=1 (default OFF, read per call — rollback seam): the persistent-
2155        // workspace twin of this walk. Same kernels, same call order, same operand bytes;
2156        // only the hc-glue allocations (mixes/gates/comb/y/h/z/post-out, ~12 alloc+free pairs
2157        // per layer per token of the census's 2,358) disappear. Byte identity ON/OFF is gated
2158        // by hc_decode_ws_gpu.rs; refusal shapes fall through to the allocating walk below.
2159        if hyper_decode_ws_on() {
2160            return self.hyper_range_decode_ws(e, topology, x, lo, hi, pos_d, pos, cache);
2161        }
2162        let n_embd = self.cfg.n_embd as usize;
2163        let eps = self.cfg.rms_eps;
2164        for il in lo..hi {
2165            let layer = &self.layers[il];
2166            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2167                format!("layer {il} carries no hyper-connection weights under an hc plan")
2168            })?;
2169
2170            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, 1, n_embd)?;
2171            let mut h = e.uninit(n_embd)?;
2172            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
2173            let mixed = match &layer.mixer {
2174                Mixer::Full(fa) => self.full_attn_decode(e, fa, &h, pos_d, pos, cache, il)?,
2175                Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il)?,
2176                Mixer::Mla(mla) if mla.tp.is_some() => {
2177                    self.mla_tp_attn_cached(e, mla, &h, pos_d, 1, il, cache, false)?
2178                }
2179                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache)?,
2180                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2181                    e,
2182                    la,
2183                    &h,
2184                    1,
2185                    eps,
2186                    cache,
2187                    il,
2188                    crate::kda::ConvArm::Decode,
2189                )?,
2190                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?,
2191            };
2192            x = crate::hyper::post(e, topology, &mixed, &x, &mix, 1, n_embd)?;
2193
2194            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, 1, n_embd)?;
2195            let mut z = e.uninit(n_embd)?;
2196            e.rms_norm(
2197                &y,
2198                layer.post_attn_norm.float_data(),
2199                &mut z,
2200                n_embd,
2201                1,
2202                eps,
2203            )?;
2204            let ffn_out = self.hyper_ffn_branch(e, layer, &z, 1, il, false)?;
2205            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, 1, n_embd)?;
2206        }
2207        Ok(x)
2208    }
2209
2210    /// The persistent-workspace twin of `hyper_range_decode` (MEMRA_HC_DECODE_WS, lever 2 of
2211    /// the decode diet). One `HyperDecodeWs` per engine (so each ppN stage owns its own,
2212    /// allocated on its own device); the walk TAKES it from the engine pool, rotates the
2213    /// stream state against `ws.xb` (an ownership swap, not a copy), and puts it back. The
2214    /// mixers and the FFN/MoE branches are the SAME calls with the SAME inputs — their
2215    /// internal allocations are untouched by this lever.
2216    #[allow(clippy::too_many_arguments)]
2217    fn hyper_range_decode_ws(
2218        &self,
2219        e: &Engine,
2220        topology: &crate::hyper::HyperTopology,
2221        x: CudaSlice<f32>,
2222        lo: usize,
2223        hi: usize,
2224        pos_d: &CudaSlice<i32>,
2225        pos: usize,
2226        cache: &mut Cache,
2227    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2228        let n_embd = self.cfg.n_embd as usize;
2229        let mut ws = match e.hyper_ws_take() {
2230            Some(ws) if ws.matches(topology, n_embd) => ws,
2231            _ => crate::hyper::HyperDecodeWs::new(e, topology, n_embd)?,
2232        };
2233        if HC_DECODE_WS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2234            eprintln!(
2235                "[hc-decode-ws] engaged streams={} hidden={n_embd} (persistent hc-glue \
2236                 workspace, per-engine pool; MEMRA_HC_DECODE_WS=1)",
2237                topology.streams
2238            );
2239        }
2240        let out =
2241            self.hyper_range_decode_ws_body(e, topology, x, lo, hi, pos_d, pos, cache, &mut ws);
2242        e.hyper_ws_put(ws);
2243        out
2244    }
2245
2246    /// The walk itself — `hyper_range_decode`'s loop with the hc glue landing in `ws`.
2247    /// KEPT CALL-FOR-CALL IN STEP with the allocating walk above: same kernels, same order
2248    /// (pre -> rms_norm -> mixer -> post -> pre -> rms_norm -> ffn -> post), so the
2249    /// byte-identity gate is a structural claim, not a coincidence.
2250    #[allow(clippy::too_many_arguments)]
2251    fn hyper_range_decode_ws_body(
2252        &self,
2253        e: &Engine,
2254        topology: &crate::hyper::HyperTopology,
2255        mut x: CudaSlice<f32>,
2256        lo: usize,
2257        hi: usize,
2258        pos_d: &CudaSlice<i32>,
2259        pos: usize,
2260        cache: &mut Cache,
2261        ws: &mut crate::hyper::HyperDecodeWs,
2262    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2263        let n_embd = self.cfg.n_embd as usize;
2264        let eps = self.cfg.rms_eps;
2265        for il in lo..hi {
2266            let layer = &self.layers[il];
2267            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2268                format!("layer {il} carries no hyper-connection weights under an hc plan")
2269            })?;
2270
2271            crate::hyper::pre_t1_ws(e, topology, &hyper.attn, &x, ws, n_embd)?;
2272            e.rms_norm(
2273                &ws.y,
2274                layer.attn_norm.float_data(),
2275                &mut ws.h,
2276                n_embd,
2277                1,
2278                eps,
2279            )?;
2280            let mixed = match &layer.mixer {
2281                Mixer::Full(fa) => self.full_attn_decode(e, fa, &ws.h, pos_d, pos, cache, il)?,
2282                Mixer::Linear(la) => self.linear_attn_decode(e, la, &ws.h, cache, il)?,
2283                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &ws.h, pos_d, 1, il, cache)?,
2284                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &ws.h, eps, cache, il)?,
2285            };
2286            crate::hyper::post_t1_ws(e, topology, &mixed, &x, ws, n_embd)?;
2287            std::mem::swap(&mut x, &mut ws.xb);
2288
2289            crate::hyper::pre_t1_ws(e, topology, &hyper.mlp, &x, ws, n_embd)?;
2290            e.rms_norm(
2291                &ws.y,
2292                layer.post_attn_norm.float_data(),
2293                &mut ws.z,
2294                n_embd,
2295                1,
2296                eps,
2297            )?;
2298            let ffn_out = self.hyper_ffn_branch(e, layer, &ws.z, 1, il, false)?;
2299            crate::hyper::post_t1_ws(e, topology, &ffn_out, &x, ws, n_embd)?;
2300            std::mem::swap(&mut x, &mut ws.xb);
2301        }
2302        Ok(x)
2303    }
2304
2305    /// One hc layer RANGE `[lo, hi)` of the BATCHED T=1 decode step: B independent sessions
2306    /// share one walk over the `[B, streams, n_embd]` stream state. The batched twin of
2307    /// `hyper_range_decode`, and the trunk of `decode_step_batch_hyper` (decode_batch.rs).
2308    ///
2309    /// SHAPE — batched where the arithmetic is row-independent, per-session where the state
2310    /// is, decode-exact where a reduction is width-dependent:
2311    ///
2312    ///   * The hc glue (`expand`/`pre_finish` kernels/`post`) is block-per-token by
2313    ///     construction (grid over t), so t=B batches it with per-row bytes unchanged.
2314    ///   * The hc mixing GEMM is the ONE width-dependent reduction in the glue
2315    ///     (cuBLASLt's n-dependent split — the lt_ndep probe), so this walk calls
2316    ///     `hyper::pre_exact`, which runs each row through the m=1 program the serial step
2317    ///     runs. rms_norm at m=B is a per-row program.
2318    ///   * The MIXERS (KDA conv ring + delta rule, MLA latent rows + kpool indexer plane,
2319    ///     and the Full/Linear classes for completeness) hold per-session recurrent or
2320    ///     latent state, so each session's row is routed to ITS OWN cache through the SAME
2321    ///     t=1 call its solo step makes — the per-seq loop is the v1 exactness doctrine
2322    ///     from this module's sibling (`decode_batch.rs` header), and the row copies in and
2323    ///     out are arithmetic-free materializations.
2324    ///   * The FFN batches at t=B: the MoE body's router is the fixed per-row program at
2325    ///     t < PRIME_MIN_T, expert dispatch is per-token, and the shexp trio rides the
2326    ///     per-column decode-exact arm at decode widths; the dense branch runs per-row so
2327    ///     each row executes the serial `hyper_ffn_branch` program verbatim.
2328    ///
2329    /// EXACTNESS BAR: row b of a B-row step must be BIT-IDENTICAL to session b decoding
2330    /// alone through `decode_step_hyper` — full-logit compare, per step. Gate:
2331    /// `glm5-hyper-batch-gate` (fixture-driven, red-armed with a swapped-row and a
2332    /// wrong-cache-slot mutation; receipts in
2333    /// `research/glm53-flash-bringup-20260827/batched-decode-gate/`).
2334    ///
2335    /// `pos_rows[bi]` is session bi's single-position device buffer, uploaded by the caller
2336    /// through THIS range's engine (the per-stage pos_d law under a pp split). `caches[bi]`
2337    /// advances exactly as its solo step would; `cache.pos` itself is bumped by the caller's
2338    /// epilogue, never here.
2339    #[allow(clippy::too_many_arguments)]
2340    pub(crate) fn hyper_batch_range_decode(
2341        &self,
2342        e: &Engine,
2343        topology: &crate::hyper::HyperTopology,
2344        mut x: CudaSlice<f32>,
2345        lo: usize,
2346        hi: usize,
2347        pos_rows: &[CudaSlice<i32>],
2348        caches: &mut [&mut Cache],
2349    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2350        let b_n = caches.len();
2351        assert_eq!(
2352            pos_rows.len(),
2353            b_n,
2354            "hyper_batch_range_decode: pos_rows built for a different batch width"
2355        );
2356        let n_embd = self.cfg.n_embd as usize;
2357        let eps = self.cfg.rms_eps;
2358        for il in lo..hi {
2359            let layer = &self.layers[il];
2360            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2361                format!("layer {il} carries no hyper-connection weights under an hc plan")
2362            })?;
2363
2364            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, b_n, n_embd)?;
2365            let mut h = e.uninit(b_n * n_embd)?;
2366            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, b_n, eps)?;
2367            // ---- mixers: per-session, each row through its OWN cache (the t=1 serial
2368            // program — cross-session contamination here is the failure mode the gate's
2369            // swapped-row mutation exists to catch) ----
2370            let mut mixed = e.uninit(b_n * n_embd)?;
2371            for bi in 0..b_n {
2372                let mut h_row = e.uninit(n_embd)?;
2373                e.dtod_copy_view(&h.slice(bi * n_embd..(bi + 1) * n_embd), &mut h_row)?;
2374                let cache: &mut Cache = &mut *caches[bi];
2375                let pos = cache.pos;
2376                let out_row = match &layer.mixer {
2377                    Mixer::Full(fa) => {
2378                        self.full_attn_decode(e, fa, &h_row, &pos_rows[bi], pos, cache, il)?
2379                    }
2380                    Mixer::Linear(la) => self.linear_attn_decode(e, la, &h_row, cache, il)?,
2381                    Mixer::Mla(mla) => {
2382                        self.mla_attn_cached(e, mla, &h_row, &pos_rows[bi], 1, il, cache)?
2383                    }
2384                    Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?,
2385                };
2386                e.copy_into(&mut mixed, bi * n_embd, &out_row, n_embd)?;
2387            }
2388            x = crate::hyper::post(e, topology, &mixed, &x, &mix, b_n, n_embd)?;
2389
2390            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, b_n, n_embd)?;
2391            let mut z = e.uninit(b_n * n_embd)?;
2392            e.rms_norm(
2393                &y,
2394                layer.post_attn_norm.float_data(),
2395                &mut z,
2396                n_embd,
2397                b_n,
2398                eps,
2399            )?;
2400            let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, b_n, il, false)?;
2401            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, b_n, n_embd)?;
2402        }
2403        Ok(x)
2404    }
2405
2406    /// The FFN branch of the batched hc decode walk (see `hyper_batch_range_decode`).
2407    ///
2408    /// MoE batches at t=B — that is the weight win this walk exists for (one stream of the
2409    /// routed experts serves B rows), and every stage of `moe_ffn_il_zq8` is per-row exact
2410    /// at decode widths (router: fixed per-row program at t < PRIME_MIN_T; experts:
2411    /// per-(token,expert) programs; shexp: the per-column decode-exact arm). The DENSE
2412    /// branch runs PER ROW through the serial `hyper_ffn_branch` instead: its
2413    /// `matmul_group` dispatch carries no per-row bit-identity contract across widths for
2414    /// every weight class this walk must serve, and a first-k-dense plan carries one such
2415    /// layer — per-row costs nothing and each row executes the solo step's program verbatim.
2416    /// `vrows` (lane/glm5-vrest): the VERIFY walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`)
2417    /// passes `true`, which lets the MoE body take the pairs-shaped batched routed-expert
2418    /// program across the t rows (`moe_vrows_pairs_q8` — bit-identical per row, fail-closed
2419    /// to the sequential loop for every unqualified shape). The batched DECODE walk
2420    /// (`decode_step_batch_hyper`) passes `false` — its priced dispatch class stays
2421    /// byte-stable; porting it is a named follow-up with its own re-price. The DENSE branch
2422    /// is per-row in both arms (its `matmul_group` dispatch carries no cross-width per-row
2423    /// bit-identity contract for every weight class this walk must serve; ~3 layers, named
2424    /// out of scope in the vrest attribution).
2425    pub(crate) fn hyper_ffn_branch_batch(
2426        &self,
2427        e: &Engine,
2428        layer: &crate::hybrid::HybridLayer,
2429        z: &CudaSlice<f32>,
2430        b_n: usize,
2431        il: usize,
2432        vrows: bool,
2433    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2434        let n_embd = self.cfg.n_embd as usize;
2435        match &layer.ffn {
2436            crate::hybrid::Ffn::Dense { .. } => {
2437                let mut out = e.uninit(b_n * n_embd)?;
2438                for bi in 0..b_n {
2439                    let mut z_row = e.uninit(n_embd)?;
2440                    e.dtod_copy_view(&z.slice(bi * n_embd..(bi + 1) * n_embd), &mut z_row)?;
2441                    let row = self.hyper_ffn_branch(e, layer, &z_row, 1, il, false)?;
2442                    e.copy_into(&mut out, bi * n_embd, &row, n_embd)?;
2443                }
2444                Ok(out)
2445            }
2446            crate::hybrid::Ffn::Moe(m) => {
2447                if vrows {
2448                    self.moe_ffn_il_zq8_vrows(e, m, z, b_n, il as u16)
2449                } else {
2450                    self.moe_ffn_il_zq8(e, m, z, None, b_n, il as u16)
2451                }
2452            }
2453        }
2454    }
2455
2456    // =============================== M2 ppN, mHC arm ===============================
2457    //
2458    // The three walks below are the hc twins of `decode_step_h_ppn`. They exist because the
2459    // GLM-5.3-Flash residency arithmetic does not close on one card: 171.2 GB of routed
2460    // experts against 2x96 GB means the second card is the only route to full residency, and
2461    // the pp door is how weights get there. Until these landed, all three hc walks refused
2462    // the door outright ("the sharded stage handoff is unwired for this residual topology"),
2463    // which is a loud refusal and was the right behaviour — a single-engine walk over
2464    // stage-sharded weights dereferences another device's pointers.
2465    //
2466    // WHAT IS DIFFERENT FROM THE GENERIC ARM, and it is exactly one thing: the payload on the
2467    // wire. The serial trunk hands `[n_embd]` (decode) or `[t, n_embd]` (prime) across a
2468    // boundary; the mHC trunk carries `streams` residual streams between layers, so the
2469    // boundary payload is `[streams, n_embd]` / `[t, streams, n_embd]`. `pp.rs`'s BoundarySlot
2470    // buffers are lazily sized from the caller's `n` and grow to the high-water mark, so no
2471    // slot-sizing change was needed for that — the wider payload just makes them wider.
2472    // `hyper::expand` runs on stage 0 (it takes no weights) and `hyper::collapse` +
2473    // output_norm + lm head on the last stage, which is where the loader already put the head
2474    // (`pp::layer_engine(e, n_trunk, n_trunk - 1)` in `hybrid.rs`) and, under
2475    // `HcCollapse::GatedHead`, the head trio.
2476    //
2477    // Per-layer state placement needed NO new contract: `pp::new_cache` already picks the
2478    // owning stage's `KvDev` per layer for all three of glm5_next's state classes
2479    // (`Recurrent` = KDA conv ring + delta-rule state, `LatentKvCache` = MLA rows + the kpool
2480    // indexer plane, `KvCache` = full attention), and the kpool `index_pool_keys` plane is
2481    // lazily allocated through the engine the mixer is called with, which under these walks is
2482    // the stage's engine. `glm5-hyper-ppn-gate` asserts the fence actually separates those
2483    // classes across stages, so that is a tested property rather than an argued one.
2484    //
2485    // NOT WIRED, and refused rather than approximated: the deferred-readback (pipelined) arm.
2486    // `decode_step_h_ppn_deferred` calls `refuse_hyper`, and this lane did not change that.
2487    //
2488    // Gate: `glm5-hyper-ppn-gate` (bit-identical logits vs the unsplit hc walk, decode and
2489    // prime, at every N/knob combination), receipts in
2490    // `research/glm53-flash-bringup-20260827/ppn-hyper-gate/`.
2491
2492    /// ppN twin of `forward_hyper`: the stateless prefill as N stage subgraphs.
2493    fn forward_hyper_ppn(
2494        &self,
2495        e: &Engine,
2496        tokens: &[u32],
2497        last_only: bool,
2498        topology: &crate::hyper::HyperTopology,
2499        fence: &[usize],
2500    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2501        let n_embd = self.cfg.n_embd as usize;
2502        let eps = self.cfg.rms_eps;
2503        let t = tokens.len();
2504        let width = topology.streams * n_embd;
2505        let trace = memra_reference::hidden_trace::enabled();
2506        if trace {
2507            memra_reference::hidden_trace::emit_tokens(tokens);
2508        }
2509        let pos: Vec<i32> = (0..t as i32).collect();
2510
2511        if crate::pp::pp2_streams_off() {
2512            // Same-stream rollback seam: one engine, one stream, an explicit copy pair per
2513            // boundary. Structurally identical to the split walk, which is the point — it is
2514            // the arm that says "the split is the split, not the streams".
2515            let pos_d = e.htod_i32(&pos)?;
2516            let embedded = self.embed(e, tokens)?;
2517            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2518            x = self.hyper_range_forward(e, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2519            for s in 1..fence.len() - 1 {
2520                let boundary_tx = e.clone_dtod(&x)?;
2521                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2522                x = self.hyper_range_forward(
2523                    e,
2524                    topology,
2525                    boundary_rx,
2526                    fence[s],
2527                    fence[s + 1],
2528                    &pos_d,
2529                    t,
2530                    trace,
2531                )?;
2532            }
2533            return self.hyper_head_logits(e, topology, &x, t, n_embd, eps, last_only);
2534        }
2535
2536        let rt = crate::pp::PpNRt::get(e)?;
2537        let n_st = fence.len() - 1;
2538        assert_eq!(
2539            rt.n_stages(),
2540            n_st,
2541            "PpNRt stage count {} != fence stages {n_st}",
2542            rt.n_stages()
2543        );
2544        // #87 REVERSE PUBLICATION: order every stage stream behind the caller's stream before
2545        // the first stage allocation (anatomy: `PpNRt::fence_stages_behind`).
2546        rt.fence_stages_behind(&e.stream())?;
2547
2548        let mut slot = {
2549            let _st0 = rt.enter(0);
2550            let e0 = rt.engine(0, e);
2551            // PER-STAGE pos_d (M2 pipelining law): each stage uploads its OWN copy on ITS
2552            // stream, so the buffer is allocated, consumed and freed on one stream.
2553            let pos_d = e0.htod_i32(&pos)?;
2554            let embedded = self.embed(e0, tokens)?;
2555            let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2556            let x =
2557                self.hyper_range_forward(e0, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2558            rt.tx(0, &x, t * width)?
2559        };
2560        for s in 1..n_st - 1 {
2561            let _st = rt.enter(s);
2562            let es = rt.engine(s, e);
2563            let pos_d = es.htod_i32(&pos)?;
2564            let x = rt.rx(s - 1, slot, t * width)?;
2565            let x = self.hyper_range_forward(
2566                es,
2567                topology,
2568                x,
2569                fence[s],
2570                fence[s + 1],
2571                &pos_d,
2572                t,
2573                trace,
2574            )?;
2575            slot = rt.tx(s, &x, t * width)?;
2576        }
2577        let _stl = rt.enter(n_st - 1);
2578        let el = rt.engine(n_st - 1, e);
2579        let pos_d = el.htod_i32(&pos)?;
2580        let x = rt.rx(n_st - 2, slot, t * width)?;
2581        let x = self.hyper_range_forward(
2582            el,
2583            topology,
2584            x,
2585            fence[n_st - 1],
2586            fence[n_st],
2587            &pos_d,
2588            t,
2589            trace,
2590        )?;
2591        self.hyper_head_logits(el, topology, &x, t, n_embd, eps, last_only)
2592    }
2593
2594    /// Trunk exit shared by `forward_hyper` and its ppN twin: collapse the stream state, apply
2595    /// `output_norm`, and project. Runs on the LAST stage's engine under the pp door, which is
2596    /// where `hybrid.rs` uploaded `output_norm`, the lm head and (under `HcCollapse::GatedHead`)
2597    /// the head trio.
2598    #[allow(clippy::too_many_arguments)]
2599    fn hyper_head_logits(
2600        &self,
2601        e: &Engine,
2602        topology: &crate::hyper::HyperTopology,
2603        x: &CudaSlice<f32>,
2604        t: usize,
2605        n_embd: usize,
2606        eps: f32,
2607        last_only: bool,
2608    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2609        let collapsed =
2610            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2611        if memra_reference::hidden_trace::enabled() {
2612            memra_reference::hidden_trace::emit_last_row(
2613                "collapse",
2614                -1,
2615                t,
2616                n_embd,
2617                &e.dtoh(&collapsed)?,
2618            );
2619        }
2620        let mut hn = e.uninit(t * n_embd)?;
2621        e.rms_norm(
2622            &collapsed,
2623            self.output_norm.float_data(),
2624            &mut hn,
2625            n_embd,
2626            t,
2627            eps,
2628        )?;
2629        let logits = if last_only {
2630            let last = e.view(&hn, t * n_embd);
2631            let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2632            let mut hlast = e.uninit(n_embd)?;
2633            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2634            e.matmul(&self.output, &hlast, 1)?
2635        } else {
2636            e.matmul(&self.output, &hn, t)?
2637        };
2638        e.dtoh(&logits)
2639    }
2640
2641    /// ppN twin of `prime_cache_hyper`: the monolithic stateful prime as N stage subgraphs.
2642    /// The returned device buffers (`h_seed`, `hiddens`) are owned by the LAST stage's engine,
2643    /// the same contract `decode_step_h_ppn` publishes for its `h_seed`.
2644    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2645    #[allow(clippy::too_many_arguments)]
2646    fn prime_cache_hyper_ppn(
2647        &self,
2648        e: &Engine,
2649        tokens: &[u32],
2650        cache: &mut Cache,
2651        queued_after: usize,
2652        topology: &crate::hyper::HyperTopology,
2653        fence: &[usize],
2654        overlay: Option<&crate::vision::EmbedOverlay>,
2655    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2656        let n_embd = self.cfg.n_embd as usize;
2657        let eps = self.cfg.rms_eps;
2658        let t = tokens.len();
2659        let width = topology.streams * n_embd;
2660        if cache.pos + t > cache.max_ctx {
2661            return Err("prime_cache: prompt exceeds cache max_ctx".into());
2662        }
2663        let seq_end = cache.pos + t + queued_after;
2664        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2665        // glm5 DFlash2 feature tap: set ONCE per call, before any stage walk — every stage
2666        // range of this chunk writes the same rows at the chunk's absolute offset.
2667        if let Some(sink) = cache.hc_taps.as_mut() {
2668            sink.base = cache.pos;
2669        }
2670
2671        if crate::pp::pp2_streams_off() {
2672            let pos_d = e.htod_i32(&pos)?;
2673            let mut embedded = self.embed(e, tokens)?;
2674            if let Some(ov) = overlay {
2675                // Mixed-embedding splice at embedding intake (the same point the streams-on
2676                // arm and prime_chunk_hyper use). The caller's overlay is already windowed
2677                // to THIS call (prefill_tick rebases spans call-relative), so chunk_off = 0.
2678                ov.splice_into(e, &mut embedded, 0, t, n_embd)?;
2679            }
2680            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2681            x = self.hyper_range_prime(
2682                e, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2683            )?;
2684            for s in 1..fence.len() - 1 {
2685                let boundary_tx = e.clone_dtod(&x)?;
2686                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2687                x = self.hyper_range_prime(
2688                    e,
2689                    topology,
2690                    boundary_rx,
2691                    fence[s],
2692                    fence[s + 1],
2693                    &pos_d,
2694                    t,
2695                    cache,
2696                    seq_end,
2697                )?;
2698            }
2699            return self.hyper_prime_tail(e, topology, &x, t, n_embd, eps, cache);
2700        }
2701
2702        {
2703            let rt = crate::pp::PpNRt::get(e)?;
2704            let n_st = fence.len() - 1;
2705            assert_eq!(
2706                rt.n_stages(),
2707                n_st,
2708                "PpNRt stage count {} != fence stages {n_st}",
2709                rt.n_stages()
2710            );
2711            let caller_stream = e.stream();
2712            rt.fence_stages_behind(&caller_stream)?;
2713            // OVERLAY RESIDENCY LAW (rewritten in lane/glm53-vision-ppn, 2026-09-01; was the
2714            // OVERLAY DEVICE LAW). The splice reads `overlay.rows` through stage 0's engine,
2715            // so those rows must live in STAGE 0's CUDA context. That is the real invariant,
2716            // and it is what is checked.
2717            //
2718            // The pre-lane check was `!std::ptr::eq(rt.engine(0, e), e)` — "stage 0 must BE
2719            // the primary engine". It refused the deployed 3-card shape outright: the worker's
2720            // primary engine follows the LAST pp stage (`worker::worker_device`, and that is
2721            // load-bearing — pinning the primary to stage 0 was the v0.72 tag-blocker-2
2722            // regressor, 112.5 -> 17.5 tok/s on spec+PP), so with MEMRA_PP_DEVICES=0,1,2 the
2723            // primary is dev2 while stage 0 owns dev0 and `PpNRt::build` hands stage 0 its
2724            // own Engine. Vision was therefore unservable on the ppN shape with the only
2725            // in-tag rollback costing ~3x decode (MEMRA_PP_STREAMS=0). The fix is to PUBLISH
2726            // the overlay into stage 0's context at construction
2727            // (`EmbedOverlay::new_published`, driven by `vision_intake_engine` below), never
2728            // to relax the check: the identity now passes because the pointers ARE in the
2729            // right domain.
2730            //
2731            // ORDERING, the seam the accrace lane taught (MEMRA_PP_EXIT_PUBLISH): rows built
2732            // on the caller's stream are covered by `fence_stages_behind` above, which orders
2733            // every stage stream behind the caller before stage 0 issues its splice; rows
2734            // published onto the intake engine were host-synchronized when they were uploaded.
2735            // Both producers are ordered before this body's first stage-0 kernel.
2736            if let Some(ov) = overlay
2737                && !ov.resident_in(rt.engine(0, e))
2738            {
2739                return Err(format!(
2740                    "vision embedding overlay rows are resident on dev{} but pp stage 0's \
2741                     embedding intake runs on dev{}: the overlay must be published into the \
2742                     intake engine's context (build it with EmbedOverlay::new_published; \
2743                     MEMRA_VISION_OVERLAY_PUBLISH=0 pins the pre-publication program, whose \
2744                     only vision-capable shape is MEMRA_PP_STREAMS=0)",
2745                    ov.ctx().ordinal(),
2746                    rt.engine(0, e).ctx().ordinal(),
2747                )
2748                .into());
2749            }
2750            let mut slot = {
2751                let _st0 = rt.enter(0);
2752                let e0 = rt.engine(0, e);
2753                let pos_d = e0.htod_i32(&pos)?;
2754                let mut embedded = self.embed(e0, tokens)?;
2755                if let Some(ov) = overlay {
2756                    // Mixed-embedding splice at stage-0 embedding intake, BEFORE stream
2757                    // expansion — the reference's execute_multimodal splice point. Stages
2758                    // s > 0 only ever see the [t, streams, hidden] boundary payload, so no
2759                    // other stage carries overlay arithmetic. chunk_off = 0: the caller's
2760                    // overlay is already windowed to this call.
2761                    ov.splice_into(e0, &mut embedded, 0, t, n_embd)?;
2762                }
2763                let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2764                let x = self.hyper_range_prime(
2765                    e0, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2766                )?;
2767                rt.tx(0, &x, t * width)?
2768            };
2769            for s in 1..n_st - 1 {
2770                let _st = rt.enter(s);
2771                let es = rt.engine(s, e);
2772                let pos_d = es.htod_i32(&pos)?;
2773                let x = rt.rx(s - 1, slot, t * width)?;
2774                let x = self.hyper_range_prime(
2775                    es,
2776                    topology,
2777                    x,
2778                    fence[s],
2779                    fence[s + 1],
2780                    &pos_d,
2781                    t,
2782                    cache,
2783                    seq_end,
2784                )?;
2785                slot = rt.tx(s, &x, t * width)?;
2786            }
2787            let out = {
2788                let _stl = rt.enter(n_st - 1);
2789                let el = rt.engine(n_st - 1, e);
2790                let pos_d = el.htod_i32(&pos)?;
2791                let x = rt.rx(n_st - 2, slot, t * width)?;
2792                let x = self.hyper_range_prime(
2793                    el,
2794                    topology,
2795                    x,
2796                    fence[n_st - 1],
2797                    fence[n_st],
2798                    &pos_d,
2799                    t,
2800                    cache,
2801                    seq_end,
2802                )?;
2803                self.hyper_prime_tail(el, topology, &x, t, n_embd, eps, cache)?
2804            };
2805            // EXIT PUBLICATION (lane/glm5-accrace 2026-09-01, the same law the batched and
2806            // spec ppN bodies carry): `hyper_prime_tail`'s dtoh drains the LAST stage only,
2807            // and the TX-wait chain reaches each earlier stage just as far as its `ev_tx`.
2808            // Every earlier stage's stream still holds the tail its stage-scope locals
2809            // enqueue on drop, and the caller resumes here to allocate (the glm5 spec
2810            // session's MTP plane warm, the worker's next step). Anatomy + receipts:
2811            // `PpNRt::publish_all_to`.
2812            rt.publish_all_to(&caller_stream)?;
2813            Ok(out)
2814        }
2815    }
2816
2817    /// ONE call of the mHC prime walk. Carries `tokens.len()` rows of stream state and of every
2818    /// per-layer transient, appends this call's rows to the mixers' own state, and advances
2819    /// `cache.pos`. `seq_end` is the REQUEST's absolute end, passed in rather than recomputed,
2820    /// so no arithmetic here is a function of how the prompt was split.
2821    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2822    fn prime_chunk_hyper(
2823        &self,
2824        e: &Engine,
2825        tokens: &[u32],
2826        cache: &mut Cache,
2827        seq_end: usize,
2828        chunk_off: usize,
2829        overlay: Option<&crate::vision::EmbedOverlay>,
2830    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2831        let topology = *self
2832            .hyper
2833            .as_ref()
2834            .ok_or("prime_chunk_hyper on a model with no HyperConnections topology")?;
2835        let n_embd = self.cfg.n_embd as usize;
2836        let t = tokens.len();
2837        let eps = self.cfg.rms_eps;
2838        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2839        let pos_d = e.htod_i32(&pos)?;
2840        // glm5 DFlash2 feature tap: chunked primes write their rows at the chunk's absolute
2841        // offset (cache.pos advances per chunk — the dflash_taps.base precedent).
2842        if let Some(sink) = cache.hc_taps.as_mut() {
2843            sink.base = cache.pos;
2844        }
2845
2846        let mut embedded = self.embed(e, tokens)?;
2847        if let Some(ov) = overlay {
2848            // Mixed-embedding splice (shared with the ppN twin — EmbedOverlay::splice_into):
2849            // image rows overwrite placeholder-token embeddings inside this chunk's
2850            // prompt-relative window [chunk_off, chunk_off+t), BEFORE stream expansion —
2851            // the reference's splice point (execute_multimodal replaces rows before
2852            // hc_expand).
2853            ov.splice_into(e, &mut embedded, chunk_off, t, n_embd)?;
2854        }
2855        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
2856
2857        for (il, layer) in self.layers.iter().enumerate() {
2858            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2859                format!("layer {il} carries no hyper-connection weights under an hc plan")
2860            })?;
2861
2862            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.attn, &x, t, n_embd)?;
2863            let mut h = e.uninit(t * n_embd)?;
2864            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2865            let mixed = match &layer.mixer {
2866                Mixer::Full(fa) => {
2867                    self.full_attn_prime(e, fa, &h, None, &pos_d, t, cache, il, seq_end)?
2868                }
2869                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2870                Mixer::Mla(mla) if mla.tp.is_some() => {
2871                    self.mla_tp_attn_cached(e, mla, &h, &pos_d, t, il, cache, false)?
2872                }
2873                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, &pos_d, t, il, cache)?,
2874                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2875                    e,
2876                    la,
2877                    &h,
2878                    t,
2879                    eps,
2880                    cache,
2881                    il,
2882                    crate::kda::ConvArm::Prefill,
2883                )?,
2884                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2885            };
2886            x = crate::hyper::post(e, &topology, &mixed, &x, &mix, t, n_embd)?;
2887
2888            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.mlp, &x, t, n_embd)?;
2889            let mut z = e.uninit(t * n_embd)?;
2890            e.rms_norm(
2891                &y,
2892                layer.post_attn_norm.float_data(),
2893                &mut z,
2894                n_embd,
2895                t,
2896                eps,
2897            )?;
2898            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2899            x = crate::hyper::post(e, &topology, &ffn_out, &x, &mix, t, n_embd)?;
2900            // glm5 DFlash2 feature tap (see hyper_range_prime — the unsplit chunk walk
2901            // taps the same completed-layer-output contraction).
2902            self.glm5_hc_tap(e, cache, &topology, il, &x, t)?;
2903        }
2904
2905        let hiddens =
2906            crate::hyper::collapse(e, &topology, self.hyper_head.as_ref(), &x, t, n_embd)?;
2907        let mut hn = e.uninit(t * n_embd)?;
2908        e.rms_norm(
2909            &hiddens,
2910            self.output_norm.float_data(),
2911            &mut hn,
2912            n_embd,
2913            t,
2914            eps,
2915        )?;
2916        let last = e.view(&hn, t * n_embd);
2917        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2918        let mut hlast = e.uninit(n_embd)?;
2919        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2920        let logits = e.matmul(&self.output, &hlast, 1)?;
2921        let host = e.dtoh(&logits)?;
2922
2923        // h_seed is the PRE-output_norm hidden of the last row (MTP-PLAN §A), taken from the
2924        // collapsed stack so it means the same thing it does on the serial path.
2925        let stack = e.view(&hiddens, t * n_embd);
2926        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2927        let mut h_seed = e.uninit(n_embd)?;
2928        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2929        cache.pos += t;
2930        Ok((host, h_seed, hiddens))
2931    }
2932
2933    /// Prime exit shared by `prime_cache_hyper` and its ppN twin: collapse, output_norm, last
2934    /// row logits, and the pre-output_norm hidden seed taken from the collapsed stack (MTP-PLAN
2935    /// §A) so it means the same thing it does on the serial path.
2936    #[allow(clippy::too_many_arguments)]
2937    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2938    fn hyper_prime_tail(
2939        &self,
2940        e: &Engine,
2941        topology: &crate::hyper::HyperTopology,
2942        x: &CudaSlice<f32>,
2943        t: usize,
2944        n_embd: usize,
2945        eps: f32,
2946        cache: &mut Cache,
2947    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2948        let hiddens = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2949        let mut hn = e.uninit(t * n_embd)?;
2950        e.rms_norm(
2951            &hiddens,
2952            self.output_norm.float_data(),
2953            &mut hn,
2954            n_embd,
2955            t,
2956            eps,
2957        )?;
2958        let last = e.view(&hn, t * n_embd);
2959        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2960        let mut hlast = e.uninit(n_embd)?;
2961        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2962        let logits = e.matmul(&self.output, &hlast, 1)?;
2963        let host = e.dtoh(&logits)?;
2964        let stack = e.view(&hiddens, t * n_embd);
2965        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2966        let mut h_seed = e.uninit(n_embd)?;
2967        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2968        cache.pos += t;
2969        Ok((host, h_seed, hiddens))
2970    }
2971
2972    /// ppN twin of `decode_step_hyper`: the T=1 step as N stage subgraphs, each on its own
2973    /// stream (and, under `MEMRA_PP_DEVICES`, its own device/engine), with the
2974    /// transport-selected boundary handoff of the `[streams, n_embd]` state at each fence cut.
2975    /// `cache.pos` is snapshotted once and advanced once.
2976    fn decode_step_hyper_ppn(
2977        &self,
2978        e: &Engine,
2979        token: u32,
2980        cache: &mut Cache,
2981        topology: &crate::hyper::HyperTopology,
2982        fence: &[usize],
2983    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2984        let n_embd = self.cfg.n_embd as usize;
2985        let eps = self.cfg.rms_eps;
2986        let pos = cache.pos;
2987        let width = topology.streams * n_embd;
2988
2989        if crate::pp::pp2_streams_off() {
2990            let pos_d = e.htod_i32(&[pos as i32])?;
2991            let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
2992            let mut x = crate::hyper::expand(e, topology, &embedded, 1, n_embd)?;
2993            x = self.hyper_range_decode(e, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
2994            for s in 1..fence.len() - 1 {
2995                let boundary_tx = e.clone_dtod(&x)?;
2996                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2997                x = self.hyper_range_decode(
2998                    e,
2999                    topology,
3000                    boundary_rx,
3001                    fence[s],
3002                    fence[s + 1],
3003                    &pos_d,
3004                    pos,
3005                    cache,
3006                )?;
3007            }
3008            return self.hyper_decode_tail(e, topology, &x, n_embd, eps, cache);
3009        }
3010
3011        let rt = crate::pp::PpNRt::get(e)?;
3012        let n_st = fence.len() - 1;
3013        assert_eq!(
3014            rt.n_stages(),
3015            n_st,
3016            "PpNRt stage count {} != fence stages {n_st}",
3017            rt.n_stages()
3018        );
3019        rt.fence_stages_behind(&e.stream())?;
3020
3021        let mut slot = {
3022            let _st0 = rt.enter(0);
3023            let e0 = rt.engine(0, e);
3024            let pos_d = e0.htod_i32(&[pos as i32])?;
3025            let embedded = e0.htod(&self.embd.gather(n_embd, &[token]))?;
3026            let x = crate::hyper::expand(e0, topology, &embedded, 1, n_embd)?;
3027            let x =
3028                self.hyper_range_decode(e0, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
3029            rt.tx(0, &x, width)?
3030        };
3031        for s in 1..n_st - 1 {
3032            let _st = rt.enter(s);
3033            let es = rt.engine(s, e);
3034            let pos_d = es.htod_i32(&[pos as i32])?;
3035            let x = rt.rx(s - 1, slot, width)?;
3036            let x = self.hyper_range_decode(
3037                es,
3038                topology,
3039                x,
3040                fence[s],
3041                fence[s + 1],
3042                &pos_d,
3043                pos,
3044                cache,
3045            )?;
3046            slot = rt.tx(s, &x, width)?;
3047        }
3048        let _stl = rt.enter(n_st - 1);
3049        let el = rt.engine(n_st - 1, e);
3050        let pos_d = el.htod_i32(&[pos as i32])?;
3051        let x = rt.rx(n_st - 2, slot, width)?;
3052        let x = self.hyper_range_decode(
3053            el,
3054            topology,
3055            x,
3056            fence[n_st - 1],
3057            fence[n_st],
3058            &pos_d,
3059            pos,
3060            cache,
3061        )?;
3062        self.hyper_decode_tail(el, topology, &x, n_embd, eps, cache)
3063    }
3064
3065    /// Decode exit shared by `decode_step_hyper` and its ppN twin. `h_seed` is the COLLAPSED
3066    /// hidden (not the pre-collapse stream state and not the post-norm row): that is what the
3067    /// serial hc step publishes, and the two must not drift.
3068    fn hyper_decode_tail(
3069        &self,
3070        e: &Engine,
3071        topology: &crate::hyper::HyperTopology,
3072        x: &CudaSlice<f32>,
3073        n_embd: usize,
3074        eps: f32,
3075        cache: &mut Cache,
3076    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3077        let h_seed = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, 1, n_embd)?;
3078        let mut hn = e.uninit(n_embd)?;
3079        e.rms_norm(
3080            &h_seed,
3081            self.output_norm.float_data(),
3082            &mut hn,
3083            n_embd,
3084            1,
3085            eps,
3086        )?;
3087        let logits = e.matmul(&self.output, &hn, 1)?;
3088        let host = e.dtoh(&logits)?;
3089        cache.pos += 1;
3090        Ok((host, h_seed))
3091    }
3092
3093    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
3094    pub fn forward(
3095        &self,
3096        e: &Engine,
3097        tokens: &[u32],
3098    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3099        if self.hyper.is_some() {
3100            return self.forward_hyper(e, tokens, false);
3101        }
3102        if self.is_gemma4_e4b() {
3103            return self.gemma4_e4b_forward(e, tokens, false);
3104        }
3105        if self.uses_gemma_program() {
3106            return self.gemma4_forward(e, tokens, false);
3107        }
3108        let cfg = &self.cfg;
3109        let n_embd = cfg.n_embd as usize;
3110        let t = tokens.len();
3111        let eps = cfg.rms_eps;
3112        let pos: Vec<i32> = (0..t as i32).collect();
3113        let pos_d = e.htod_i32(&pos)?;
3114
3115        let mut x = self.embed(e, tokens)?; // [T, n_embd]
3116
3117        for (il, layer) in self.layers.iter().enumerate() {
3118            // attn_norm
3119            let mut h = e.uninit(t * n_embd)?;
3120            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3121
3122            let mixed = match &layer.mixer {
3123                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
3124                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
3125                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3126                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
3127            };
3128
3129            // residual 1
3130            let mut x1 = e.uninit(t * n_embd)?;
3131            e.add(&x, &mixed, &mut x1, t * n_embd)?;
3132
3133            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
3134            let mut z = e.uninit(t * n_embd)?;
3135            e.rms_norm(
3136                &x1,
3137                layer.post_attn_norm.float_data(),
3138                &mut z,
3139                n_embd,
3140                t,
3141                eps,
3142            )?;
3143            let ffn_out = match &layer.ffn {
3144                crate::hybrid::Ffn::Dense {
3145                    ffn_gate,
3146                    ffn_up,
3147                    ffn_down,
3148                } => {
3149                    let n_ff = ffn_gate.out_features();
3150                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3151                    let up = g2.pop().unwrap();
3152                    let gate = g2.pop().unwrap();
3153                    let mut act = e.uninit(t * n_ff)?;
3154                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
3155                    // both the dense MLP and the shared expert, and its limit is
3156                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
3157                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
3158                    Self::ffn_act_lim(
3159                        e,
3160                        &self.cfg,
3161                        &gate,
3162                        &up,
3163                        1.0,
3164                        1.0,
3165                        self.cfg.clamp_shexp_at(il as u32),
3166                        &mut act,
3167                        t * n_ff,
3168                    )?;
3169                    e.matmul(ffn_down, &act, t)?
3170                }
3171                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
3172            };
3173            let mut x2 = e.uninit(t * n_embd)?;
3174            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3175            x = x2;
3176        }
3177
3178        let mut hn = e.uninit(t * n_embd)?;
3179        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3180        let logits = e.matmul(&self.output, &hn, t)?;
3181        e.dtoh(&logits)
3182    }
3183
3184    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
3185    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
3186    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
3187    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
3188    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
3189    pub fn forward_last(
3190        &self,
3191        e: &Engine,
3192        tokens: &[u32],
3193    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3194        if self.hyper.is_some() {
3195            return self.forward_hyper(e, tokens, true);
3196        }
3197        if self.uses_gemma_program() {
3198            return self.gemma4_forward(e, tokens, true);
3199        }
3200        let cfg = &self.cfg;
3201        let n_embd = cfg.n_embd as usize;
3202        let t = tokens.len();
3203        let eps = cfg.rms_eps;
3204        let pos: Vec<i32> = (0..t as i32).collect();
3205        let pos_d = e.htod_i32(&pos)?;
3206
3207        let mut x = self.embed(e, tokens)?; // [T, n_embd]
3208        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
3209        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
3210        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
3211        let anat = Self::prime_anatomy_on();
3212        let mut anat_last = if anat {
3213            e.stream().synchronize()?;
3214            Some(std::time::Instant::now())
3215        } else {
3216            None
3217        };
3218        macro_rules! anat_mark {
3219            ($slot:expr) => {
3220                if let Some(ts) = anat_last.as_mut() {
3221                    e.stream().synchronize()?;
3222                    Self::prime_anatomy_slots()[$slot].fetch_add(
3223                        ts.elapsed().as_nanos() as u64,
3224                        std::sync::atomic::Ordering::Relaxed,
3225                    );
3226                    *ts = std::time::Instant::now();
3227                }
3228            };
3229        }
3230        for (il, layer) in self.layers.iter().enumerate() {
3231            let mut h = e.uninit(t * n_embd)?;
3232            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3233            if probe {
3234                e.stream().synchronize()?;
3235                eprintln!("[probe] L{il} norm ok");
3236            }
3237            anat_mark!(4);
3238            let mixed = match &layer.mixer {
3239                Mixer::Full(fa) => {
3240                    let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
3241                    anat_mark!(0);
3242                    y
3243                }
3244                Mixer::Linear(la) => {
3245                    let y = self.linear_attn(e, la, &h, t)?;
3246                    anat_mark!(1);
3247                    y
3248                }
3249                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3250                Mixer::Kda(la) => {
3251                    let y = crate::kda::kda_attn(e, la, &h, t, eps)?;
3252                    // KDA shares the linear-mixer anatomy slot: same mixer class, one bucket.
3253                    anat_mark!(1);
3254                    y
3255                }
3256            };
3257            if probe {
3258                e.stream().synchronize()?;
3259                eprintln!("[probe] L{il} mixer ok");
3260            }
3261            let mut x1 = e.uninit(t * n_embd)?;
3262            e.add(&x, &mixed, &mut x1, t * n_embd)?;
3263            let mut z = e.uninit(t * n_embd)?;
3264            e.rms_norm(
3265                &x1,
3266                layer.post_attn_norm.float_data(),
3267                &mut z,
3268                n_embd,
3269                t,
3270                eps,
3271            )?;
3272            anat_mark!(4);
3273            let ffn_out = match &layer.ffn {
3274                crate::hybrid::Ffn::Dense {
3275                    ffn_gate,
3276                    ffn_up,
3277                    ffn_down,
3278                } => {
3279                    let n_ff = ffn_gate.out_features();
3280                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3281                    let up = g2.pop().unwrap();
3282                    let gate = g2.pop().unwrap();
3283                    let mut act = e.uninit(t * n_ff)?;
3284                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
3285                    Self::ffn_act_lim(
3286                        e,
3287                        &self.cfg,
3288                        &gate,
3289                        &up,
3290                        1.0,
3291                        1.0,
3292                        self.cfg.clamp_shexp_at(il as u32),
3293                        &mut act,
3294                        t * n_ff,
3295                    )?;
3296                    let y = e.matmul(ffn_down, &act, t)?;
3297                    anat_mark!(3);
3298                    y
3299                }
3300                crate::hybrid::Ffn::Moe(m) => {
3301                    let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
3302                    anat_mark!(2);
3303                    y
3304                }
3305            };
3306            if probe {
3307                e.stream().synchronize()?;
3308                eprintln!("[probe] L{il} ffn ok");
3309            }
3310            let mut x2 = e.uninit(t * n_embd)?;
3311            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3312            x = x2;
3313        }
3314        if anat {
3315            let s = Self::prime_anatomy_slots();
3316            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
3317            eprintln!(
3318                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
3319                 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
3320                ms(0),
3321                ms(1),
3322                ms(2),
3323                ms(3),
3324                ms(4)
3325            );
3326        }
3327        // norm over all T, then slice the LAST row and run lm_head on that single row.
3328        let mut hn = e.uninit(t * n_embd)?;
3329        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3330        let last = e.view(&hn, t * n_embd); // [T, n_embd]
3331        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
3332        let mut hlast = e.uninit(n_embd)?;
3333        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3334        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
3335        e.dtoh(&logits)
3336    }
3337
3338    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
3339    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
3340    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
3341    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
3342    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
3343    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
3344    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
3345    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
3346    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
3347    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
3348    ///       argmax gate is the accuracy authority, exactly as for forward_last);
3349    ///   (c) `cache.pos`/KV len/len_d advance by T.
3350    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
3351    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
3352    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
3353    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
3354    ///
3355    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
3356    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
3357    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
3358    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
3359    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
3360    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
3361    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
3362    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
3363    /// differently under load — research/tick-seg-20260807, receipt in
3364    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
3365    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
3366    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
3367    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
3368    /// caller that SPLITS one request across calls passes the remainder.
3369    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3370    pub fn prime_cache(
3371        &self,
3372        e: &Engine,
3373        tokens: &[u32],
3374        cache: &mut Cache,
3375        queued_after: usize,
3376    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3377        self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
3378    }
3379
3380    /// The engine that owns EMBEDDING INTAKE for the current placement — the engine whose
3381    /// context a vision overlay's rows must live in (`EmbedOverlay::new_published`).
3382    ///
3383    /// It is NOT always the primary engine. Under a per-stage-stream ppN split the embedding
3384    /// happens on stage 0 (`prime_cache_hyper_ppn` embeds inside the stage-0 scope and every
3385    /// later stage sees only the expanded stream state), and stage 0 gets its own Engine
3386    /// whenever its device differs from the primary's — which is the deployed 3-card shape,
3387    /// because the worker's primary engine follows the LAST stage. On a single-device
3388    /// placement, with the door shut, or on the `MEMRA_PP_STREAMS=0` seam, intake is the
3389    /// primary engine and this returns `e` unchanged (byte-identical to the pre-lane path).
3390    ///
3391    /// This mirrors `prime_cache_hyper`'s own door test deliberately, and the ppN prime's
3392    /// residency refusal is the enforcement: if this ever picks the wrong engine, the prime
3393    /// fails CLOSED with a named error instead of peer-reading an overlay.
3394    pub fn vision_intake_engine<'a>(
3395        &self,
3396        e: &'a Engine,
3397    ) -> Result<&'a Engine, Box<dyn std::error::Error>> {
3398        if self.hyper.is_some()
3399            && !crate::pp::pp2_streams_off()
3400            && crate::pp::pp_cuts(self.layers.len()).is_some()
3401        {
3402            let rt = crate::pp::PpNRt::get(e)?;
3403            return Ok(rt.engine(0, e));
3404        }
3405        Ok(e)
3406    }
3407
3408    /// `prime_cache` with a vision embedding overlay (lane/vision): image merger outputs
3409    /// replace the `<|image_pad|>` token embeddings at prompt-relative positions before the
3410    /// trunk walk — the mixed-embedding prime. Text-only callers use `prime_cache` (overlay
3411    /// None, byte-identical path). Scope: the serial chunk walk, the single-engine hyper walk,
3412    /// and the hyper ppN twin (splice at stage-0 embedding intake; lane/glm5-vision-default-on,
3413    /// gated by glm5-hyper-ppn-gate's overlay arm). The serial PP-2 pipelined prime and
3414    /// gemma4 E4B still refuse loudly.
3415    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3416    pub fn prime_cache_overlaid(
3417        &self,
3418        e: &Engine,
3419        tokens: &[u32],
3420        cache: &mut Cache,
3421        queued_after: usize,
3422        overlay: Option<&crate::vision::EmbedOverlay>,
3423    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3424        cache.ensure_usable("prime_cache")?;
3425        if self.hyper.is_some() {
3426            // The mixed-embedding splice lands BEFORE stream expansion (the same point
3427            // the reference's execute_multimodal replaces rows — before hc_expand), so
3428            // the hyper walk needs no overlay-specific arithmetic (lane/glm5-vision).
3429            return self.prime_cache_hyper(e, tokens, cache, queued_after, overlay);
3430        }
3431        let _pp_walk =
3432            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
3433                let rt = crate::pp::PpNRt::get(e)?;
3434                Some(rt.acquire_walk("prime_cache")?)
3435            } else {
3436                None
3437            };
3438        let n_embd = self.cfg.n_embd as usize;
3439        let t = tokens.len();
3440        // MEMRA_PRIME_TROWS=1: prefill through the same-session t-row walk (per-row t=1
3441        // program = the tokenwise-prime ORACLE class) — replaces the host-canonical
3442        // per-token step-TP prime. Text-only fresh primes; anything else falls through.
3443        // MEMRA_STEP_GEMM_PRIME: prime the prompt through the batched GEMM path, CHUNKED.
3444        // The batch entry supplies both halves of the fast prime — the GEMM trunk at m = chunk
3445        // and the grouped NVFP4 MoE — which is why routing only the MoE through the ordinary
3446        // chunk loop measured 26.9 s against 3.7 s here. Chunking keeps the transients bounded:
3447        // a whole 32k prompt in one call would build a 262144-pair CSR and ~4.3 GB of partials
3448        // per rank, the blow-up the chunked prime exists to prevent.
3449        //
3450        // CONTINUATION (lane/gemm-suffix, 2026-08-28): the entry NO LONGER requires
3451        // `cache.pos == 0`. The batch core has been continuation-capable since 7700e0b6
3452        // (positions carry each sequence's base; the fresh-prompt guard narrowed to B > 1),
3453        // and d99b2ea3 named this outer guard as the remaining blocker in its own message.
3454        // Every multi-turn suffix and every tick remainder was paying the walk's measured
3455        // ~7.2 ms/token against this path's ~1.0 ms/token, which is why session-affinity
3456        // reuse measured a 1.012x wash on a growing conversation.
3457        // ONE DEFECT HAD TO BE FIXED FIRST, and it was LIVE before this lift:
3458        // `step35_prime_batch_layers` passed `ts[s]` — the CHUNK's length — as `seq_end`.
3459        // `seq_end` is the REQUEST's absolute end position and it steers step35's SWA arm
3460        // (`seq_end > win`, win = 512 on step37). A chunk SHORTER than the window at a
3461        // NONZERO base therefore selected the UNWINDOWED FA arm over a view that the `off`
3462        // trim leaves at ~win-1+t rows: it attended OUTSIDE the sliding window. That was
3463        // already reachable with no continuation at all — a fresh prompt of 4096+k for k in
3464        // [PRIME_MIN_T, 512) ends in a trailing chunk of exactly that shape. `seq_end` is now
3465        // threaded from here (request-absolute, `+ queued_after`, computed ONCE before the
3466        // chunk loop, chunk-size-invariant exactly as on the walk), which is what makes the
3467        // suffix arm expressible at all rather than merely reachable.
3468        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
3469        let seq_end = if legacy_calllocal {
3470            cache.pos + t
3471        } else {
3472            cache.pos + t + queued_after
3473        };
3474        // MEMRA_STEP_GEMM_PRIME_SUFFIX is the SUFFIX-ONLY seam: off (its default in this
3475        // commit) leaves continuations on the walk while fresh primes stay on the fast path;
3476        // MEMRA_STEP_GEMM_PRIME=0 is the whole-path seam. The `seq_end` threading above is
3477        // deliberately NOT behind either door — it is a correctness fix for the fresh path too.
3478        if overlay.is_none()
3479            && (cache.pos == 0 || step_gemm_prime_suffix_on())
3480            && t >= PRIME_MIN_T
3481            && crate::step_gemm_prime_on()
3482            && self.uses_sliding_gated_moe_program()
3483        {
3484            let n_embd = self.cfg.n_embd as usize;
3485            let base = cache.pos;
3486            let width = crate::cache::PRIME_CHUNK_MAX_TOKENS;
3487            let mut hiddens = e.uninit(t * n_embd)?;
3488            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3489            let mut start = 0usize;
3490            while start < t {
3491                // A trailing chunk below the walk floor folds into the previous one; every chunk
3492                // this entry sees must clear PRIME_MIN_T on its own.
3493                let mut end = (start + width).min(t);
3494                if t - end > 0 && t - end < PRIME_MIN_T {
3495                    end = t;
3496                }
3497                let mut out = self.step35_prime_cache_batch(
3498                    e,
3499                    &[&tokens[start..end]],
3500                    &mut [cache],
3501                    &[seq_end],
3502                )?;
3503                if out.len() != 1 {
3504                    return Err("B=1 batched prime returned a non-singleton".into());
3505                }
3506                let (logits, h_seed, hidden) = out.remove(0);
3507                e.copy_into(
3508                    &mut hiddens,
3509                    start * n_embd,
3510                    &hidden,
3511                    (end - start) * n_embd,
3512                )?;
3513                last = Some((logits, h_seed));
3514                start = end;
3515            }
3516            let (logits, h_seed) = last.expect("prime produced no chunk");
3517            // ENGAGEMENT RECEIPT, both directions. `base` is the discriminator: base=0 is a
3518            // fresh prime (this line existed before the lift), base>0 is a SUFFIX riding the
3519            // GEMM trunk — the arm this lane added. The declining twin below counts the other
3520            // direction, so a log that shows neither line is an instrument fault, not a pass.
3521            eprintln!(
3522                "[gemm-prime] ENGAGED t={t} base={base} seq_end={seq_end} chunks<={width} (GEMM trunk + grouped MoE)"
3523            );
3524            return Ok((logits, h_seed, hiddens));
3525        }
3526        if self.uses_sliding_gated_moe_program() {
3527            eprintln!(
3528                "[gemm-prime] WALK t={t} base={} seq_end={seq_end} (batched prime declined)",
3529                cache.pos
3530            );
3531        }
3532        if overlay.is_none()
3533            && let Some(out) = self.step35_prime_trows(e, tokens, cache)?
3534        {
3535            return Ok(out);
3536        }
3537        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
3538        // session cache — every chunk (including the first) takes the continuation arm
3539        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
3540        assert!(
3541            t >= PRIME_MIN_T,
3542            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
3543        );
3544        assert!(
3545            cache.pos + t <= cache.max_ctx,
3546            "prime_cache: prompt exceeds cache max_ctx"
3547        );
3548
3549        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
3550        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
3551        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
3552        // each chunk runs the full layer stack with transients sized to the chunk, appending its
3553        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
3554        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
3555        // exactly the state carry it was built for). Full-attn chunks after the first attend to
3556        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
3557        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
3558        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
3559        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
3560        if self.is_gemma4_e4b() || self.uses_gemma_program() {
3561            if self.is_gemma4_e4b() {
3562                if overlay.is_some() {
3563                    return Err(
3564                        "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
3565                    );
3566                }
3567                return self.gemma4_e4b_prime(e, tokens, cache);
3568            }
3569            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
3570            // An overlay takes the masked-prefill arm: image rows splice in unscaled
3571            // (gemma4.cpp:182 — embd batches skip the sqrt(n_embd) scale) and the image
3572            // spans become bidirectional attention islands (lane/gemma-vision).
3573            return self.gemma4_prime(e, tokens, cache, overlay);
3574        }
3575        if crate::pp::prime_pipe_on()
3576            && crate::pp::prime_pp_on()
3577            && !crate::pp::pp2_streams_off()
3578            && crate::pp::pp_cuts(self.layers.len())
3579                .is_some_and(|fence| matches!(fence.len(), 4 | 5))
3580        {
3581            crate::pp::pp_wave_on()
3582                .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3583        }
3584        let ranges = prime_chunk_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
3585        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
3586        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
3587        // the prefill's ARITHMETIC, so two rigs with different values produced different
3588        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
3589        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
3590        // (VERDICT.md) — and it is NOT what docs originally said:
3591        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
3592        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
3593        //     output head), so growing a chunk cannot move an existing row's value.
3594        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
3595        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
3596        //     not describe our leak.
3597        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
3598        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
3599        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
3600        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
3601        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
3602        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
3603        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
3604        // the source — every row is in one numeric class, so the chunk size no longer steers
3605        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
3606        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
3607        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
3608        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
3609        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
3610        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
3611        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
3612        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
3613        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
3614        // across calls, the request still ends at the same absolute position, whatever the tick
3615        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
3616        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
3617        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
3618        // default. Read per call, not cached (the probe flips it in-process between arms). Never
3619        // on in a measured default run.
3620        // `seq_end` (and its MEMRA_PRIME_CALLLOCAL seam) is computed ONCE above the batched
3621        // entry so both prime arms read the identical request-absolute value.
3622        if ranges.len() == 1 {
3623            return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
3624        }
3625        // PIPELINED PP PRIME. PP-2 retains its independently-qualified two-stage schedule;
3626        // PP-3/4 require the explicit MEMRA_PP_WAVE=1 persistent-stage wavefront. The serial
3627        // split stays reachable through MEMRA_PRIME_PIPE=0 and is the exactness oracle.
3628        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
3629            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
3630                if overlay.is_some() {
3631                    return Err(
3632                        "vision embedding overlay + pipelined PP prime unsupported (v1); \
3633                         run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
3634                            .into(),
3635                    );
3636                }
3637                if crate::pp::pp_multi_stream_same_device() {
3638                    return Err(
3639                        "prime chunk pipeline refused with 2 stage streams on one device — \
3640                         that concurrent-stream placement remains quarantined by the deferred \
3641                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
3642                         the serial split."
3643                            .into(),
3644                    );
3645                }
3646                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3647            }
3648            if let Some(fence) =
3649                crate::pp::pp_cuts(self.layers.len()).filter(|f| matches!(f.len(), 4 | 5))
3650            {
3651                let wave_on = crate::pp::pp_wave_on()
3652                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3653                let stages = fence.len() - 1;
3654                if crate::pp::pp_wave_route_enabled(
3655                    wave_on,
3656                    crate::pp::pp2_overlap(),
3657                    stages,
3658                    ranges.len(),
3659                ) {
3660                    if overlay.is_some() {
3661                        return Err(
3662                            "vision embedding overlay + pipelined PP prime unsupported; \
3663                             run the serial prime (MEMRA_PP_WAVE=0 or MEMRA_PRIME_PIPE=0)"
3664                                .into(),
3665                        );
3666                    }
3667                    let rt = crate::pp::PpNRt::get(e)?;
3668                    let double_slot = crate::pp::pp2_overlap();
3669                    crate::pp::pp_wave_eligibility(
3670                        stages,
3671                        double_slot,
3672                        rt.host_bounce_active(),
3673                        rt.repeated_stage_device(),
3674                    )
3675                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3676                    return self
3677                        .prime_cache_ppn_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3678                }
3679            }
3680        }
3681        let mut hiddens = e.uninit(t * n_embd)?;
3682        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3683        for &(start, end) in &ranges {
3684            // chunked prime writes tap rows at the chunk's absolute offset
3685            if let Some(taps) = cache.dflash_taps.as_mut() {
3686                taps.base = start;
3687            }
3688            let (l, hs, x) =
3689                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
3690            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
3691            last = Some((l, hs));
3692        }
3693        let (logits, h_seed) = last.unwrap();
3694        Ok((logits, h_seed, hiddens))
3695    }
3696
3697    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
3698    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
3699    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
3700    /// norm, lm head, and caller hidden-stack copy as the serial split.
3701    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3702    fn prime_cache_pp2_pipelined(
3703        &self,
3704        e: &Engine,
3705        tokens: &[u32],
3706        cache: &mut Cache,
3707        seq_end: usize,
3708        ranges: &[(usize, usize)],
3709        fence: &[usize],
3710    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3711        debug_assert_eq!(fence.len(), 3);
3712        debug_assert!(ranges.len() >= 2);
3713        let rt = crate::pp::PpNRt::get(e)?;
3714        assert_eq!(
3715            rt.n_stages(),
3716            2,
3717            "prime pipeline requires exactly two PP stages"
3718        );
3719        let n_embd = self.cfg.n_embd as usize;
3720        let t = tokens.len();
3721        let initial_base = cache.pos;
3722        let caller_stream = e.stream();
3723
3724        // #87 reverse publication before any new stage allocation, then prewarm both
3725        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
3726        // after stage 1(N) is queued would synchronize that stream and erase the first
3727        // overlap on a two-chunk prompt.
3728        rt.fence_stages_behind(&caller_stream)?;
3729        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
3730        rt.prepare_overlap_slots(0, max_payload)?;
3731
3732        let mut hiddens = e.uninit(t * n_embd)?;
3733        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3734        let mut stage_caches = PrimeCacheStages::new(cache, fence);
3735        let (cache0, cache1) = stage_caches.pp2_parts();
3736        let (first_start, first_end) = ranges[0];
3737        let mut slot = self.prime_pp2_stage0_enqueue(
3738            e,
3739            rt,
3740            &tokens[first_start..first_end],
3741            cache0,
3742            seq_end,
3743            fence,
3744            initial_base + first_start,
3745            true,
3746        )?;
3747        cache0.pos = initial_base + first_end;
3748
3749        for (i, &(start, end)) in ranges.iter().enumerate() {
3750            let base = initial_base + start;
3751            debug_assert_eq!(
3752                cache1.pos, base,
3753                "stage 1 must drain chunks in original position order"
3754            );
3755            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
3756                let next_base = initial_base + next_start;
3757                debug_assert_eq!(
3758                    cache0.pos, next_base,
3759                    "stage 0 must issue chunks in original position order"
3760                );
3761                let cache0_stage = &mut *cache0;
3762                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
3763                // on one host thread therefore serialize even if the calls are ordered as
3764                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
3765                // stage 1 consumes slot N while stage 0 produces slot N+1.
3766                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
3767                    let stage0 = scope.spawn(move || -> Result<usize, String> {
3768                        let next = self
3769                            .prime_pp2_stage0_enqueue(
3770                                e,
3771                                rt,
3772                                &tokens[next_start..next_end],
3773                                cache0_stage,
3774                                seq_end,
3775                                fence,
3776                                next_base,
3777                                true,
3778                            )
3779                            .map_err(|err| err.to_string())?;
3780                        cache0_stage.pos = initial_base + next_end;
3781                        Ok(next)
3782                    });
3783                    let x = self.prime_pp2_stage1_enqueue(
3784                        e,
3785                        rt,
3786                        slot,
3787                        end - start,
3788                        cache1,
3789                        seq_end,
3790                        fence,
3791                        base,
3792                        true,
3793                    )?;
3794                    let out = {
3795                        rt.bind_stage(1)?;
3796                        let _st1 = rt.enter(1);
3797                        let e1 = rt.engine(1, e);
3798                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3799                    };
3800                    let next = match stage0.join() {
3801                        Ok(result) => {
3802                            result.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?
3803                        }
3804                        Err(payload) => std::panic::resume_unwind(payload),
3805                    };
3806                    Ok((out, Some(next)))
3807                })?
3808            } else {
3809                let x = self.prime_pp2_stage1_enqueue(
3810                    e,
3811                    rt,
3812                    slot,
3813                    end - start,
3814                    cache1,
3815                    seq_end,
3816                    fence,
3817                    base,
3818                    true,
3819                )?;
3820                let out = {
3821                    rt.bind_stage(1)?;
3822                    let _st1 = rt.enter(1);
3823                    let e1 = rt.engine(1, e);
3824                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3825                };
3826                (out, None)
3827            };
3828
3829            rt.publish_to(1, &caller_stream)?;
3830            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
3831            last = Some((out.0, out.1));
3832            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3833
3834            if let Some(next) = next_slot {
3835                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
3836                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
3837                // Stage 0(N+1) is already queued before this wait is appended, so its
3838                // overlap with stage 1(N) is preserved.
3839                rt.fence_stages_behind(&caller_stream)?;
3840                slot = next;
3841            }
3842        }
3843
3844        debug_assert_eq!(cache0.pos, initial_base + t);
3845        debug_assert_eq!(cache1.pos, initial_base + t);
3846        let (logits, h_seed) = last.unwrap();
3847        stage_caches.commit();
3848        Ok((logits, h_seed, hiddens))
3849    }
3850
3851    /// PP-3/4 prime wavefront: one prompt microchunk is one wave. One scoped host worker owns each
3852    /// non-head stage for the whole walk; the caller thread owns the head stage. Forward boundary
3853    /// messages preserve wave order, while reverse exact-wave acknowledgements are sent only after
3854    /// downstream `rx` has recorded `ev_rx`. An upstream stage therefore cannot cycle back to either
3855    /// shared slot before that slot's current generation has a host-observed release point.
3856    ///
3857    /// PP-2 remains on its independently qualified scheduler above. This path is reachable only
3858    /// through MEMRA_PP_WAVE=1 and the topology gate.
3859    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3860    fn prime_cache_ppn_pipelined(
3861        &self,
3862        e: &Engine,
3863        tokens: &[u32],
3864        cache: &mut Cache,
3865        seq_end: usize,
3866        ranges: &[(usize, usize)],
3867        fence: &[usize],
3868    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3869        let stages = fence.len().saturating_sub(1);
3870        debug_assert!((3..=4).contains(&stages));
3871        debug_assert!(ranges.len() >= 2);
3872        let rt = crate::pp::PpNRt::get(e)?;
3873        assert_eq!(
3874            rt.n_stages(),
3875            stages,
3876            "prime wavefront PpNRt/fence stage mismatch"
3877        );
3878        let n_embd = self.cfg.n_embd as usize;
3879        let initial_base = cache.pos;
3880        let caller_stream = e.stream();
3881        let primary_context = crate::pp::PrimaryContextRestore::new(e);
3882
3883        rt.fence_stages_behind(&caller_stream)?;
3884        let max_payload = ranges
3885            .iter()
3886            .map(|(start, end)| (end - start) * n_embd)
3887            .max()
3888            .unwrap_or(0);
3889        for boundary in 0..stages - 1 {
3890            rt.prepare_overlap_slots(boundary, max_payload)?;
3891        }
3892
3893        let mut stage_caches = PrimeCacheStages::new(cache, fence);
3894        let waves: Vec<_> = ranges
3895            .iter()
3896            .map(|&(start, end)| PrimePpWave {
3897                start,
3898                end,
3899                tokens: &tokens[start..end],
3900            })
3901            .collect();
3902        let mut forward_senders = Vec::with_capacity(stages - 1);
3903        let mut forward_receivers = Vec::with_capacity(stages - 1);
3904        let mut release_senders = Vec::with_capacity(stages - 1);
3905        let mut release_receivers = Vec::with_capacity(stages - 1);
3906        for _ in 0..stages - 1 {
3907            let (forward_sender, forward_receiver) = std::sync::mpsc::channel();
3908            let (release_sender, release_receiver) = std::sync::mpsc::channel();
3909            forward_senders.push(Some(forward_sender));
3910            forward_receivers.push(Some(forward_receiver));
3911            release_senders.push(Some(release_sender));
3912            release_receivers.push(Some(release_receiver));
3913        }
3914        let mut stage_channels = Vec::with_capacity(stages - 1);
3915        for stage in 0..stages - 1 {
3916            stage_channels.push(Some(PrimePpStageChannels {
3917                incoming: (stage > 0).then(|| forward_receivers[stage - 1].take().unwrap()),
3918                release_upstream: (stage > 0).then(|| release_senders[stage - 1].take().unwrap()),
3919                outgoing: forward_senders[stage].take().unwrap(),
3920                released_downstream: release_receivers[stage].take().unwrap(),
3921            }));
3922        }
3923        let head_incoming = forward_receivers[stages - 2].take().unwrap();
3924        let head_release = release_senders[stages - 2].take().unwrap();
3925        // allow: one-shot composite type; naming it would hide the shape that matters here —
3926        // this is the head stage's per-wave output, indexed by wave.
3927        #[allow(clippy::type_complexity)]
3928        let mut results: Vec<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>> =
3929            std::iter::repeat_with(|| None).take(waves.len()).collect();
3930        let walk_result = std::thread::scope(|scope| -> Result<(), Box<dyn std::error::Error>> {
3931            let waves_ref = &waves;
3932            let mut handles = Vec::with_capacity(stages - 1);
3933            // `stage` is a STAGE ID, not merely an index: it indexes two different containers
3934            // (`stage_channels`, `stage_caches.stages()`) and is passed to the worker as the
3935            // stage it owns. An iterator form would keep only one of the three uses.
3936            #[allow(clippy::needless_range_loop)]
3937            for stage in 0..stages - 1 {
3938                let channels = stage_channels[stage].take().unwrap();
3939                let cache_state = &stage_caches.stages()[stage];
3940                handles.push(scope.spawn(move || -> Result<(), String> {
3941                    let result = self.prime_ppn_wave_worker(
3942                        e,
3943                        rt,
3944                        waves_ref,
3945                        cache_state,
3946                        channels.incoming.as_ref(),
3947                        channels.release_upstream.as_ref(),
3948                        &channels.outgoing,
3949                        &channels.released_downstream,
3950                        stage,
3951                        seq_end,
3952                        fence,
3953                        initial_base,
3954                    );
3955                    if let Err(error) = &result {
3956                        channels.notify_failure(&error.to_string());
3957                    }
3958                    result.map_err(|error| error.to_string())
3959                }));
3960            }
3961
3962            let mut head_panic = None;
3963            let head_result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(
3964                || -> Result<(), Box<dyn std::error::Error>> {
3965                    let mut head_cache = stage_caches.stages()[stages - 1]
3966                        .lock()
3967                        .map_err(|_| "prime PP head cache lock poisoned")?;
3968                    for (wave_index, wave) in waves_ref.iter().enumerate() {
3969                        let incoming = recv_prime_pp_signal(
3970                            &head_incoming,
3971                            PrimePpWaveSlot {
3972                                wave: wave_index,
3973                                slot: 0,
3974                            },
3975                            false,
3976                            "prime PP head input",
3977                        )?;
3978                        results[wave_index] = Some(self.prime_ppn_wave_final(
3979                            e,
3980                            rt,
3981                            wave,
3982                            &mut head_cache,
3983                            incoming,
3984                            &head_release,
3985                            seq_end,
3986                            fence,
3987                            initial_base,
3988                        )?);
3989                    }
3990                    Ok(())
3991                },
3992            )) {
3993                Ok(result) => result,
3994                Err(payload) => {
3995                    head_panic = Some(payload);
3996                    Err("prime PP head-stage host walker panicked".into())
3997                }
3998            };
3999            if let Err(error) = &head_result {
4000                let _ = head_release.send(PrimePpSignal::Error(error.to_string()));
4001            }
4002            let mut first_error = head_result.err().map(|error| error.to_string());
4003            let mut worker_panic = None;
4004            for handle in handles {
4005                match handle.join() {
4006                    Ok(Ok(())) => {}
4007                    Ok(Err(error)) => {
4008                        first_error.get_or_insert(error);
4009                    }
4010                    Err(payload) => {
4011                        if worker_panic.is_none() {
4012                            worker_panic = Some(payload);
4013                        }
4014                    }
4015                }
4016            }
4017            if let Some(payload) = head_panic {
4018                std::panic::resume_unwind(payload);
4019            }
4020            if let Some(payload) = worker_panic {
4021                std::panic::resume_unwind(payload);
4022            }
4023            if let Some(error) = first_error {
4024                return Err(error.into());
4025            }
4026            Ok(())
4027        });
4028        let publish_result = if walk_result.is_ok() {
4029            Some(rt.publish_to(stages - 1, &caller_stream))
4030        } else {
4031            None
4032        };
4033        let restore_result = primary_context.restore();
4034        walk_result?;
4035        if let Some(result) = publish_result {
4036            result?;
4037        }
4038        restore_result?;
4039        let mut hiddens = e.uninit(tokens.len() * n_embd)?;
4040        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
4041        for (wave_index, (wave, result)) in waves.iter().zip(results).enumerate() {
4042            debug_assert_eq!((wave.start, wave.end), ranges[wave_index]);
4043            let out = result.ok_or("prime PP wavefront completed without a head-stage result")?;
4044            e.copy_into(
4045                &mut hiddens,
4046                wave.start * n_embd,
4047                &out.2,
4048                (wave.end - wave.start) * n_embd,
4049            )?;
4050            last = Some((out.0, out.1));
4051            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4052        }
4053        stage_caches.commit();
4054        drop(stage_caches);
4055
4056        static LOGGED: std::sync::Once = std::sync::Once::new();
4057        LOGGED.call_once(|| {
4058            eprintln!(
4059                "[pp-wave] PP{stages} prime wavefront engaged: microchunks={} \
4060                 (experimental, MEMRA_PP_WAVE=1)",
4061                ranges.len(),
4062            );
4063        });
4064        let (logits, h_seed) = last.expect("prime PP wavefront produced no microchunk");
4065        crate::pp::record_pp_wave_tick();
4066        Ok((logits, h_seed, hiddens))
4067    }
4068
4069    #[allow(clippy::too_many_arguments)]
4070    fn prime_ppn_wave_worker(
4071        &self,
4072        e: &Engine,
4073        rt: &crate::pp::PpNRt,
4074        waves: &[PrimePpWave<'_>],
4075        cache: &std::sync::Mutex<Cache>,
4076        incoming: Option<&std::sync::mpsc::Receiver<PrimePpSignal>>,
4077        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
4078        outgoing: &std::sync::mpsc::Sender<PrimePpSignal>,
4079        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
4080        stage: usize,
4081        seq_end: usize,
4082        fence: &[usize],
4083        initial_base: usize,
4084    ) -> Result<(), Box<dyn std::error::Error>> {
4085        debug_assert_eq!(incoming.is_some(), stage > 0);
4086        debug_assert_eq!(release_upstream.is_some(), stage > 0);
4087        let mut cache = cache
4088            .lock()
4089            .map_err(|_| "prime PP cache stage lock poisoned")?;
4090        let mut credits = PrimePpWaveCredits::default();
4091        for (wave_index, wave) in waves.iter().enumerate() {
4092            let incoming = match incoming {
4093                Some(receiver) => Some(recv_prime_pp_signal(
4094                    receiver,
4095                    PrimePpWaveSlot {
4096                        wave: wave_index,
4097                        slot: 0,
4098                    },
4099                    false,
4100                    "prime PP stage input",
4101                )?),
4102                None => None,
4103            };
4104            let sent = self.prime_ppn_wave_stage(
4105                e,
4106                rt,
4107                wave,
4108                &mut cache,
4109                stage,
4110                incoming,
4111                release_upstream,
4112                &mut credits,
4113                released_downstream,
4114                seq_end,
4115                fence,
4116                initial_base,
4117            )?;
4118            send_prime_pp_signal(outgoing, PrimePpSignal::Slot(sent), "prime PP stage output")?;
4119        }
4120        while let Some(expected) = credits.pending.front().copied() {
4121            let released = recv_prime_pp_signal(
4122                released_downstream,
4123                expected,
4124                true,
4125                "prime PP final slot release",
4126            )?;
4127            credits.record_release(released)?;
4128        }
4129        Ok(())
4130    }
4131
4132    #[allow(clippy::too_many_arguments)]
4133    fn prime_ppn_wave_stage(
4134        &self,
4135        e: &Engine,
4136        rt: &crate::pp::PpNRt,
4137        wave: &PrimePpWave<'_>,
4138        cache: &mut Cache,
4139        stage: usize,
4140        incoming: Option<PrimePpWaveSlot>,
4141        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
4142        credits: &mut PrimePpWaveCredits,
4143        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
4144        seq_end: usize,
4145        fence: &[usize],
4146        initial_base: usize,
4147    ) -> Result<PrimePpWaveSlot, Box<dyn std::error::Error>> {
4148        debug_assert!(stage + 1 < fence.len() - 1);
4149        let t = wave.end - wave.start;
4150        let base = initial_base + wave.start;
4151        debug_assert_eq!(cache.pos, base, "prime PP stage advanced out of order");
4152        let n_embd = self.cfg.n_embd as usize;
4153        let payload = t * n_embd;
4154        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4155        rt.bind_stage(stage)?;
4156        let _stage = rt.enter(stage);
4157        let engine = rt.engine(stage, e);
4158        let positions_d = engine.htod_i32(&positions)?;
4159        let x = if stage == 0 {
4160            debug_assert!(incoming.is_none());
4161            self.embed(engine, wave.tokens)?
4162        } else {
4163            let incoming = incoming.ok_or("prime PP stage has no incoming boundary slot")?;
4164            let x = rt.rx(stage - 1, incoming.slot, payload)?;
4165            send_prime_pp_signal(
4166                release_upstream.ok_or("prime PP stage has no upstream release channel")?,
4167                PrimePpSignal::Slot(incoming),
4168                "prime PP upstream slot release",
4169            )?;
4170            x
4171        };
4172        let x = {
4173            let _wave_cell = crate::pp::enter_pp_wave_cell();
4174            let _overlap = crate::pp::enter_prime_pipe_stage();
4175            self.prime_layers(
4176                engine,
4177                x,
4178                fence[stage],
4179                fence[stage + 1],
4180                &positions_d,
4181                t,
4182                base,
4183                cache,
4184                seq_end,
4185            )?
4186        };
4187        if let Some(expected) = credits.release_required() {
4188            let released =
4189                recv_prime_pp_signal(released_downstream, expected, true, "prime PP slot credit")?;
4190            credits.record_release(released)?;
4191        }
4192        let sent = PrimePpWaveSlot {
4193            wave: credits.next_wave,
4194            slot: rt.tx_pipelined(stage, &x, payload)?,
4195        };
4196        credits.record_send(sent)?;
4197        cache.pos = base + t;
4198        Ok(sent)
4199    }
4200
4201    #[allow(clippy::too_many_arguments)]
4202    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4203    fn prime_ppn_wave_final(
4204        &self,
4205        e: &Engine,
4206        rt: &crate::pp::PpNRt,
4207        wave: &PrimePpWave<'_>,
4208        cache: &mut Cache,
4209        incoming: PrimePpWaveSlot,
4210        release_upstream: &std::sync::mpsc::Sender<PrimePpSignal>,
4211        seq_end: usize,
4212        fence: &[usize],
4213        initial_base: usize,
4214    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4215        let stage = fence.len() - 2;
4216        let t = wave.end - wave.start;
4217        let base = initial_base + wave.start;
4218        debug_assert_eq!(cache.pos, base, "prime PP head stage advanced out of order");
4219        let n_embd = self.cfg.n_embd as usize;
4220        let payload = t * n_embd;
4221        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4222        rt.bind_stage(stage)?;
4223        let _stage = rt.enter(stage);
4224        let engine = rt.engine(stage, e);
4225        let positions_d = engine.htod_i32(&positions)?;
4226        let x = rt.rx(stage - 1, incoming.slot, payload)?;
4227        send_prime_pp_signal(
4228            release_upstream,
4229            PrimePpSignal::Slot(incoming),
4230            "prime PP head slot release",
4231        )?;
4232        let _wave_cell = crate::pp::enter_pp_wave_cell();
4233        let _overlap = crate::pp::enter_prime_pipe_stage();
4234        let x = self.prime_layers(
4235            engine,
4236            x,
4237            fence[stage],
4238            fence[stage + 1],
4239            &positions_d,
4240            t,
4241            base,
4242            cache,
4243            seq_end,
4244        )?;
4245        self.prime_chunk_epilogue(engine, x, t, cache)
4246    }
4247
4248    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
4249    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
4250    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
4251    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
4252    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
4253    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
4254    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
4255        if Engine::gdn_db_on()
4256            && Engine::gdn_chunked_enabled()
4257            && t >= 16
4258            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
4259            && num_k * 2 == num_v
4260        {
4261            num_k
4262        } else {
4263            num_v
4264        }
4265    }
4266
4267    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
4268    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
4269    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
4270    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
4271    fn f16out_on(e: &Engine, t: usize) -> bool {
4272        crate::f16_ffi::pp_f16_enabled()
4273            && t >= 16
4274            && !e.verify_exact_on()
4275            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
4276    }
4277
4278    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
4279    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
4280    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
4281    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
4282    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
4283    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
4284    /// see one entry, byte-identical behavior.
4285    pub fn prime_slabs_get(
4286        &self,
4287        e: &Engine,
4288        t: usize,
4289        n_embd: usize,
4290        n_ff_max: usize,
4291    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
4292        let mut slabs = self.prime_slabs.lock().unwrap();
4293        let dev = e.ctx().ordinal();
4294        let need_new = match slabs.get(&dev) {
4295            None => true,
4296            Some(sl) => sl.lock().unwrap().t_cap < t,
4297        };
4298        if need_new {
4299            slabs.insert(
4300                dev,
4301                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
4302                    t_cap: t,
4303                    h: e.uninit(t * n_embd)?,
4304                    x1: e.uninit(t * n_embd)?,
4305                    z: e.uninit(t * n_embd)?,
4306                    act: e.uninit(t * n_ff_max)?,
4307                    xa: e.uninit(t * n_embd)?,
4308                    xb: e.uninit(t * n_embd)?,
4309                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
4310                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
4311                    gate: e.uninit(t * n_ff_max)?,
4312                    up: e.uninit(t * n_ff_max)?,
4313                    ffn_out: e.uninit(t * n_embd)?,
4314                    seg_glue: Vec::new(),
4315                    mixed: e.uninit(t * n_embd)?,
4316                    seg_mid: Vec::new(),
4317                    seg_t: 0,
4318                })),
4319            );
4320        }
4321        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
4322    }
4323
4324    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
4325    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
4326    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
4327    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4328    fn prime_chunk(
4329        &self,
4330        e: &Engine,
4331        tokens: &[u32],
4332        cache: &mut Cache,
4333        seq_end: usize,
4334        chunk_off: usize,
4335        overlay: Option<&crate::vision::EmbedOverlay>,
4336    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4337        if crate::pp::pp_host_bounce_active()
4338            && (self.uses_gemma_program() || !crate::pp::prime_pp_on())
4339        {
4340            return Err(
4341                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
4342                 has no active prime stage split and would peer-read remote weights; keep \
4343                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
4344                    .into(),
4345            );
4346        }
4347        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
4348        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
4349        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
4350        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
4351        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
4352        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
4353        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
4354        // loader is off and there is nothing remote to split for.
4355        if !self.uses_gemma_program()
4356            && !crate::pp::pp2_streams_off()
4357            && crate::pp::prime_pp_on()
4358            && let Some(fence) = crate::pp::pp_cuts(self.layers.len())
4359        {
4360            if overlay.is_some() {
4361                return Err("vision embedding overlay + PP prime unsupported (v1); \
4362                         run single-device or MEMRA_PRIME_PP=0"
4363                    .into());
4364            }
4365            return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
4366        }
4367        if crate::pp::pp_host_bounce_active() {
4368            return Err(
4369                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
4370                 refusing an unsplit remote-weight walk"
4371                    .into(),
4372            );
4373        }
4374        let t = tokens.len();
4375        let base = cache.pos;
4376        debug_assert!(
4377            seq_end >= base + t,
4378            "prime_chunk: seq_end must cover this chunk"
4379        );
4380        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
4381        let pos_d = e.htod_i32(&pos)?;
4382
4383        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
4384        if let Some(ov) = overlay {
4385            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
4386            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
4387            // Images larger than one prime chunk straddle boundaries, hence the clipping.
4388            //
4389            // Through the SHARED `EmbedOverlay::splice_into` (lane/glm53-vision-ppn): this
4390            // site used to carry its own byte-identical copy of that loop, which meant the
4391            // overlay residency law had to be re-stated per call site. One implementation,
4392            // one law, and the splice point cannot drift between arms.
4393            ov.splice_into(e, &mut x_embed, chunk_off, t, self.cfg.n_embd as usize)?;
4394        }
4395        let x = self.prime_layers(
4396            e,
4397            x_embed,
4398            0,
4399            self.layers.len(),
4400            &pos_d,
4401            t,
4402            base,
4403            cache,
4404            seq_end,
4405        )?;
4406        self.prime_chunk_epilogue(e, x, t, cache)
4407    }
4408
4409    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
4410    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
4411    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
4412    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
4413    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
4414    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
4415    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
4416    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
4417    ///     the plain add (materialize) and the next stage hoists its own first norm — the
4418    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
4419    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
4420    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
4421    ///     each stage walks through its own resident transients;
4422    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
4423    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
4424    #[allow(clippy::too_many_arguments)]
4425    fn prime_layers(
4426        &self,
4427        e: &Engine,
4428        x_in: CudaSlice<f32>,
4429        lo: usize,
4430        hi: usize,
4431        pos_d: &CudaSlice<i32>,
4432        t: usize,
4433        base: usize,
4434        cache: &mut Cache,
4435        seq_end: usize,
4436    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4437        let cfg = &self.cfg;
4438        let n_embd = cfg.n_embd as usize;
4439        let eps = cfg.rms_eps;
4440        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
4441        // standalone convert launches). Only when the f16 lane serves and T reaches the
4442        // GEMM tier; bit-identical either way.
4443        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
4444        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
4445        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
4446        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
4447        // live prefix is fully overwritten before use; x ping-pongs xa<->xb; the inactive
4448        // capacity tail must stay behind checked views. The hidden-stack return clones the
4449        // final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
4450        let n_ff_max = self
4451            .layers
4452            .iter()
4453            .map(|l| match &l.ffn {
4454                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
4455                _ => n_embd,
4456            })
4457            .max()
4458            .unwrap_or(n_embd)
4459            .max(n_embd);
4460        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
4461        let slab = if use_slabs {
4462            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
4463        } else {
4464            None
4465        };
4466        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
4467        let mut x_own; // fallback storage when slabs are off
4468        type SlabRefs<'a> = (
4469            &'a mut CudaSlice<f32>,
4470            &'a mut CudaSlice<f32>,
4471            &'a mut CudaSlice<f32>,
4472            &'a mut CudaSlice<f32>,
4473            &'a mut CudaSlice<u8>,
4474            &'a mut CudaSlice<u8>,
4475            &'a mut CudaSlice<f32>,
4476            &'a mut CudaSlice<f32>,
4477            &'a mut CudaSlice<f32>,
4478        );
4479        let (mut x_cur, mut x_nxt, sl): (
4480            &mut CudaSlice<f32>,
4481            &mut CudaSlice<f32>,
4482            Option<SlabRefs>,
4483        );
4484        #[allow(clippy::type_complexity)]
4485        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4486        let mut seg: Option<(
4487            &mut Vec<Option<cudarc::driver::CudaGraph>>,
4488            &mut Vec<Option<cudarc::driver::CudaGraph>>,
4489            &mut CudaSlice<f32>,
4490            &mut usize,
4491        )> = None;
4492        let mut x_own2;
4493        match slab_guard.as_mut() {
4494            Some(g) => {
4495                let slabs = &mut **g;
4496                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
4497                let PrimeSlabs {
4498                    xa,
4499                    xb,
4500                    h,
4501                    x1,
4502                    z,
4503                    act,
4504                    h16,
4505                    z16,
4506                    gate,
4507                    up,
4508                    ffn_out,
4509                    seg_glue,
4510                    mixed,
4511                    seg_mid,
4512                    seg_t,
4513                    ..
4514                } = slabs;
4515                x_cur = xa;
4516                x_nxt = xb;
4517                seg = Some((seg_glue, seg_mid, mixed, seg_t));
4518                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
4519            }
4520            None => {
4521                x_own = x_in;
4522                x_own2 = e.uninit(t * n_embd)?;
4523                x_cur = &mut x_own;
4524                x_nxt = &mut x_own2;
4525                sl = None;
4526            }
4527        }
4528        let mut alloc_h;
4529        let mut alloc_x1;
4530        let mut alloc_z;
4531        let mut alloc_act;
4532        let mut alloc_h16;
4533        let mut alloc_z16;
4534        let mut alloc_gate;
4535        let mut alloc_up;
4536        let mut alloc_fo;
4537        let (h, x1, z, act): (
4538            &mut CudaSlice<f32>,
4539            &mut CudaSlice<f32>,
4540            &mut CudaSlice<f32>,
4541            &mut CudaSlice<f32>,
4542        );
4543        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
4544        let (sl_gate, sl_up, sl_fo): (
4545            &mut CudaSlice<f32>,
4546            &mut CudaSlice<f32>,
4547            &mut CudaSlice<f32>,
4548        );
4549        match sl {
4550            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
4551                h = a;
4552                x1 = b;
4553                z = c;
4554                act = d;
4555                h16 = e16;
4556                z16 = f16b;
4557                sl_gate = g;
4558                sl_up = u;
4559                sl_fo = fo;
4560            }
4561            None => {
4562                alloc_h = e.uninit(t * n_embd)?;
4563                alloc_x1 = e.uninit(t * n_embd)?;
4564                alloc_z = e.uninit(t * n_embd)?;
4565                alloc_act = e.uninit(t * n_ff_max)?;
4566                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4567                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4568                alloc_gate = e.uninit(t * n_ff_max)?;
4569                alloc_up = e.uninit(t * n_ff_max)?;
4570                alloc_fo = e.uninit(t * n_embd)?;
4571                h = &mut alloc_h;
4572                x1 = &mut alloc_x1;
4573                z = &mut alloc_z;
4574                act = &mut alloc_act;
4575                h16 = &mut alloc_h16;
4576                z16 = &mut alloc_z16;
4577                sl_gate = &mut alloc_gate;
4578                sl_up = &mut alloc_up;
4579                sl_fo = &mut alloc_fo;
4580            }
4581        }
4582        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
4583        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
4584        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
4585        // first prime at this t (capture does not execute -> launch right after).
4586        let n_layers = self.layers.len();
4587        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
4588        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
4589        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
4590        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
4591        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
4592        // machinery stays (byte-identical) as their foundation.
4593        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
4594        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
4595        // step35 rides its own mixer through the normal per-layer arm below.
4596        let use_seg = f16fuse
4597            && seg.is_some()
4598            && !self.uses_sliding_gated_moe_program()
4599            && lo == 0
4600            && hi == n_layers
4601            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1")
4602            // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): below the
4603            // driver-free floor this prime call takes the eager fused else-arm (the
4604            // byte-identical twin the opt-in was gated against) instead of replaying
4605            // the S-mid/S-glue segment graphs into an exhausted card. Probe runs only
4606            // when the opt-in flag is armed (short-circuit order).
4607            && {
4608                let ok = crate::spec::graph_launch_headroom_ok(e);
4609                if !ok {
4610                    static NOTED: std::sync::Once = std::sync::Once::new();
4611                    NOTED.call_once(|| crate::spec::graph_replay_suspended_note("prime-seg"));
4612                }
4613                ok
4614            };
4615        if let Some((sg, sm, _, st)) = seg.as_mut()
4616            && **st != t
4617        {
4618            sg.clear();
4619            sg.extend((0..n_layers).map(|_| None));
4620            sm.clear();
4621            sm.extend((0..n_layers).map(|_| None));
4622            **st = t;
4623        }
4624        {
4625            let layer_lo = &self.layers[lo];
4626            if f16fuse {
4627                e.rms_norm_f16out(
4628                    x_cur,
4629                    layer_lo.attn_norm.float_data(),
4630                    h,
4631                    h16,
4632                    n_embd,
4633                    t,
4634                    eps,
4635                )?;
4636            } else {
4637                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
4638            }
4639        }
4640        let anat = Self::prime_anatomy_on();
4641        let mut anat_last = if anat {
4642            e.stream().synchronize()?;
4643            Some(std::time::Instant::now())
4644        } else {
4645            None
4646        };
4647        // Closes the region that just ENDED into `slot`, restarting the clock.
4648        macro_rules! anat_mark {
4649            ($slot:expr) => {
4650                if let Some(ts) = anat_last.as_mut() {
4651                    e.stream().synchronize()?;
4652                    Self::prime_anatomy_slots()[$slot].fetch_add(
4653                        ts.elapsed().as_nanos() as u64,
4654                        std::sync::atomic::Ordering::Relaxed,
4655                    );
4656                    *ts = std::time::Instant::now();
4657                }
4658            };
4659        }
4660        for il in lo..hi {
4661            let layer = &self.layers[il];
4662            let hx16 = if f16fuse { Some(&*h16) } else { None };
4663            if use_seg {
4664                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
4665                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
4666                let (pre, pre16, w_out) = match &layer.mixer {
4667                    Mixer::Full(fa) => {
4668                        let g3 = match hx16 {
4669                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
4670                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
4671                        };
4672                        let (pre, pre16) =
4673                            self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
4674                        (pre, pre16, &fa.wo)
4675                    }
4676                    Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("core-split prime"),
4677                    Mixer::Kda(_) => {
4678                        crate::hybrid::kda_path_unimplemented("core-split captured prime")
4679                    }
4680                    Mixer::Linear(la) => {
4681                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
4682                        let g4 = match hx16 {
4683                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
4684                            None => e.matmul_group(&ws, h, t)?,
4685                        };
4686                        let (pre, pre16) =
4687                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
4688                        (pre, pre16, &la.ssm_out)
4689                    }
4690                };
4691                {
4692                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
4693                    let pre_n = pre.len() / t;
4694                    let xh_pre = match pre16 {
4695                        Some(x) => x,
4696                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
4697                    };
4698                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
4699                        let y = e.matmul(w_out, &pre, t)?;
4700                        e.copy_into(mslab, 0, &y, t * n_embd)?;
4701                    }
4702                    if sm[il].is_none() {
4703                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4704                        let w_post = layer.post_attn_norm.float_data();
4705                        e.stream().synchronize()?;
4706                        e.stream()
4707                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4708                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4709                            e.add(x_cur, mslab, x1, t * n_embd)?;
4710                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
4711                            Ok(())
4712                        })();
4713                        let g = e.stream().end_capture(
4714                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
4715                        r?;
4716                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
4717                    }
4718                    sm[il].as_ref().unwrap().launch()?;
4719                }
4720            } else {
4721                let mixed = match &layer.mixer {
4722                    Mixer::Full(fa) => {
4723                        let y =
4724                            self.full_attn_prime(e, fa, h, hx16, pos_d, t, cache, il, seq_end)?;
4725                        anat_mark!(0);
4726                        y
4727                    }
4728                    Mixer::Linear(la) => {
4729                        let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
4730                        anat_mark!(1);
4731                        y
4732                    }
4733                    Mixer::Mla(mla) => self.mla_attn_cached(e, mla, h, pos_d, t, il, cache)?,
4734                    Mixer::Kda(la) => {
4735                        let y = crate::kda::kda_prime_cached(e, la, h, t, eps, cache, il)?;
4736                        anat_mark!(1);
4737                        y
4738                    }
4739                };
4740                if f16fuse {
4741                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
4742                    // bit-identical) — the standalone add pass disappears.
4743                    e.add_rms_norm_f16out(
4744                        x_cur,
4745                        &mixed,
4746                        layer.post_attn_norm.float_data(),
4747                        x1,
4748                        z,
4749                        z16,
4750                        n_embd,
4751                        t,
4752                        eps,
4753                    )?;
4754                } else {
4755                    e.add(x_cur, &mixed, x1, t * n_embd)?;
4756                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
4757                }
4758                anat_mark!(4);
4759            }
4760            let zx16 = if f16fuse { Some(&*z16) } else { None };
4761            match &layer.ffn {
4762                crate::hybrid::Ffn::Dense {
4763                    ffn_gate,
4764                    ffn_up,
4765                    ffn_down,
4766                } => {
4767                    let n_ff = ffn_gate.out_features();
4768                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
4769                    // the allocating group + copy when a mirror is missing.
4770                    let mut into_ok = false;
4771                    if let Some(xh) = zx16 {
4772                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
4773                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
4774                    }
4775                    if !into_ok {
4776                        let mut g2 = match zx16 {
4777                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
4778                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
4779                        };
4780                        let up_y = g2.pop().unwrap();
4781                        let gate_y = g2.pop().unwrap();
4782                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
4783                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
4784                    }
4785                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
4786                    // operand in-epilogue; non-silu activations keep the standalone convert.
4787                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
4788                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
4789                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
4790                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
4791                    {
4792                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
4793                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
4794                        Some(a16)
4795                    } else {
4796                        Self::ffn_act_lim(
4797                            e,
4798                            &self.cfg,
4799                            sl_gate,
4800                            sl_up,
4801                            1.0,
4802                            1.0,
4803                            d_lim,
4804                            act,
4805                            t * n_ff,
4806                        )?;
4807                        None
4808                    };
4809                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
4810                    let xh_act = match act16 {
4811                        Some(x) => x,
4812                        None => e.f16_act(act, t * n_ff, n_ff)?,
4813                    };
4814                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
4815                        let y = e.matmul(ffn_down, &*act, t)?;
4816                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4817                    }
4818                }
4819                crate::hybrid::Ffn::Moe(m) => {
4820                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
4821                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4822                    anat_mark!(2);
4823                }
4824            }
4825            if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
4826                anat_mark!(3);
4827            }
4828            if use_seg && il + 1 < hi {
4829                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
4830                let w_next = self.layers[il + 1].attn_norm.float_data();
4831                let (sg, _, _, _) = seg.as_mut().unwrap();
4832                if sg[il].is_none() {
4833                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4834                    e.stream().synchronize()?;
4835                    e.stream()
4836                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4837                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4838                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4839                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
4840                        Ok(())
4841                    })();
4842                    let g = e.stream().end_capture(
4843                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
4844                    );
4845                    r?;
4846                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
4847                }
4848                sg[il].as_ref().unwrap().launch()?;
4849            } else {
4850                if il + 1 < hi {
4851                    let w_next = self.layers[il + 1].attn_norm.float_data();
4852                    if f16fuse {
4853                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
4854                    } else {
4855                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4856                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
4857                    }
4858                } else {
4859                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4860                }
4861            }
4862            anat_mark!(4);
4863            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
4864            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
4865            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
4866            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
4867            // unset (the default) costs one OnceLock read per layer.
4868            if let Some(path) = Self::prime_trace_path() {
4869                let row = base + t - 1;
4870                let host = e.dtoh(x_nxt)?;
4871                let last = &host[(t - 1) * n_embd..t * n_embd];
4872                use std::io::Write as _;
4873                let mut f = std::fs::OpenOptions::new()
4874                    .create(true)
4875                    .append(true)
4876                    .open(path)?;
4877                let mut h64: u64 = 0xcbf29ce484222325;
4878                for v in last {
4879                    h64 ^= v.to_bits() as u64;
4880                    h64 = h64.wrapping_mul(0x100000001b3);
4881                }
4882                writeln!(
4883                    f,
4884                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
4885                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
4886                    last[0], last[1], last[2]
4887                )?;
4888            }
4889            // dflash/dspark tap (no-op when no sink armed): post-layer residual rows for
4890            // drafter conditioning — the qwen twin of the gemma4 tap sites.
4891            self.dflash_tap(e, cache, il, x_nxt, t)?;
4892            std::mem::swap(&mut x_cur, &mut x_nxt);
4893        }
4894        if anat {
4895            let s = Self::prime_anatomy_slots();
4896            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
4897            eprintln!(
4898                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
4899                 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
4900                ms(0),
4901                ms(1),
4902                ms(2),
4903                ms(3),
4904                ms(4)
4905            );
4906        }
4907        // hidden-stack return: clone the final x out of the slab
4908        let mut x = e.uninit(t * n_embd)?;
4909        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
4910        drop(slab_guard);
4911        Ok(x)
4912    }
4913
4914    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
4915    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
4916    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
4917    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
4918    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4919    fn prime_chunk_epilogue(
4920        &self,
4921        e: &Engine,
4922        x: CudaSlice<f32>,
4923        t: usize,
4924        cache: &mut Cache,
4925    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4926        let n_embd = self.cfg.n_embd as usize;
4927        let eps = self.cfg.rms_eps;
4928        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
4929        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
4930        // the post-norm copy happens after hn exists).
4931        let mut h_seed = e.uninit(n_embd)?;
4932        if !crate::spec::spec_hpost() {
4933            e.copy_view_into(
4934                &mut h_seed,
4935                0,
4936                &x.slice((t - 1) * n_embd..t * n_embd),
4937                n_embd,
4938            )?;
4939        }
4940        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
4941        let mut hn = e.uninit(t * n_embd)?;
4942        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4943        if crate::spec::spec_hpost() {
4944            e.copy_view_into(
4945                &mut h_seed,
4946                0,
4947                &hn.slice((t - 1) * n_embd..t * n_embd),
4948                n_embd,
4949            )?;
4950        }
4951        let last = e.view(&hn, t * n_embd);
4952        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
4953        let mut hlast = e.uninit(n_embd)?;
4954        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
4955        let logits = e.matmul(&self.output, &hlast, 1)?;
4956        cache.pos += t;
4957        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
4958        // post-norm stack hn (MEMRA_SPEC_HPOST).
4959        Ok((
4960            e.dtoh(&logits)?,
4961            h_seed,
4962            if crate::spec::spec_hpost() { hn } else { x },
4963        ))
4964    }
4965
4966    /// Post-final-norm hidden state of one row of a prime-returned hidden stack — the
4967    /// embedding-pooling read (lane/embed-serve). `hiddens` is `prime_cache`'s third
4968    /// return: the pre-norm stack by default, but ALREADY post-norm under
4969    /// MEMRA_SPEC_HPOST (see `prime_chunk_epilogue`), so the norm is applied only in
4970    /// the default shape. Returns the host f32 row (`n_embd` wide).
4971    pub fn hidden_postnorm_row(
4972        &self,
4973        e: &Engine,
4974        hiddens: &CudaSlice<f32>,
4975        row: usize,
4976    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4977        let n_embd = self.cfg.n_embd as usize;
4978        let mut x1 = e.uninit(n_embd)?;
4979        e.copy_view_into(
4980            &mut x1,
4981            0,
4982            &hiddens.slice(row * n_embd..(row + 1) * n_embd),
4983            n_embd,
4984        )?;
4985        if crate::spec::spec_hpost() {
4986            return e.dtoh(&x1);
4987        }
4988        let mut hn = e.uninit(n_embd)?;
4989        e.rms_norm(
4990            &x1,
4991            self.output_norm.float_data(),
4992            &mut hn,
4993            n_embd,
4994            1,
4995            self.cfg.rms_eps,
4996        )?;
4997        e.dtoh(&hn)
4998    }
4999
5000    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
5001    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
5002    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
5003    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
5004    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
5005    /// prefill kernels. Structure mirrors the verify split exactly:
5006    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
5007    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
5008    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
5009    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
5010    ///                  there via the sharded loader) → `publish_to`
5011    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
5012    /// round's stage-freed buffers must not be reused under the caller's queued reads);
5013    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
5014    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
5015    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
5016    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
5017    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
5018    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
5019    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
5020    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
5021    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
5022    /// and its liveness counter is bumped here — the gate goes green with this function.
5023    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5024    fn prime_chunk_ppn(
5025        &self,
5026        e: &Engine,
5027        tokens: &[u32],
5028        cache: &mut Cache,
5029        seq_end: usize,
5030        fence: &[usize],
5031    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5032        let rt = crate::pp::PpNRt::get(e)?;
5033        let n_st = fence.len() - 1;
5034        assert_eq!(
5035            rt.n_stages(),
5036            n_st,
5037            "PpNRt stage count {} != fence stages {n_st}",
5038            rt.n_stages()
5039        );
5040        let n_embd = self.cfg.n_embd as usize;
5041        let t = tokens.len();
5042        let base = cache.pos;
5043        debug_assert!(
5044            seq_end >= base + t,
5045            "prime_chunk_ppn: seq_end must cover this chunk"
5046        );
5047        let payload = t * n_embd;
5048        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
5049        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
5050        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
5051        let caller_stream = e.stream();
5052        rt.fence_stages_behind(&caller_stream)?;
5053
5054        if n_st == 2 {
5055            let slot =
5056                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
5057            let x =
5058                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
5059            let out = {
5060                rt.bind_stage(1)?;
5061                let _st1 = rt.enter(1);
5062                let e1 = rt.engine(1, e);
5063                self.prime_chunk_epilogue(e1, x, t, cache)?
5064            };
5065            rt.publish_to(1, &caller_stream)?;
5066            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5067            return Ok(out);
5068        }
5069
5070        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5071
5072        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
5073        let mut slot = {
5074            let _st0 = rt.enter(0);
5075            let e0 = rt.engine(0, e);
5076            let pos_d = e0.htod_i32(&pos)?;
5077            let x = self.embed(e0, tokens)?;
5078            let x =
5079                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
5080            rt.tx(0, &x, payload)?
5081            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5082        };
5083
5084        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5085        for s in 1..n_st - 1 {
5086            let _st = rt.enter(s);
5087            let es = rt.engine(s, e);
5088            let pos_d = es.htod_i32(&pos)?;
5089            let x = rt.rx(s - 1, slot, payload)?;
5090            let x = self.prime_layers(
5091                es,
5092                x,
5093                fence[s],
5094                fence[s + 1],
5095                &pos_d,
5096                t,
5097                base,
5098                cache,
5099                seq_end,
5100            )?;
5101            slot = rt.tx(s, &x, payload)?;
5102        }
5103
5104        // ---- LAST STAGE: RX + final range + the shared epilogue ----
5105        let _stl = rt.enter(n_st - 1);
5106        let el = rt.engine(n_st - 1, e);
5107        let pos_d = el.htod_i32(&pos)?;
5108        let x = rt.rx(n_st - 2, slot, payload)?;
5109        let x = self.prime_layers(
5110            el,
5111            x,
5112            fence[n_st - 1],
5113            fence[n_st],
5114            &pos_d,
5115            t,
5116            base,
5117            cache,
5118            seq_end,
5119        )?;
5120        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
5121        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
5122        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
5123        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
5124        // stage stream host-side, but the law is stated in events, not in a dtoh side
5125        // effect a later deferred form would remove.
5126        rt.publish_to(n_st - 1, &caller_stream)?;
5127        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5128        Ok(out)
5129    }
5130
5131    #[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
5132    fn prime_pp2_stage0_enqueue(
5133        &self,
5134        e: &Engine,
5135        rt: &crate::pp::PpNRt,
5136        tokens: &[u32],
5137        cache: &mut Cache,
5138        seq_end: usize,
5139        fence: &[usize],
5140        base: usize,
5141        pipelined: bool,
5142    ) -> Result<usize, Box<dyn std::error::Error>> {
5143        let t = tokens.len();
5144        let n_embd = self.cfg.n_embd as usize;
5145        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5146        rt.bind_stage(0)?;
5147        let _st0 = rt.enter(0);
5148        let e0 = rt.engine(0, e);
5149        let pos_d = e0.htod_i32(&pos)?;
5150        let x = self.embed(e0, tokens)?;
5151        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5152        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
5153        if pipelined {
5154            rt.tx_pipelined(0, &x, t * n_embd)
5155        } else {
5156            rt.tx(0, &x, t * n_embd)
5157        }
5158    }
5159
5160    #[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
5161    fn prime_pp2_stage1_enqueue(
5162        &self,
5163        e: &Engine,
5164        rt: &crate::pp::PpNRt,
5165        slot: usize,
5166        t: usize,
5167        cache: &mut Cache,
5168        seq_end: usize,
5169        fence: &[usize],
5170        base: usize,
5171        pipelined: bool,
5172    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5173        let n_embd = self.cfg.n_embd as usize;
5174        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5175        rt.bind_stage(1)?;
5176        let _st1 = rt.enter(1);
5177        let e1 = rt.engine(1, e);
5178        let pos_d = e1.htod_i32(&pos)?;
5179        let x = rt.rx(0, slot, t * n_embd)?;
5180        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5181        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
5182    }
5183
5184    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
5185    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
5186    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
5187    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
5188    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
5189    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
5190    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
5191    /// bookkeeping still runs on the host per call — the real replay path moves the write
5192    /// slot to the len_d device counter (increment 3).
5193    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
5194    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
5195    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
5196    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
5197    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
5198    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
5199    #[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
5200    pub fn prime_chunk_captured(
5201        &self,
5202        e: &Engine,
5203        x_in: &CudaSlice<f32>,
5204        pos_d: &CudaSlice<i32>,
5205        t: usize,
5206        cache: &mut Cache,
5207        len_d: &CudaSlice<i32>,
5208        logits_out: &mut CudaSlice<f32>,
5209        h_seed_out: &mut CudaSlice<f32>,
5210    ) -> Result<(), Box<dyn std::error::Error>> {
5211        self.refuse_hyper("prime_chunk_captured")?;
5212        cache.ensure_usable("prime_chunk_captured")?;
5213        let cfg = &self.cfg;
5214        let n_embd = cfg.n_embd as usize;
5215        let eps = cfg.rms_eps;
5216        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
5217        let mut x = e.uninit(t * n_embd)?;
5218        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
5219        for (il, layer) in self.layers.iter().enumerate() {
5220            let mut h = e.uninit(t * n_embd)?;
5221            let mut hx16: Option<CudaSlice<u8>> = None;
5222            if f16fuse {
5223                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5224                e.rms_norm_f16out(
5225                    &x,
5226                    layer.attn_norm.float_data(),
5227                    &mut h,
5228                    &mut b16,
5229                    n_embd,
5230                    t,
5231                    eps,
5232                )?;
5233                hx16 = Some(b16);
5234            } else {
5235                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5236            }
5237            let mixed = match &layer.mixer {
5238                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
5239                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
5240                // come from the caller (see step35_attn_pre_wo's doc note).
5241                Mixer::Full(fa) => {
5242                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
5243                }
5244                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("captured-graph prime"),
5245                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("captured prime chunk"),
5246                Mixer::Linear(la) => {
5247                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
5248                    let g4 = match hx16.as_ref() {
5249                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
5250                        None => e.matmul_group(&ws, &h, t)?,
5251                    };
5252                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
5253                }
5254            };
5255            let mut x1 = e.uninit(t * n_embd)?;
5256            e.add(&x, &mixed, &mut x1, t * n_embd)?;
5257            let mut z = e.uninit(t * n_embd)?;
5258            let mut zx16: Option<CudaSlice<u8>> = None;
5259            if f16fuse {
5260                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5261                e.rms_norm_f16out(
5262                    &x1,
5263                    layer.post_attn_norm.float_data(),
5264                    &mut z,
5265                    &mut b16,
5266                    n_embd,
5267                    t,
5268                    eps,
5269                )?;
5270                zx16 = Some(b16);
5271            } else {
5272                e.rms_norm(
5273                    &x1,
5274                    layer.post_attn_norm.float_data(),
5275                    &mut z,
5276                    n_embd,
5277                    t,
5278                    eps,
5279                )?;
5280            }
5281            let ffn_out = match &layer.ffn {
5282                crate::hybrid::Ffn::Dense {
5283                    ffn_gate,
5284                    ffn_up,
5285                    ffn_down,
5286                } => {
5287                    let n_ff = ffn_gate.out_features();
5288                    let mut g2 = match &zx16 {
5289                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
5290                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
5291                    };
5292                    let up = g2.pop().unwrap();
5293                    let gate = g2.pop().unwrap();
5294                    let mut act = e.uninit(t * n_ff)?;
5295                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
5296                    Self::ffn_act_lim(
5297                        e,
5298                        &self.cfg,
5299                        &gate,
5300                        &up,
5301                        1.0,
5302                        1.0,
5303                        self.cfg.clamp_shexp_at(il as u32),
5304                        &mut act,
5305                        t * n_ff,
5306                    )?;
5307                    e.matmul(ffn_down, &act, t)?
5308                }
5309                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
5310            };
5311            let mut x2 = e.uninit(t * n_embd)?;
5312            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5313            x = x2;
5314        }
5315        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
5316        if !crate::spec::spec_hpost() {
5317            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
5318        }
5319        let mut hn = e.uninit(t * n_embd)?;
5320        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5321        if crate::spec::spec_hpost() {
5322            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
5323        }
5324        let mut hlast = e.uninit(n_embd)?;
5325        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
5326        let logits = e.matmul(&self.output, &hlast, 1)?;
5327        let nv = logits.len();
5328        e.copy_into(logits_out, 0, &logits, nv)?;
5329        Ok(())
5330    }
5331
5332    fn step35_prime_batch_on() -> bool {
5333        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
5334    }
5335
5336    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
5337    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
5338    #[allow(clippy::too_many_arguments)]
5339    /// `seq_ends[s]`: sequence s's REQUEST-absolute end position — NOT `ts[s]`. It is the
5340    /// only thing step35's SWA arm keys on, so a chunk-local value here decides the attention
5341    /// kernel from the chunk size (and, below the 512-row window at a nonzero base, drops the
5342    /// window mask entirely). See the batched entry's note in `prime_cache_overlaid`.
5343    #[allow(clippy::too_many_arguments)]
5344    fn step35_prime_batch_layers(
5345        &self,
5346        e: &Engine,
5347        mut x: CudaSlice<f32>,
5348        lo: usize,
5349        hi: usize,
5350        ts: &[usize],
5351        offs: &[usize],
5352        seq_ends: &[usize],
5353        pos_ds: &[CudaSlice<i32>],
5354        caches: &mut [&mut Cache],
5355    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5356        let cfg = &self.cfg;
5357        let n_embd = cfg.n_embd as usize;
5358        let eps = cfg.rms_eps;
5359        let b = ts.len();
5360        let total: usize = ts.iter().sum();
5361        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
5362
5363        let split = |e: &Engine,
5364                     y: &CudaSlice<f32>,
5365                     dim: usize|
5366         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5367            let mut out = Vec::with_capacity(b);
5368            for s in 0..b {
5369                let mut ys = e.uninit(ts[s] * dim)?;
5370                e.copy_view_into(
5371                    &mut ys,
5372                    0,
5373                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5374                    ts[s] * dim,
5375                )?;
5376                out.push(ys);
5377            }
5378            Ok(out)
5379        };
5380
5381        // MEMRA_PRIME_PROF=1: per-phase wall inside the prime, sync-bounded (absolute time
5382        // inflates; the SPLIT is the signal). Two inspection passes failed to find where a
5383        // 3.8 s/4096-token chunk goes against a ~0.55 s compute budget, and nsys cannot capture
5384        // through the server's worker, so the walk measures itself.
5385        let prof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
5386        let mut ph = [0f64; 4]; // 0 norm+qkv, 1 attn, 2 o_proj+norm, 3 moe
5387        let mark = |e: &Engine, acc: usize, t0: &mut std::time::Instant, ph: &mut [f64; 4]| {
5388            if prof {
5389                let _ = e.stream().synchronize();
5390                ph[acc] += t0.elapsed().as_secs_f64() * 1e3;
5391                *t0 = std::time::Instant::now();
5392            }
5393        };
5394        let mut pt = std::time::Instant::now();
5395        for il in lo..hi {
5396            let layer = &self.layers[il];
5397            let Mixer::Full(fa) = &layer.mixer else {
5398                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
5399            };
5400
5401            let mut h = e.uninit(total * n_embd)?;
5402            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5403            if f16fuse {
5404                e.rms_norm_f16out(
5405                    &x,
5406                    layer.attn_norm.float_data(),
5407                    &mut h,
5408                    &mut hx16,
5409                    n_embd,
5410                    total,
5411                    eps,
5412                )?;
5413            } else {
5414                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
5415            }
5416
5417            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
5418            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
5419            // application stay verbatim.
5420            let gate_w = fa
5421                .attn_gate
5422                .as_ref()
5423                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
5424            let mut g4 = if f16fuse {
5425                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
5426            } else {
5427                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
5428            };
5429            let gate = g4.pop().unwrap();
5430            let mut parts: Vec<Vec<CudaSlice<f32>>> =
5431                (0..b).map(|_| Vec::with_capacity(3)).collect();
5432            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
5433                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
5434                    parts[s].push(ys);
5435                }
5436            }
5437            let gates = split(e, &gate, gate_w.out_features())?;
5438            let geometry = self.step35_geom(il);
5439            let hd = geometry.head_dim_k as usize;
5440            let nh = geometry.n_head as usize;
5441            let mut ag_cat = e.uninit(total * nh * hd)?;
5442            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
5443                mark(e, 0, &mut pt, &mut ph);
5444                let ag = self.step35_attn_pre_wo(
5445                    e,
5446                    fa,
5447                    g3s,
5448                    None,
5449                    Some(&gate),
5450                    &pos_ds[s],
5451                    ts[s],
5452                    Some(&mut *caches[s]),
5453                    il,
5454                    seq_ends[s],
5455                )?;
5456                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
5457            }
5458            mark(e, 1, &mut pt, &mut ph);
5459            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
5460
5461            let mut x1 = e.uninit(total * n_embd)?;
5462            let mut z = e.uninit(total * n_embd)?;
5463            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5464            if f16fuse {
5465                e.add_rms_norm_f16out(
5466                    &x,
5467                    &mixed,
5468                    layer.post_attn_norm.float_data(),
5469                    &mut x1,
5470                    &mut z,
5471                    &mut zx16,
5472                    n_embd,
5473                    total,
5474                    eps,
5475                )?;
5476            } else {
5477                e.add(&x, &mixed, &mut x1, total * n_embd)?;
5478                e.rms_norm(
5479                    &x1,
5480                    layer.post_attn_norm.float_data(),
5481                    &mut z,
5482                    n_embd,
5483                    total,
5484                    eps,
5485                )?;
5486            }
5487
5488            mark(e, 2, &mut pt, &mut ph);
5489            let ffn_out = match &layer.ffn {
5490                crate::hybrid::Ffn::Dense {
5491                    ffn_gate,
5492                    ffn_up,
5493                    ffn_down,
5494                } => {
5495                    let n_ff = ffn_gate.out_features();
5496                    let mut g2 = if f16fuse {
5497                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
5498                    } else {
5499                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
5500                    };
5501                    let up = g2.pop().unwrap();
5502                    let gate = g2.pop().unwrap();
5503                    let mut act = e.uninit(total * n_ff)?;
5504                    let d_lim = cfg.clamp_shexp_at(il as u32);
5505                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
5506                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
5507                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
5508                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
5509                            Some(y) => y,
5510                            None => e.matmul(ffn_down, &act, total)?,
5511                        }
5512                    } else {
5513                        Self::ffn_act_lim(
5514                            e,
5515                            cfg,
5516                            &gate,
5517                            &up,
5518                            1.0,
5519                            1.0,
5520                            d_lim,
5521                            &mut act,
5522                            total * n_ff,
5523                        )?;
5524                        e.matmul(ffn_down, &act, total)?
5525                    }
5526                }
5527                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
5528            };
5529            let mut x2 = e.uninit(total * n_embd)?;
5530            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
5531            x = x2;
5532            mark(e, 3, &mut pt, &mut ph);
5533        }
5534        if prof {
5535            eprintln!(
5536                "[prime-prof] t={total} layers={} norm+qkv={:.0}ms attn={:.0}ms o_proj={:.0}ms moe={:.0}ms",
5537                hi - lo,
5538                ph[0],
5539                ph[1],
5540                ph[2],
5541                ph[3]
5542            );
5543        }
5544        Ok(x)
5545    }
5546
5547    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5548    fn step35_prime_batch_epilogue(
5549        &self,
5550        e: &Engine,
5551        x: CudaSlice<f32>,
5552        ts: &[usize],
5553        offs: &[usize],
5554        caches: &mut [&mut Cache],
5555    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5556        let n_embd = self.cfg.n_embd as usize;
5557        let total: usize = ts.iter().sum();
5558        let mut hn = e.uninit(total * n_embd)?;
5559        e.rms_norm(
5560            &x,
5561            self.output_norm.float_data(),
5562            &mut hn,
5563            n_embd,
5564            total,
5565            self.cfg.rms_eps,
5566        )?;
5567
5568        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
5569        let mut out = Vec::with_capacity(ts.len());
5570        for s in 0..ts.len() {
5571            let mut hidden = e.uninit(ts[s] * n_embd)?;
5572            e.copy_view_into(
5573                &mut hidden,
5574                0,
5575                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
5576                ts[s] * n_embd,
5577            )?;
5578            let last0 = (offs[s] + ts[s] - 1) * n_embd;
5579            let mut h_seed = e.uninit(n_embd)?;
5580            e.copy_view_into(
5581                &mut h_seed,
5582                0,
5583                &hidden_src.slice(last0..last0 + n_embd),
5584                n_embd,
5585            )?;
5586            // Exactness-first: the serial reference runs the output head at m=1.
5587            let mut hlast = e.uninit(n_embd)?;
5588            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
5589            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
5590            caches[s].pos += ts[s];
5591            out.push((logits, h_seed, hidden));
5592        }
5593        Ok(out)
5594    }
5595
5596    /// `seq_ends[s]` = sequence s's REQUEST-absolute end position (`cache.pos + prompt_len
5597    /// + queued_after`, computed once before any chunk loop). Only step35's SWA arm reads it,
5598    /// and it must NOT be this chunk's own length: see the note on the batched entry in
5599    /// `prime_cache_overlaid` for the window the chunk-local value opened.
5600    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5601    fn step35_prime_cache_batch(
5602        &self,
5603        e: &Engine,
5604        prompts: &[&[u32]],
5605        caches: &mut [&mut Cache],
5606        seq_ends: &[usize],
5607    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5608        assert_eq!(
5609            seq_ends.len(),
5610            prompts.len(),
5611            "step35 batched prime: one seq_end per sequence"
5612        );
5613        validate_step_prime_batch_modes(
5614            step_tp_prefill_enabled()?,
5615            step_ep_grouped_prefill_enabled()?,
5616        )?;
5617        if crate::pp::pp_host_bounce_active()
5618            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
5619        {
5620            return Err(
5621                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
5622                 stage split; refusing an unsplit remote-weight walk"
5623                    .into(),
5624            );
5625        }
5626        if !Self::step35_prime_batch_on() {
5627            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
5628        }
5629        // Continuation chunks are admitted (positions above carry each sequence's base). The
5630        // remaining restriction is genuine: a CROSS-REQUEST batch mixing sequences at different
5631        // positions still needs per-request queued_after to place its KV, so B > 1 keeps the
5632        // fresh-prompt rule.
5633        if prompts.len() > 1 && caches.iter().any(|c| c.pos != 0) {
5634            return Err(
5635                "step35 batched prime supports continuation only at B=1; a cross-request batch \
5636                 at mixed positions requires per-request queued_after"
5637                    .into(),
5638            );
5639        }
5640
5641        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5642        for &t in &ts {
5643            assert!(
5644                t >= PRIME_MIN_T,
5645                "step35 batched prime needs T >= {PRIME_MIN_T}"
5646            );
5647        }
5648        for (s, c) in caches.iter().enumerate() {
5649            // POS-INCLUSIVE, like the walk's assert: a continuation chunk's rows land at
5650            // c.pos.., so the fresh-only `ts[s] <= max_ctx` form under-checked it.
5651            assert!(
5652                c.pos + ts[s] <= c.max_ctx,
5653                "step35 batched prime exceeds cache max_ctx"
5654            );
5655            assert!(
5656                seq_ends[s] >= c.pos + ts[s],
5657                "step35 batched prime: seq_end must cover this chunk"
5658            );
5659        }
5660        let mut transaction = CacheTaintGuard::arm(caches);
5661        // MEMRA_STEP35_PRIME_BATCH_TSEND=1: CANARY SEAM restoring the pre-fix chunk-local
5662        // `seq_end` (this chunk's own length, which `ts[s]` used to supply here). It is suffix-
5663        // and chunk-VARIANT by construction, so the suffix byte-identity gate MUST break under
5664        // it. That is how the defect is DEMONSTRATED rather than argued: one binary, one seam,
5665        // the legacy arm fails cold-vs-rewound identity and the default arm passes. Read per
5666        // call; never on in a measured default run.
5667        let legacy_tsend = std::env::var("MEMRA_STEP35_PRIME_BATCH_TSEND").as_deref() == Ok("1");
5668        let seq_ends_eff: Vec<usize> = if legacy_tsend {
5669            ts.clone()
5670        } else {
5671            seq_ends.to_vec()
5672        };
5673        let offs: Vec<usize> = ts
5674            .iter()
5675            .scan(0usize, |a, &t| {
5676                let o = *a;
5677                *a += t;
5678                Some(o)
5679            })
5680            .collect();
5681        let total: usize = ts.iter().sum();
5682        let payload = total * self.cfg.n_embd as usize;
5683        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5684        // Positions start at each sequence's CURRENT cache position, not 0, so this entry can
5685        // prime a continuation chunk. The attention core already supports it: step35_attn_pre_wo
5686        // with Some(cache) is PRIME mode — it appends this chunk's post-rope K / raw V and
5687        // attends THROUGH the cache view — so only the hardcoded 0..t and the guard below ever
5688        // restricted it to fresh prompts.
5689        let positions: Vec<Vec<i32>> = ts
5690            .iter()
5691            .zip(caches.iter())
5692            .map(|(&t, c)| {
5693                let base = c.pos as i32;
5694                (0..t as i32).map(|i| base + i).collect()
5695            })
5696            .collect();
5697        let upload_positions =
5698            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
5699                positions
5700                    .iter()
5701                    .map(|p| e.htod_i32(p))
5702                    .collect::<Result<_, _>>()
5703            };
5704
5705        static ONCE: std::sync::Once = std::sync::Once::new();
5706        ONCE.call_once(|| {
5707            eprintln!(
5708                "[step35-prime-batch] first concat prime: B={} tokens={total}",
5709                prompts.len()
5710            );
5711        });
5712
5713        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
5714            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5715                let rt = crate::pp::PpNRt::get(e)?;
5716                let n_st = fence.len() - 1;
5717                assert_eq!(
5718                    rt.n_stages(),
5719                    n_st,
5720                    "step35 prime batch stage count mismatch"
5721                );
5722                let caller_stream = e.stream();
5723                rt.fence_stages_behind(&caller_stream)?;
5724
5725                let mut slot = {
5726                    let _st0 = rt.enter(0);
5727                    let e0 = rt.engine(0, e);
5728                    let pos_ds = upload_positions(e0)?;
5729                    let x = self.embed(e0, &cat_tokens)?;
5730                    let x = self.step35_prime_batch_layers(
5731                        e0,
5732                        x,
5733                        fence[0],
5734                        fence[1],
5735                        &ts,
5736                        &offs,
5737                        &seq_ends_eff,
5738                        &pos_ds,
5739                        caches,
5740                    )?;
5741                    rt.tx(0, &x, payload)?
5742                };
5743                for s in 1..n_st - 1 {
5744                    let _st = rt.enter(s);
5745                    let es = rt.engine(s, e);
5746                    let pos_ds = upload_positions(es)?;
5747                    let x = rt.rx(s - 1, slot, payload)?;
5748                    let x = self.step35_prime_batch_layers(
5749                        es,
5750                        x,
5751                        fence[s],
5752                        fence[s + 1],
5753                        &ts,
5754                        &offs,
5755                        &seq_ends_eff,
5756                        &pos_ds,
5757                        caches,
5758                    )?;
5759                    slot = rt.tx(s, &x, payload)?;
5760                }
5761
5762                let _stl = rt.enter(n_st - 1);
5763                let el = rt.engine(n_st - 1, e);
5764                let pos_ds = upload_positions(el)?;
5765                let x = rt.rx(n_st - 2, slot, payload)?;
5766                let x = self.step35_prime_batch_layers(
5767                    el,
5768                    x,
5769                    fence[n_st - 1],
5770                    fence[n_st],
5771                    &ts,
5772                    &offs,
5773                    &seq_ends_eff,
5774                    &pos_ds,
5775                    caches,
5776                )?;
5777                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
5778                rt.publish_to(n_st - 1, &caller_stream)?;
5779                crate::pp::STEP35_PRIME_BATCH_SPLITS
5780                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5781                out
5782            } else {
5783                let pos_ds = upload_positions(e)?;
5784                let x = self.embed(e, &cat_tokens)?;
5785                let x = self.step35_prime_batch_layers(
5786                    e,
5787                    x,
5788                    0,
5789                    self.layers.len(),
5790                    &ts,
5791                    &offs,
5792                    &seq_ends_eff,
5793                    &pos_ds,
5794                    caches,
5795                )?;
5796                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5797            }
5798        } else {
5799            let pos_ds = upload_positions(e)?;
5800            let x = self.embed(e, &cat_tokens)?;
5801            let x = self.step35_prime_batch_layers(
5802                e,
5803                x,
5804                0,
5805                self.layers.len(),
5806                &ts,
5807                &offs,
5808                &seq_ends_eff,
5809                &pos_ds,
5810                caches,
5811            )?;
5812            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5813        };
5814        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5815        transaction.commit();
5816        Ok(out)
5817    }
5818
5819    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
5820    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
5821    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
5822    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
5823    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
5824    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
5825    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
5826    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
5827    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
5828    /// over the quantized past; Linear: the stateful pad_view twin — the same state
5829    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
5830    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
5831    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
5832    /// back to single-chunk serving).
5833    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
5834    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
5835    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5836    pub fn prime_cache_batch(
5837        &self,
5838        e: &Engine,
5839        prompts: &[&[u32]],
5840        caches: &mut [&mut Cache],
5841    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5842        self.refuse_hyper("prime_cache_batch")?;
5843        for cache in caches.iter() {
5844            cache.ensure_usable("prime_cache_batch")?;
5845        }
5846        if crate::pp::pp_cuts(self.layers.len()).is_some()
5847            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
5848        {
5849            return Err("pipeline rewrite is not qualified for batched prime".into());
5850        }
5851        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::CarriedPrime) {
5852            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
5853                return Err("neither batched-prime nor eager rewrite is qualified".into());
5854            }
5855            if prompts.len() != caches.len() {
5856                return Err("prime fallback prompt/cache shape mismatch".into());
5857            }
5858            static ONCE: std::sync::Once = std::sync::Once::new();
5859            ONCE.call_once(|| {
5860                eprintln!(
5861                    "[rewrite] carried-prime.v1 unqualified; using individual native eager primes"
5862                );
5863            });
5864            let mut transaction = CacheTaintGuard::arm(caches);
5865            let result: Result<Vec<_>, Box<dyn std::error::Error>> = prompts
5866                .iter()
5867                .copied()
5868                .zip(caches.iter_mut())
5869                .map(|(prompt, cache)| self.prime_cache(e, prompt, cache, 0))
5870                .collect();
5871            if result.is_ok() {
5872                transaction.commit();
5873            }
5874            return result;
5875        }
5876        let _pp_walk =
5877            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5878                let rt = crate::pp::PpNRt::get(e)?;
5879                Some(rt.acquire_walk("prime_cache_batch")?)
5880            } else {
5881                None
5882            };
5883        let cfg = &self.cfg;
5884        let n_embd = cfg.n_embd as usize;
5885        let eps = cfg.rms_eps;
5886        let b = prompts.len();
5887        assert!(b >= 1 && b == caches.len());
5888        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
5889        let carried = pos0s.iter().any(|&p| p > 0);
5890        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
5891        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
5892        // generic concat attn core below (uniform geometry, no per-layer swa window, no
5893        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
5894        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
5895        if self.uses_gemma_program() {
5896            return Err(
5897                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
5898                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
5899                    .into(),
5900            );
5901        }
5902        // Step35 has a dedicated concat walk: the generic core below cannot express its
5903        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
5904        if self.uses_sliding_gated_moe_program() {
5905            // The cross-request driver hands whole requests (no chunk loop of its own), so each
5906            // sequence's request-absolute end IS its base plus its prompt length — the value
5907            // `ts[s]` happened to equal for the fresh B>=1 batches this caller admits, which is
5908            // why this arm is bit-for-bit unchanged by the seq_end threading.
5909            let seq_ends: Vec<usize> = caches
5910                .iter()
5911                .zip(prompts.iter())
5912                .map(|(c, p)| c.pos + p.len())
5913                .collect();
5914            return self.step35_prime_cache_batch(e, prompts, caches, &seq_ends);
5915        }
5916        if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5917            let rt = crate::pp::PpNRt::get(e)?;
5918            if rt.cross_device() {
5919                return Err(
5920                    "prime_cache_batch: generic dense concat prime has no cross-device PP split; use individual prime_cache calls"
5921                        .into(),
5922                );
5923            }
5924        }
5925        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5926        for &t in &ts {
5927            assert!(
5928                t >= PRIME_MIN_T,
5929                "prime_cache_batch needs T >= {PRIME_MIN_T}"
5930            );
5931        }
5932        for (s, c) in caches.iter().enumerate() {
5933            assert!(
5934                c.pos + ts[s] <= c.max_ctx,
5935                "prime_cache_batch: prompt exceeds cache max_ctx"
5936            );
5937        }
5938        let mut transaction = CacheTaintGuard::arm(caches);
5939        let total: usize = ts.iter().sum();
5940        let offs: Vec<usize> = ts
5941            .iter()
5942            .scan(0usize, |a, &t| {
5943                let o = *a;
5944                *a += t;
5945                Some(o)
5946            })
5947            .collect();
5948        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
5949        let pos_ds: Vec<CudaSlice<i32>> = ts
5950            .iter()
5951            .zip(&pos0s)
5952            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
5953            .collect::<Result<_, _>>()?;
5954        // split a concat [total, dim] buffer into per-seq copies
5955        let split = |e: &Engine,
5956                     y: &CudaSlice<f32>,
5957                     dim: usize|
5958         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5959            let mut out = Vec::with_capacity(b);
5960            for s in 0..b {
5961                let mut ys = e.uninit(ts[s] * dim)?;
5962                e.copy_view_into(
5963                    &mut ys,
5964                    0,
5965                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5966                    ts[s] * dim,
5967                )?;
5968                out.push(ys);
5969            }
5970            Ok(out)
5971        };
5972
5973        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5974        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
5975        for (il, layer) in self.layers.iter().enumerate() {
5976            let mut h = e.uninit(total * n_embd)?;
5977            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5978            e.rms_norm_f16out(
5979                &x,
5980                layer.attn_norm.float_data(),
5981                &mut h,
5982                &mut hx16,
5983                n_embd,
5984                total,
5985                eps,
5986            )?;
5987            // mixer: projection GROUP on the concat (m = total), stateful core per seq
5988            let mut mixed = e.uninit(total * n_embd)?;
5989            match &layer.mixer {
5990                Mixer::Full(fa) => {
5991                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
5992                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
5993                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
5994                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
5995                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
5996                    // back to the per-seq dispatch.
5997                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
5998                    let (n_head, n_head_kv, head_dim) = (
5999                        geometry.n_head as usize,
6000                        geometry.n_head_kv as usize,
6001                        geometry.head_dim_k as usize,
6002                    );
6003                    let fa_scale = geometry.attention_scale();
6004                    let use_favl = !carried
6005                        && (2..=8).contains(&b)
6006                        && (head_dim == 256 || head_dim == 128)
6007                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
6008                        && std::env::var("MEMRA_NOFA").is_err()
6009                        && std::env::var("MEMRA_FA_FLOOR").is_err()
6010                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
6011                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
6012                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
6013                    if use_favl {
6014                        let (qf_w, kf_w, vf_w) = (
6015                            fa.wq.out_features(),
6016                            fa.wk.out_features(),
6017                            fa.wv.out_features(),
6018                        );
6019                        // Same bounds contract as `Engine::q_gate_split`, applied to the varlen
6020                        // twin's PER-TOKEN stride. `attn_pre_vl8` takes raw device pointers so it
6021                        // cannot check its own extents; `qf_w` is the wq out-features that set
6022                        // them, and `q_gate_split_vl` reads 2*head_dim per head out of it.
6023                        memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
6024                        struct APre {
6025                            q: CudaSlice<f32>,
6026                            gate: Option<CudaSlice<f32>>,
6027                            qn: CudaSlice<f32>,
6028                            kn: CudaSlice<f32>,
6029                        }
6030                        let mut aps = Vec::with_capacity(b);
6031                        for &t in ts.iter().take(b) {
6032                            aps.push(APre {
6033                                q: e.uninit(t * n_head * head_dim)?,
6034                                gate: Some(e.uninit(t * n_head * head_dim)?),
6035                                qn: e.uninit(t * n_head * head_dim)?,
6036                                kn: e.uninit(t * n_head_kv * head_dim)?,
6037                            });
6038                        }
6039                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
6040                            let kvl = caches[0].kv[il].as_ref().unwrap();
6041                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
6042                        };
6043                        let pargs: Vec<crate::AttnPreVl> = (0..b)
6044                            .map(|s| {
6045                                let (o, t) = (offs[s], ts[s]);
6046                                let kvl = caches[s].kv[il].as_ref().unwrap();
6047                                assert!(
6048                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
6049                                    "prime_cache_batch attn vl: fresh + capacity"
6050                                );
6051                                crate::AttnPreVl {
6052                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
6053                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
6054                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
6055                                    q: e.addr_f32(&aps[s].q),
6056                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
6057                                    qn: e.addr_f32(&aps[s].qn),
6058                                    kn: e.addr_f32(&aps[s].kn),
6059                                    kc: e.addr_u8(&kvl.k),
6060                                    vc: e.addr_u8(&kvl.v),
6061                                    t: t as i32,
6062                                    pad: 0,
6063                                }
6064                            })
6065                            .collect();
6066                        e.attn_pre_vl8(
6067                            &pargs,
6068                            fa.q_norm.float_data(),
6069                            fa.k_norm.float_data(),
6070                            head_dim,
6071                            geometry.n_rot as usize,
6072                            n_head,
6073                            n_head_kv,
6074                            self.cfg.rms_eps,
6075                            geometry.rope_base,
6076                            1.0,
6077                            kv_dim_k,
6078                            kv_dim_v,
6079                            ktb,
6080                            vtb,
6081                        )?;
6082                        for s in 0..b {
6083                            let kvl = caches[s].kv[il].as_mut().unwrap();
6084                            kvl.len += ts[s];
6085                            let new_len = kvl.len as i32;
6086                            e.set_i32_one(&mut kvl.len_d, new_len)?;
6087                        }
6088                        let mut attns = Vec::with_capacity(b);
6089                        let mut mirrors = Vec::with_capacity(b);
6090                        for &t in ts.iter().take(b) {
6091                            attns.push(e.uninit(t * n_head * head_dim)?);
6092                            let n = t * n_head_kv * head_dim;
6093                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
6094                        }
6095                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
6096                        // promoted single-seq config is on; else the mma favl.
6097                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
6098                            Ok("0") => false,
6099                            // Same refusal as the single-seq twin (lib.rs fa_prefill): the
6100                            // batched bf16 stage reaches func("f32_to_bf16_bulk"), absent on a
6101                            // portable build.
6102                            Ok("1") => {
6103                                crate::refuse_portable_force(
6104                                    "MEMRA_FA3=1",
6105                                    "the sm_90a fa3/bf16 kernels",
6106                                );
6107                                true
6108                            }
6109                            _ => cfg!(memra_hopper_mma),
6110                        };
6111                        if fa3_on {
6112                            let mut q16s = Vec::with_capacity(b);
6113                            let mut v16s = Vec::with_capacity(b);
6114                            for s in 0..b {
6115                                let t = ts[s];
6116                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
6117                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
6118                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6119                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
6120                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6121                                e.f32_to_bf16_v(
6122                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
6123                                    &mut v16,
6124                                    t * n_head_kv * head_dim,
6125                                )?;
6126                                q16s.push(q16);
6127                                v16s.push((k16, v16));
6128                            }
6129                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
6130                            let mut kp = qp;
6131                            let mut vp = qp;
6132                            let mut op = [core::ptr::null_mut::<f32>(); 8];
6133                            let mut tsv = [0i32; 8];
6134                            for s in 0..b {
6135                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
6136                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
6137                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
6138                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
6139                                tsv[s] = ts[s] as i32;
6140                            }
6141                            let rc = unsafe {
6142                                crate::fa3_vl_raw(
6143                                    qp.as_ptr(),
6144                                    kp.as_ptr(),
6145                                    vp.as_ptr(),
6146                                    op.as_ptr(),
6147                                    tsv.as_ptr(),
6148                                    b as i32,
6149                                    n_head as i32,
6150                                    n_head_kv as i32,
6151                                    head_dim as i32,
6152                                    fa_scale,
6153                                    e.stream().cu_stream() as *mut core::ffi::c_void,
6154                                )
6155                            };
6156                            if rc != 0 {
6157                                return Err(format!("memra_fa3_vl rc={rc}").into());
6158                            }
6159                        } else {
6160                            let fargs: Vec<crate::FaSeqVl> = (0..b)
6161                                .map(|s| crate::FaSeqVl {
6162                                    q: e.addr_f32(&aps[s].qn),
6163                                    k16: e.addr_u8(&mirrors[s].0),
6164                                    v16: e.addr_u8(&mirrors[s].1),
6165                                    o: e.addr_f32(&attns[s]),
6166                                    kf: e.addr_f32(&aps[s].kn),
6167                                    vf: e.addr_f32v(
6168                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
6169                                    ),
6170                                    t: ts[s] as i32,
6171                                    pad: 0,
6172                                })
6173                                .collect();
6174                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
6175                        }
6176                        for (s, attn) in attns.into_iter().enumerate() {
6177                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
6178                                e,
6179                                attn,
6180                                &aps[s].gate,
6181                                ts[s],
6182                                n_head,
6183                                head_dim,
6184                            )?;
6185                            let mut done = false;
6186                            if let Some(xh) = &ag16 {
6187                                done = e.try_f16_gemm_pre_into_off(
6188                                    &fa.wo,
6189                                    xh,
6190                                    ts[s],
6191                                    &mut mixed,
6192                                    offs[s] * n_embd,
6193                                )?;
6194                            }
6195                            if !done {
6196                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6197                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6198                            }
6199                        }
6200                    } else {
6201                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
6202                            (0..b).map(|_| Vec::new()).collect();
6203                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
6204                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
6205                                parts[s].push(ys);
6206                            }
6207                        }
6208                        for (s, g3s) in parts.into_iter().enumerate() {
6209                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
6210                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
6211                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
6212                            )?;
6213                            let mut done = false;
6214                            if let Some(xh) = &ag16 {
6215                                done = e.try_f16_gemm_pre_into_off(
6216                                    &fa.wo,
6217                                    xh,
6218                                    ts[s],
6219                                    &mut mixed,
6220                                    offs[s] * n_embd,
6221                                )?;
6222                            }
6223                            if !done {
6224                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6225                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6226                            }
6227                        }
6228                    }
6229                }
6230                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched cache prime"),
6231                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched prime"),
6232                Mixer::Linear(la) => {
6233                    // task #16: NO split copies (cores read row-offset views of the concat
6234                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
6235                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
6236                    // varlen K5 launch for all sequences.
6237                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6238                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
6239                    let outs =
6240                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
6241                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
6242                        let (o, t) = (offs[s], ts[s]);
6243                        let mut done = false;
6244                        if let Some(xh) = &gn16 {
6245                            done = e.try_f16_gemm_pre_into_off(
6246                                &la.ssm_out,
6247                                xh,
6248                                t,
6249                                &mut mixed,
6250                                o * n_embd,
6251                            )?;
6252                        }
6253                        if !done {
6254                            let m = e.matmul(&la.ssm_out, &gn, t)?;
6255                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
6256                        }
6257                    }
6258                }
6259            }
6260            let mut x1 = e.uninit(total * n_embd)?;
6261            let mut z = e.uninit(total * n_embd)?;
6262            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6263            e.add_rms_norm_f16out(
6264                &x,
6265                &mixed,
6266                layer.post_attn_norm.float_data(),
6267                &mut x1,
6268                &mut z,
6269                &mut zx16,
6270                n_embd,
6271                total,
6272                eps,
6273            )?;
6274            let ffn_out = match &layer.ffn {
6275                crate::hybrid::Ffn::Dense {
6276                    ffn_gate,
6277                    ffn_up,
6278                    ffn_down,
6279                } => {
6280                    let n_ff = ffn_gate.out_features();
6281                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
6282                    let up = g2.pop().unwrap();
6283                    let gate = g2.pop().unwrap();
6284                    let mut act = e.uninit(total * n_ff)?;
6285                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
6286                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
6287                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
6288                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
6289                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
6290                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
6291                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
6292                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
6293                            Some(y) => y,
6294                            None => e.matmul(ffn_down, &act, total)?,
6295                        }
6296                    } else {
6297                        Self::ffn_act_lim(
6298                            e,
6299                            &self.cfg,
6300                            &gate,
6301                            &up,
6302                            1.0,
6303                            1.0,
6304                            d_lim,
6305                            &mut act,
6306                            total * n_ff,
6307                        )?;
6308                        e.matmul(ffn_down, &act, total)?
6309                    }
6310                }
6311                crate::hybrid::Ffn::Moe(m) => {
6312                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
6313                }
6314            };
6315            let mut x2 = e.uninit(total * n_embd)?;
6316            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
6317            x = x2;
6318        }
6319        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
6320        let mut hn = e.uninit(total * n_embd)?;
6321        e.rms_norm(
6322            &x,
6323            self.output_norm.float_data(),
6324            &mut hn,
6325            n_embd,
6326            total,
6327            eps,
6328        )?;
6329        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
6330        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
6331        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
6332        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
6333        // argmax battery arbitrates, same as every other prefill GEMM change.
6334        let mut hcat = e.uninit(b * n_embd)?;
6335        for s in 0..b {
6336            let last0 = (offs[s] + ts[s] - 1) * n_embd;
6337            e.copy_view_into(
6338                &mut hcat,
6339                s * n_embd,
6340                &hn.slice(last0..last0 + n_embd),
6341                n_embd,
6342            )?;
6343        }
6344        let logits_cat = if b >= 2 {
6345            e.try_f16_gemm(&self.output, &hcat, b)?
6346        } else {
6347            None
6348        };
6349        let logits_host: Option<Vec<f32>> = match &logits_cat {
6350            Some(lc) => Some(e.dtoh(lc)?),
6351            None => None,
6352        };
6353        let n_vocab = self.output.out_features();
6354        let mut hidden_all = if crate::spec::spec_hpost() {
6355            split(e, &hn, n_embd)?
6356        } else {
6357            split(e, &x, n_embd)?
6358        };
6359        let mut out = Vec::with_capacity(b);
6360        for s in 0..b {
6361            let last0 = (offs[s] + ts[s] - 1) * n_embd;
6362            let mut h_seed = e.uninit(n_embd)?;
6363            if !crate::spec::spec_hpost() {
6364                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
6365            } else {
6366                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6367            }
6368            let logits = match &logits_host {
6369                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
6370                None => {
6371                    let mut hlast = e.uninit(n_embd)?;
6372                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6373                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
6374                }
6375            };
6376            caches[s].pos += ts[s];
6377            out.push((logits, h_seed, hidden_all.remove(0)));
6378        }
6379        transaction.commit();
6380        Ok(out)
6381    }
6382
6383    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
6384    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
6385    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
6386    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
6387    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
6388    ///
6389    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
6390    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
6391    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
6392    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
6393    #[allow(clippy::too_many_arguments)]
6394    fn full_attn_prime(
6395        &self,
6396        e: &Engine,
6397        fa: &FullAttnLayer,
6398        h: &CudaSlice<f32>,
6399        hx: Option<&CudaSlice<u8>>,
6400        pos_d: &CudaSlice<i32>,
6401        t: usize,
6402        cache: &mut Cache,
6403        il: usize,
6404        seq_end: usize,
6405    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6406        if self.uses_sliding_gated_moe_program() {
6407            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
6408        }
6409        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
6410        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
6411        // this single-seq path composes proj+core identically (byte-for-byte the old body).
6412        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
6413        let g3 = match hx {
6414            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
6415            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
6416        };
6417        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
6418    }
6419
6420    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
6421    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
6422    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
6423    #[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
6424    fn full_attn_prime_core(
6425        &self,
6426        e: &Engine,
6427        fa: &FullAttnLayer,
6428        g3: Vec<CudaSlice<f32>>,
6429        pos_d: &CudaSlice<i32>,
6430        t: usize,
6431        cache: &mut Cache,
6432        il: usize,
6433    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6434        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
6435        if let Some(xh) = &ag16
6436            && let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)?
6437        {
6438            return Ok(y);
6439        }
6440        e.matmul(&fa.wo, &attn_g, t)
6441    }
6442
6443    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6444    #[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
6445    fn full_attn_prime_core_inner(
6446        &self,
6447        e: &Engine,
6448        fa: &FullAttnLayer,
6449        g3: Vec<CudaSlice<f32>>,
6450        pos_d: &CudaSlice<i32>,
6451        t: usize,
6452        cache: &mut Cache,
6453        il: usize,
6454    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6455        let cfg = &self.cfg;
6456        let geometry = cfg.full_attention_geometry_at(il as u32);
6457        let n_head = geometry.n_head as usize;
6458        let n_head_kv = geometry.n_head_kv as usize;
6459        let head_dim = geometry.head_dim_k as usize;
6460        let scale = geometry.attention_scale();
6461        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
6462        let AttnPre { q, k, v, gate } = pre;
6463        let mut attn = e.uninit(t * n_head * head_dim)?;
6464        self.full_attn_prime_fa_dispatch(
6465            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
6466        )?;
6467        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
6468    }
6469
6470    /// task #18 (attn side): projections tail through KV append — everything before the
6471    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
6472    /// present BEFORE this chunk's append (base_len; 0 == fresh).
6473    #[allow(clippy::type_complexity)]
6474    #[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
6475    fn full_attn_prime_pre_fa(
6476        &self,
6477        e: &Engine,
6478        fa: &FullAttnLayer,
6479        mut g3: Vec<CudaSlice<f32>>,
6480        pos_d: &CudaSlice<i32>,
6481        t: usize,
6482        cache: &mut Cache,
6483        il: usize,
6484    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
6485        let cfg = &self.cfg;
6486        let geometry = cfg.full_attention_geometry_at(il as u32);
6487        let n_head = geometry.n_head as usize;
6488        let n_head_kv = geometry.n_head_kv as usize;
6489        let head_dim = geometry.head_dim_k as usize;
6490        let eps = cfg.rms_eps;
6491
6492        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
6493        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
6494        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
6495        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6496        let v = g3.pop().unwrap();
6497        let mut k = g3.pop().unwrap();
6498        let qf = g3.pop().unwrap();
6499        let (mut q, gate) = if gated {
6500            let mut q = e.uninit(t * n_head * head_dim)?;
6501            let mut gate = e.uninit(t * n_head * head_dim)?;
6502            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
6503            (q, Some(gate))
6504        } else {
6505            (qf, None)
6506        };
6507
6508        let mut qn = e.uninit(t * n_head * head_dim)?;
6509        e.rms_norm(
6510            &q,
6511            fa.q_norm.float_data(),
6512            &mut qn,
6513            head_dim,
6514            n_head * t,
6515            eps,
6516        )?;
6517        q = qn;
6518        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6519        e.rms_norm(
6520            &k,
6521            fa.k_norm.float_data(),
6522            &mut kn,
6523            head_dim,
6524            n_head_kv * t,
6525            eps,
6526        )?;
6527        k = kn;
6528        let rope_dims = geometry.n_rot as usize;
6529        e.rope_neox(
6530            &mut q,
6531            pos_d,
6532            head_dim,
6533            rope_dims,
6534            n_head,
6535            t,
6536            geometry.rope_base,
6537            1.0,
6538        )?;
6539        e.rope_neox(
6540            &mut k,
6541            pos_d,
6542            head_dim,
6543            rope_dims,
6544            n_head_kv,
6545            t,
6546            geometry.rope_base,
6547            1.0,
6548        )?;
6549
6550        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
6551        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
6552        {
6553            let kvl = cache.kv[il].as_mut().unwrap();
6554            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
6555            e.append_kv_quantized_rows(
6556                &k,
6557                &v,
6558                &mut kvl.k,
6559                &mut kvl.v,
6560                kvl.len,
6561                t,
6562                kvl.kv_dim_k,
6563                kvl.kv_dim_v,
6564                kvl.k_tok_bytes,
6565                kvl.v_tok_bytes,
6566                crate::Engine::kv_fp8_on(),
6567            )?;
6568            kvl.len += t;
6569            let new_len = kvl.len as i32;
6570            e.set_i32_one(&mut kvl.len_d, new_len)?;
6571        }
6572
6573        let base_len = {
6574            let kvl = cache.kv[il].as_ref().unwrap();
6575            kvl.len - t // KV rows present BEFORE this chunk's append above
6576        };
6577        Ok((AttnPre { q, k, v, gate }, base_len))
6578    }
6579
6580    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
6581    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
6582    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
6583    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
6584    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
6585    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
6586    #[allow(clippy::too_many_arguments)]
6587    fn full_attn_prime_fa_dispatch(
6588        &self,
6589        e: &Engine,
6590        q: &CudaSlice<f32>,
6591        k: &CudaSlice<f32>,
6592        v: &CudaSlice<f32>,
6593        attn: &mut CudaSlice<f32>,
6594        base_len: usize,
6595        t: usize,
6596        cache: &mut Cache,
6597        il: usize,
6598        head_dim: usize,
6599        n_head: usize,
6600        n_head_kv: usize,
6601        scale: f32,
6602    ) -> Result<(), Box<dyn std::error::Error>> {
6603        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
6604        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
6605        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
6606        // attend through the quantized cache exactly like every later chunk (quantize-then-
6607        // attend). One numeric class for every row => the chunk size cannot decide where a
6608        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
6609        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
6610        // pin-the-boundary approach).
6611        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
6612        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
6613        // with the fix unconditional, only re-introducing the class edge can prove the gate
6614        // still detects the mechanism. Never on in a measured default run.
6615        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
6616            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6617                e.sdpa_naive(
6618                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6619                )?;
6620            } else {
6621                e.fa_prefill(
6622                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6623                )?;
6624            }
6625            return Ok(());
6626        }
6627        let kvl = cache.kv[il].as_ref().unwrap();
6628        let t_kv = base_len + t;
6629        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6630        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6631        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
6632        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
6633        // same numeric class, so the uniform contract holds on the fallback too.
6634        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6635            e.sdpa_naive_quantized_view(
6636                q,
6637                &k_view,
6638                &v_view,
6639                attn,
6640                head_dim,
6641                n_head,
6642                n_head_kv,
6643                t,
6644                t_kv,
6645                scale,
6646                true,
6647                kvl.k_tok_bytes,
6648                kvl.v_tok_bytes,
6649            )?;
6650            return Ok(());
6651        }
6652        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
6653        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
6654        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
6655        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
6656        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
6657        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
6658        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
6659        let deqw = std::env::var("MEMRA_PRIME_DEQW")
6660            .map(|v| v != "0")
6661            .unwrap_or(true);
6662        if deqw {
6663            e.fa_prefill_view_ws(
6664                q,
6665                &k_view,
6666                &v_view,
6667                attn,
6668                head_dim,
6669                n_head,
6670                n_head_kv,
6671                t,
6672                t_kv,
6673                scale,
6674                true,
6675                kvl.k_tok_bytes,
6676                kvl.v_tok_bytes,
6677                crate::Engine::kv_fp8_on(),
6678            )?;
6679        } else {
6680            e.fa_prefill_view(
6681                q,
6682                &k_view,
6683                &v_view,
6684                attn,
6685                head_dim,
6686                n_head,
6687                n_head_kv,
6688                t,
6689                t_kv,
6690                scale,
6691                true,
6692                kvl.k_tok_bytes,
6693                kvl.v_tok_bytes,
6694                crate::Engine::kv_fp8_on(),
6695            )?;
6696        }
6697        Ok(())
6698    }
6699
6700    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
6701    /// (bit-identical composition) and hands wo its fp16 operand directly.
6702    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6703    fn full_attn_prime_post_fa(
6704        &self,
6705        e: &Engine,
6706        attn: CudaSlice<f32>,
6707        gate: &Option<CudaSlice<f32>>,
6708        t: usize,
6709        n_head: usize,
6710        head_dim: usize,
6711    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6712        let (attn_g, ag16) = match gate {
6713            Some(gate) => {
6714                let n = t * n_head * head_dim;
6715                let mut ag = e.uninit(n)?;
6716                if Self::f16out_on(e, t) {
6717                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
6718                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
6719                    (ag, Some(a16))
6720                } else {
6721                    let mut gsig = e.uninit(n)?;
6722                    e.sigmoid(gate, &mut gsig, n)?;
6723                    e.mul(&attn, &gsig, &mut ag, n)?;
6724                    (ag, None)
6725                }
6726            }
6727            None => (attn, None),
6728        };
6729        Ok((attn_g, ag16))
6730    }
6731
6732    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
6733    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
6734    /// carried THROUGH the cache like the spec verify does: carried-ring conv
6735    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
6736    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
6737    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
6738    #[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
6739    fn linear_attn_prime(
6740        &self,
6741        e: &Engine,
6742        la: &LinearAttnLayer,
6743        h: &CudaSlice<f32>,
6744        hx: Option<&CudaSlice<u8>>,
6745        t: usize,
6746        cache: &mut Cache,
6747        il: usize,
6748    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6749        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
6750        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6751        let g4 = match hx {
6752            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
6753            None => e.matmul_group(&ws, h, t)?,
6754        };
6755        self.linear_attn_prime_core(e, la, g4, t, cache, il)
6756    }
6757
6758    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
6759    fn linear_attn_prime_core(
6760        &self,
6761        e: &Engine,
6762        la: &LinearAttnLayer,
6763        mut g4: Vec<CudaSlice<f32>>,
6764        t: usize,
6765        cache: &mut Cache,
6766        il: usize,
6767    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6768        self.linear_attn_prime_core_pad(e, la, std::mem::take(&mut g4), t, cache, il, None)
6769    }
6770
6771    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
6772    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
6773    /// conv ring writes back from the true tail. None = classic path, byte-identical.
6774    #[allow(clippy::too_many_arguments)]
6775    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6776    fn linear_attn_prime_core_pad_inner(
6777        &self,
6778        e: &Engine,
6779        la: &LinearAttnLayer,
6780        mut g4: Vec<CudaSlice<f32>>,
6781        t: usize,
6782        cache: &mut Cache,
6783        il: usize,
6784        pad_len: Option<&CudaSlice<i32>>,
6785    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6786        // shim over the view twin (task #16): full-range views of the owned buffers.
6787        let geometry = la.geometry;
6788        let d_state = geometry.key_head_dim as usize;
6789        let num_k = geometry.key_heads as usize;
6790        let num_v = geometry.value_heads as usize;
6791        let key_dim = d_state * num_k;
6792        let value_dim = geometry.value_head_dim as usize * num_v;
6793        let conv_dim = key_dim * 2 + value_dim;
6794        let alpha = g4.pop().unwrap(); // [T, num_v]
6795        let beta_raw = g4.pop().unwrap(); // [T, num_v]
6796        let z = g4.pop().unwrap(); // [T, value_dim]
6797        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
6798        self.linear_attn_prime_core_pad_view(
6799            e,
6800            la,
6801            &qkv_mixed.slice(0..t * conv_dim),
6802            &z.slice(0..t * value_dim),
6803            &beta_raw.slice(0..t * num_v),
6804            &alpha.slice(0..t * num_v),
6805            t,
6806            cache,
6807            il,
6808            pad_len,
6809        )
6810    }
6811
6812    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
6813    /// shared verbatim by the per-seq scan path and the varlen batched path.
6814    #[allow(clippy::too_many_arguments)]
6815    fn linear_attn_gdn_prep(
6816        &self,
6817        e: &Engine,
6818        la: &LinearAttnLayer,
6819        qkv_mixed: &cudarc::driver::CudaView<f32>,
6820        beta_raw: &cudarc::driver::CudaView<f32>,
6821        alpha: &cudarc::driver::CudaView<f32>,
6822        t: usize,
6823        cache: &mut Cache,
6824        il: usize,
6825        pad_len: Option<&CudaSlice<i32>>,
6826    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
6827        let cfg = &self.cfg;
6828        let geometry = la.geometry;
6829        let d_state = geometry.key_head_dim as usize;
6830        let num_k = geometry.key_heads as usize;
6831        let num_v = geometry.value_heads as usize;
6832        let d_conv = geometry.conv_kernel as usize;
6833        let key_dim = d_state * num_k; // 2048
6834        let value_dim = geometry.value_head_dim as usize * num_v;
6835        let conv_dim = key_dim * 2 + value_dim; // 8192
6836        let eps = cfg.rms_eps;
6837        debug_assert!(
6838            t >= d_conv - 1,
6839            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
6840        );
6841
6842        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
6843        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
6844        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
6845        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
6846        let rl = cache.recur[il].as_mut().unwrap();
6847        let hk = Self::gdn_hk(e, t, num_v, num_k);
6848        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
6849        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
6850        let mut q_g = e.uninit(d_state * hk * t)?;
6851        let mut k_g = e.uninit(d_state * hk * t)?;
6852        let mut v_g = e.uninit(d_state * num_v * t)?;
6853        if conv_fuse {
6854            e.ssm_conv1d_gdn_state_pad(
6855                qkv_mixed,
6856                &mut rl.conv_state,
6857                la.ssm_conv1d.float_data(),
6858                &mut q_g,
6859                &mut k_g,
6860                &mut v_g,
6861                conv_dim,
6862                t,
6863                d_conv,
6864                d_state,
6865                num_v,
6866                num_k,
6867                key_dim,
6868                hk,
6869                pad_len,
6870            )?;
6871        } else {
6872            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
6873            e.ssm_conv1d_tm_state_pad_v(
6874                qkv_mixed,
6875                &mut rl.conv_state,
6876                la.ssm_conv1d.float_data(),
6877                &mut conv_out,
6878                conv_dim,
6879                t,
6880                d_conv,
6881                pad_len,
6882            )?;
6883            e.qkv_to_gdn_repack(
6884                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
6885            )?;
6886        }
6887        let mut q_l2 = e.uninit(d_state * hk * t)?;
6888        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
6889        // Emitted only where a consumer exists (the wgmma config) — on other arches the
6890        // alloc + epilogue stores would be pure waste.
6891        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
6892            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6893            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
6894            Some(qb)
6895        } else {
6896            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
6897            None
6898        };
6899        let mut k_l2 = e.uninit(d_state * hk * t)?;
6900        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
6901        let kb16 = if Engine::l2_v2_on(d_state) {
6902            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6903            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
6904            Some(kb)
6905        } else {
6906            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
6907            None
6908        };
6909        let mut beta = e.uninit(t * num_v)?;
6910        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
6911        let mut g_log = e.uninit(t * num_v)?;
6912        e.gdn_glog_v(
6913            alpha,
6914            la.ssm_dt.float_data(),
6915            la.ssm_a.float_data(),
6916            &mut g_log,
6917            num_v,
6918            t,
6919        )?;
6920        if let Some(len_d) = pad_len {
6921            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
6922        }
6923        Ok(GdnPrep {
6924            hk,
6925            q_l2,
6926            k_l2,
6927            v_g,
6928            beta,
6929            g_log,
6930            kb16,
6931            qb16,
6932        })
6933    }
6934
6935    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
6936    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
6937    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
6938    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
6939    #[allow(clippy::too_many_arguments)]
6940    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6941    fn linear_attn_prime_core_batch(
6942        &self,
6943        e: &Engine,
6944        la: &LinearAttnLayer,
6945        g4: &[CudaSlice<f32>],
6946        offs: &[usize],
6947        ts: &[usize],
6948        caches: &mut [&mut Cache],
6949        il: usize,
6950    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
6951        let geometry = la.geometry;
6952        let d_state = geometry.key_head_dim as usize;
6953        let num_k = geometry.key_heads as usize;
6954        let num_v = geometry.value_heads as usize;
6955        let d_conv = geometry.conv_kernel as usize;
6956        let key_dim = d_state * num_k;
6957        let value_dim = geometry.value_head_dim as usize * num_v;
6958        let conv_dim = key_dim * 2 + value_dim;
6959        let eps = self.cfg.rms_eps;
6960        let scale = 1.0 / (d_state as f32).sqrt();
6961        let b = ts.len();
6962        let c = Engine::gdn_chunk_size();
6963        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
6964        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
6965        let carried = caches.iter().any(|c| c.pos > 0);
6966        let use_vl = !carried
6967            && (2..=8).contains(&b)
6968            && Engine::gdn_chunked_enabled()
6969            && ts.iter().all(|&t| t >= 16)
6970            && e.gdn_mma_enabled(c)
6971            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
6972        if !use_vl {
6973            return (0..b)
6974                .map(|s| {
6975                    let (o, t) = (offs[s], ts[s]);
6976                    self.linear_attn_prime_core_pad_view(
6977                        e,
6978                        la,
6979                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
6980                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
6981                        &g4[2].slice(o * num_v..(o + t) * num_v),
6982                        &g4[3].slice(o * num_v..(o + t) * num_v),
6983                        t,
6984                        caches[s],
6985                        il,
6986                        None,
6987                    )
6988                })
6989                .collect();
6990        }
6991        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
6992        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
6993        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
6994        struct SeqBufs {
6995            conv_out: CudaSlice<f32>,
6996            q_g: CudaSlice<f32>,
6997            k_g: CudaSlice<f32>,
6998            v_g: CudaSlice<f32>,
6999            q_l2: CudaSlice<f32>,
7000            k_l2: CudaSlice<f32>,
7001            beta: CudaSlice<f32>,
7002            g_log: CudaSlice<f32>,
7003            gn: CudaSlice<f32>,
7004            gn16: CudaSlice<u8>,
7005        }
7006        let f16o = Self::f16out_on(e, 16);
7007        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
7008        let mut sb = Vec::with_capacity(b);
7009        let mut pres = Vec::with_capacity(b);
7010        for &t in ts.iter().take(b) {
7011            sb.push(SeqBufs {
7012                conv_out: e.uninit(conv_dim * t)?,
7013                q_g: e.uninit(d_state * hk * t)?,
7014                k_g: e.uninit(d_state * hk * t)?,
7015                v_g: e.uninit(d_state * num_v * t)?,
7016                q_l2: e.uninit(d_state * hk * t)?,
7017                k_l2: e.uninit(d_state * hk * t)?,
7018                beta: e.uninit(t * num_v)?,
7019                g_log: e.uninit(t * num_v)?,
7020                gn: e.uninit(d_state * num_v * t)?,
7021                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
7022            });
7023            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
7024        }
7025        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
7026            .map(|s| {
7027                let (o, t) = (offs[s], ts[s]);
7028                let rl = caches[s].recur[il].as_ref().unwrap();
7029                crate::GdnPrepVl {
7030                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
7031                    conv_state: e.addr_f32(&rl.conv_state),
7032                    conv_out: e.addr_f32(&sb[s].conv_out),
7033                    q_g: e.addr_f32(&sb[s].q_g),
7034                    k_g: e.addr_f32(&sb[s].k_g),
7035                    v_g: e.addr_f32(&sb[s].v_g),
7036                    q_l2: e.addr_f32(&sb[s].q_l2),
7037                    k_l2: e.addr_f32(&sb[s].k_l2),
7038                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
7039                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
7040                    beta: e.addr_f32(&sb[s].beta),
7041                    g_log: e.addr_f32(&sb[s].g_log),
7042                    o: e.addr_f32(&pres[s].o),
7043                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
7044                    gn: e.addr_f32(&sb[s].gn),
7045                    gn16: e.addr_u8(&sb[s].gn16),
7046                    kb16: if Engine::l2_v2_on(d_state) {
7047                        e.addr_u8(&pres[s].kb16)
7048                    } else {
7049                        0
7050                    },
7051                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
7052                        e.addr_u8(&pres[s].qb16)
7053                    } else {
7054                        0
7055                    },
7056                    t: t as i32,
7057                    pad: 0,
7058                }
7059            })
7060            .collect();
7061        let args: Vec<crate::GdnSeqVl> = (0..b)
7062            .map(|s| {
7063                let rl = caches[s].recur[il].as_ref().unwrap();
7064                crate::GdnSeqVl {
7065                    kb16: e.addr_u8(&pres[s].kb16),
7066                    gcum: e.addr_f32(&pres[s].gcum),
7067                    beta: e.addr_f32(&sb[s].beta),
7068                    u: e.addr_f32(&pres[s].u),
7069                    wb16: e.addr_u8(&pres[s].wb16),
7070                    y: e.addr_u8(&pres[s].y16),
7071                    ssnap: e.addr_u8(&pres[s].ssnap16),
7072                    state_in: e.addr_f32(&rl.ssm_state),
7073                    state_out: e.addr_f32(&rl.ssm_state_alt),
7074                    q: e.addr_f32(&sb[s].q_l2),
7075                    p: e.addr_f32(&pres[s].p),
7076                    o: e.addr_f32(&pres[s].o),
7077                    k: e.addr_f32(&sb[s].k_l2),
7078                    v: e.addr_f32(&sb[s].v_g),
7079                    g: e.addr_f32(&sb[s].g_log),
7080                    a: e.addr_f32(&pres[s].a),
7081                    w: e.addr_f32(&pres[s].w),
7082                    t: ts[s] as i32,
7083                    nc: pres[s].nc as i32,
7084                }
7085            })
7086            .collect();
7087        e.gdn_prep_vl8(
7088            &prep_args,
7089            la.ssm_conv1d.float_data(),
7090            la.ssm_dt.float_data(),
7091            la.ssm_a.float_data(),
7092            conv_dim,
7093            d_conv,
7094            d_state,
7095            num_v,
7096            num_k,
7097            key_dim,
7098            hk,
7099            eps,
7100        )?;
7101        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
7102        // both standalone mirror launches vanish on the default config.
7103        if !Engine::l2_v2_on(d_state) {
7104            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
7105        }
7106        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
7107        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
7108            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
7109            if !Engine::l2_v2_on(d_state) {
7110                for s in 0..b {
7111                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
7112                }
7113            }
7114            let mut wa = [crate::GdnWVl::default(); 8];
7115            for s in 0..b {
7116                wa[s] = crate::GdnWVl {
7117                    qb16: e.addr_u8(&pres[s].qb16),
7118                    pb16: e.addr_u8(&pres[s].pb16),
7119                };
7120            }
7121            Some(crate::GdnWVl8(wa))
7122        } else {
7123            None
7124        };
7125        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
7126        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
7127        if f16o {
7128            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
7129        }
7130        // per-seq state swap (+ non-f16out tail fallback)
7131        let mut out = Vec::with_capacity(b);
7132        for (s, bufs) in sb.into_iter().enumerate() {
7133            let rl = caches[s].recur[il].as_mut().unwrap();
7134            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7135            let (o, t) = (offs[s], ts[s]);
7136            let SeqBufs { mut gn, gn16, .. } = bufs;
7137            if f16o {
7138                out.push((gn, Some(gn16)));
7139            } else {
7140                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
7141                e.gated_rmsnorm_zv(
7142                    &pres[s].o,
7143                    la.ssm_norm.float_data(),
7144                    &z_v,
7145                    &mut gn,
7146                    d_state,
7147                    num_v * t,
7148                    eps,
7149                )?;
7150                out.push((gn, None));
7151            }
7152        }
7153        Ok(out)
7154    }
7155
7156    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
7157    /// views of the CONCAT projection outputs directly (no per-seq split copies).
7158    /// Same kernels, same values, byte-identical to the Vec shim above.
7159    #[allow(clippy::too_many_arguments)]
7160    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7161    fn linear_attn_prime_core_pad_view(
7162        &self,
7163        e: &Engine,
7164        la: &LinearAttnLayer,
7165        qkv_mixed: &cudarc::driver::CudaView<f32>,
7166        z: &cudarc::driver::CudaView<f32>,
7167        beta_raw: &cudarc::driver::CudaView<f32>,
7168        alpha: &cudarc::driver::CudaView<f32>,
7169        t: usize,
7170        cache: &mut Cache,
7171        il: usize,
7172        pad_len: Option<&CudaSlice<i32>>,
7173    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7174        let cfg = &self.cfg;
7175        let geometry = la.geometry;
7176        let d_state = geometry.key_head_dim as usize;
7177        let num_v = geometry.value_heads as usize;
7178        let eps = cfg.rms_eps;
7179        let scale = 1.0 / (d_state as f32).sqrt();
7180
7181        let prep =
7182            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
7183
7184        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
7185        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
7186        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
7187        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
7188        // verify keep the sequential kernel).
7189        let mut o = e.uninit(d_state * num_v * t)?;
7190        let rl = cache.recur[il].as_mut().unwrap();
7191        {
7192            let crate::cache::RecurLayer {
7193                ssm_state,
7194                ssm_state_alt,
7195                ..
7196            } = rl;
7197            e.gdn_scan_prefill(
7198                &prep.q_l2,
7199                &prep.k_l2,
7200                &prep.v_g,
7201                &prep.g_log,
7202                &prep.beta,
7203                prep.kb16.as_ref(),
7204                prep.qb16.as_ref(),
7205                ssm_state,
7206                ssm_state_alt,
7207                &mut o,
7208                num_v,
7209                t,
7210                scale,
7211                prep.hk,
7212            )?;
7213        }
7214        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7215
7216        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
7217        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
7218        let mut gn = e.uninit(d_state * num_v * t)?;
7219        let gn16 = if Self::f16out_on(e, t) {
7220            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
7221            e.gated_rmsnorm_f16out_zv(
7222                &o,
7223                la.ssm_norm.float_data(),
7224                z,
7225                &mut gn,
7226                &mut g16,
7227                d_state,
7228                num_v * t,
7229                eps,
7230            )?;
7231            Some(g16)
7232        } else {
7233            e.gated_rmsnorm_zv(
7234                &o,
7235                la.ssm_norm.float_data(),
7236                z,
7237                &mut gn,
7238                d_state,
7239                num_v * t,
7240                eps,
7241            )?;
7242            None
7243        };
7244        Ok((gn, gn16))
7245    }
7246
7247    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
7248    #[allow(clippy::too_many_arguments)]
7249    fn linear_attn_prime_core_pad(
7250        &self,
7251        e: &Engine,
7252        la: &LinearAttnLayer,
7253        g4: Vec<CudaSlice<f32>>,
7254        t: usize,
7255        cache: &mut Cache,
7256        il: usize,
7257        pad_len: Option<&CudaSlice<i32>>,
7258    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7259        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
7260        if let Some(xh) = &gn16
7261            && let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)?
7262        {
7263            return Ok(y);
7264        }
7265        e.matmul(&la.ssm_out, &gn, t)
7266    }
7267
7268    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
7269    ///
7270    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
7271    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
7272    pub fn full_attn(
7273        &self,
7274        e: &Engine,
7275        fa: &FullAttnLayer,
7276        h: &CudaSlice<f32>,
7277        pos_d: &CudaSlice<i32>,
7278        t: usize,
7279        il: usize,
7280    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7281        if self.uses_sliding_gated_moe_program() {
7282            return self.step35_attn(e, fa, h, pos_d, t, il);
7283        }
7284        let cfg = &self.cfg;
7285        let _n_embd = cfg.n_embd as usize;
7286        let geometry = cfg.full_attention_geometry_at(il as u32);
7287        let n_head = geometry.n_head as usize;
7288        let n_head_kv = geometry.n_head_kv as usize;
7289        let head_dim = geometry.head_dim_k as usize;
7290        let eps = cfg.rms_eps;
7291        let scale = geometry.attention_scale();
7292
7293        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
7294        // gate — wq out = n_head*head_dim, no split (see prime-path note).
7295        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7296        // A load-time full-attention TP plan owns the same Q/K/V projections for every
7297        // architecture. Fall back to the original grouped owner-device projection when this
7298        // layer has no TP sidecar.
7299        let mut g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
7300            Some(g3) => g3,
7301            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7302        };
7303        let v = g3.pop().unwrap();
7304        let mut k = g3.pop().unwrap();
7305        let qf = g3.pop().unwrap();
7306        let (mut q, gate) = if gated {
7307            let mut q = e.uninit(t * n_head * head_dim)?;
7308            let mut gate = e.uninit(t * n_head * head_dim)?;
7309            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7310            (q, Some(gate))
7311        } else {
7312            (qf, None)
7313        };
7314
7315        // QK-norm (per head_dim row), then partial RoPE.
7316        let mut qn = e.uninit(t * n_head * head_dim)?;
7317        e.rms_norm(
7318            &q,
7319            fa.q_norm.float_data(),
7320            &mut qn,
7321            head_dim,
7322            n_head * t,
7323            eps,
7324        )?;
7325        q = qn;
7326        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7327        e.rms_norm(
7328            &k,
7329            fa.k_norm.float_data(),
7330            &mut kn,
7331            head_dim,
7332            n_head_kv * t,
7333            eps,
7334        )?;
7335        k = kn;
7336        let rope_dims = geometry.n_rot as usize;
7337        e.rope_neox(
7338            &mut q,
7339            pos_d,
7340            head_dim,
7341            rope_dims,
7342            n_head,
7343            t,
7344            geometry.rope_base,
7345            1.0,
7346        )?;
7347        e.rope_neox(
7348            &mut k,
7349            pos_d,
7350            head_dim,
7351            rope_dims,
7352            n_head_kv,
7353            t,
7354            geometry.rope_base,
7355            1.0,
7356        )?;
7357
7358        // SDPA
7359        let mut attn = e.uninit(t * n_head * head_dim)?;
7360        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
7361        // falls back to naive sdpa.
7362        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
7363            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
7364            e.sdpa_naive(
7365                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7366            )?;
7367        } else {
7368            e.fa_prefill(
7369                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7370            )?;
7371        }
7372
7373        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
7374        let attn_g = match &gate {
7375            Some(gate) => {
7376                let mut gsig = e.uninit(t * n_head * head_dim)?;
7377                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
7378                let mut ag = e.uninit(t * n_head * head_dim)?;
7379                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
7380                ag
7381            }
7382            None => attn,
7383        };
7384
7385        // O follows the same generic load-time TP plan as Q/K/V.
7386        self.full_attn_o(e, fa, &attn_g, t)
7387    }
7388
7389    /// The absorb/decompress operands (`attn_k_b` / `attn_v_b`) are 3D and are ALWAYS the Float
7390    /// arm — on every checkpoint dtype, not just an f32 fixture. Two guards upstream make that a
7391    /// property rather than a hope: `GpuTensor::load_from_source` refuses any quantized non-2D
7392    /// tensor by name (`row_bytes` is derived from `ne[1]`, the MIDDLE axis of a 3D tensor), and
7393    /// `MlaAttnLayer::load` audits residency at load. A quantized `kv_b_proj` is dequantized at
7394    /// the source by `TransformKind::MlaKeyUpSplit`/`MlaValueUpSplit`. This is the last backstop:
7395    /// fail NAMING the constraint rather than through `float_data()`'s norm-flavoured panic.
7396    fn mla_split_operand<'w>(
7397        w: &'w crate::model::GpuTensor,
7398        name: &str,
7399        il: usize,
7400    ) -> &'w CudaSlice<f32> {
7401        match w {
7402            crate::model::GpuTensor::Float { data, .. } => data,
7403            _ => panic!(
7404                "layer {il}: MLA conversion-split operand {name} is not f32-resident. The 3D \
7405                 (d_nope|kv_rank, kv_rank|d_v, n_head) splits have no quantized resident layout: \
7406                 a quantized 3D tensor mis-derives row_bytes in the generic 2D Quant arm, so the \
7407                 source must dequantize the fused kv_b_proj (TensorTransform::SplitMlaKv). \
7408                 Reaching this means both the loader rank guard and MlaAttnLayer::load's \
7409                 residency audit were bypassed"
7410            ),
7411        }
7412    }
7413
7414    /// MLA (multi-head latent attention) mixer core, ABSORBED form — the one arm that serves
7415    /// prefill, chunked prefill and decode (see `cu/mla_attn.cu` FORM CHOICE).
7416    ///
7417    /// `latent` is the layer's latent KV plane; this call APPENDS its own `t` rows at row
7418    /// `slot` and then attends rows `0..slot + t`, which is exactly the oracle's convention
7419    /// that the queries are the LAST `t_q` rows of the cache (`crate::mla::MlaInputs`).
7420    /// Returns the post-`wo` block output [t, n_embd].
7421    #[allow(clippy::too_many_arguments)]
7422    // allow: the parameter list mirrors the kernel/FFI/call contract (rows_exact is the
7423    // verify-batch matmul-class selector, lane/glm5-verify-batch); bundling into a struct
7424    // is a refactor, not a lint fix
7425    fn mla_attn_core(
7426        &self,
7427        e: &Engine,
7428        mla: &crate::hybrid::MlaAttnLayer,
7429        h: &CudaSlice<f32>,
7430        pos_d: &CudaSlice<i32>,
7431        t: usize,
7432        il: usize,
7433        latent: &mut CudaSlice<f32>,
7434        index_plane: Option<IndexerPlanes<'_>>,
7435        slot: usize,
7436        rows_exact: bool,
7437    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7438        let attn = self.mla_attn_core_pre_wo(
7439            e,
7440            mla,
7441            h,
7442            pos_d,
7443            t,
7444            il,
7445            latent,
7446            index_plane,
7447            slot,
7448            rows_exact,
7449        )?;
7450        // Verify-batch wo seam (lane/glm5-verify-batch): the rows arm routes the output
7451        // projection decode-exact, same as every projection inside the core — the wo
7452        // dispatch moved here with the TP split, its routing did not change.
7453        if rows_exact {
7454            e.matmul_rows_exact(&mla.wo, &attn, t)
7455        } else {
7456            e.matmul(&mla.wo, &attn, t)
7457        }
7458    }
7459
7460    /// [`mla_attn_core`] up to (and excluding) the output projection: returns the
7461    /// per-head attention output `[t, n_head * d_v]`. Split out for the glm5 TP-2 seam,
7462    /// whose column-parallel `wo` runs over the cross-rank GATHERED heads — the plain path
7463    /// is the wrapper above, byte-for-byte the pre-split body (the wo matmul moved,
7464    /// nothing else).
7465    #[allow(clippy::too_many_arguments)]
7466    fn mla_attn_core_pre_wo(
7467        &self,
7468        e: &Engine,
7469        mla: &crate::hybrid::MlaAttnLayer,
7470        h: &CudaSlice<f32>,
7471        pos_d: &CudaSlice<i32>,
7472        t: usize,
7473        il: usize,
7474        latent: &mut CudaSlice<f32>,
7475        index_plane: Option<IndexerPlanes<'_>>,
7476        slot: usize,
7477        rows_exact: bool,
7478    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7479        let g = mla.geom;
7480        let cfg = &self.cfg;
7481        let eps = cfg.rms_eps;
7482        let base = cfg.rope_freq_base;
7483        let (nh, dn, dr, dv, r) = (g.n_head, g.d_nope, g.d_rope, g.d_v, g.kv_rank);
7484        assert_eq!(
7485            g.latent_dim,
7486            r + dr,
7487            "layer {il}: MlaGeom latent_dim disagrees with kv_rank + d_rope"
7488        );
7489        let t_kv = slot + t;
7490        // Verify-batch matmul seam (lane/glm5-verify-batch): rows_exact routes every
7491        // projection through the decode-exact classes so each of the t rows is
7492        // bit-identical to the t=1 decode program (matmul_rows_exact contract); false =
7493        // the unchanged dispatch for every other caller (prime keeps its classes).
7494        let mm = |w: &crate::model::GpuTensor,
7495                  x: &CudaSlice<f32>|
7496         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7497            if rows_exact {
7498                e.matmul_rows_exact(w, x, t)
7499            } else {
7500                e.matmul(w, x, t)
7501            }
7502        };
7503
7504        // --- q path: wq_a -> q_a_norm -> wq_b -> per-head [nope | rope] ---
7505        let q_a = mm(&mla.wq_a, h)?;
7506        let q_lora = mla.wq_b.in_features();
7507        let mut q_an = e.uninit(t * q_lora)?;
7508        e.rms_norm(&q_a, mla.q_a_norm.float_data(), &mut q_an, q_lora, t, eps)?;
7509        let q = mm(&mla.wq_b, &q_an)?;
7510        // Per head the row is [nope | rope] contiguous, so t*nh rows of width dn+dr split with
7511        // the same kernel the latent row uses — the two layouts are the same shape.
7512        let mut q_nope = e.uninit(t * nh * dn)?;
7513        let mut q_pe = e.uninit((t * nh * dr).max(1))?;
7514        e.mla_split_latent(&q, &mut q_nope, &mut q_pe, t * nh, dn, dr)?;
7515        // NoPE (glm5_next, rope_head_dim 0): no rope plane exists. The launcher is a no-op, but
7516        // the allocation above is still non-empty so nothing downstream holds a null slice.
7517        e.mla_rope_interleaved(&mut q_pe, pos_d, t, nh, dr, base)?;
7518
7519        // --- kv path: wkv_a -> [c_kv (rms-normed) | k_pe (roped, NOT normed)] ---
7520        let kv = mm(&mla.wkv_a, h)?;
7521        let mut c_kv = e.uninit(t * r)?;
7522        let mut k_pe = e.uninit((t * dr).max(1))?;
7523        e.mla_split_latent(&kv, &mut c_kv, &mut k_pe, t, r, dr)?;
7524        let mut c_kv_n = e.uninit(t * r)?;
7525        e.rms_norm(&c_kv, mla.kv_a_norm.float_data(), &mut c_kv_n, r, t, eps)?;
7526        e.mla_rope_interleaved(&mut k_pe, pos_d, t, 1, dr, base)?;
7527        e.mla_append_latent(latent, &c_kv_n, &k_pe, slot, t, r, dr)?;
7528
7529        // --- DSA k-pool selection, BEFORE attending: the indexer's own state row for each of
7530        // this call's tokens is appended first, so a query sees itself exactly as the latent
7531        // plane already lets it (the reference concatenates into the indexer cache, then scores).
7532        let gathered = match (&mla.index, index_plane) {
7533            (Some(indexer), Some(plane)) => {
7534                Some(self.mla_kpool_select(e, indexer, h, &q_an, plane, t, slot, il, rows_exact)?)
7535            }
7536            (Some(_), None) => {
7537                return Err(format!(
7538                    "layer {il} declares a DSA k-pool indexer but no indexer state plane was \
7539                     supplied — the ModelPlan must declare StatePlan::LatentKvCache with a \
7540                     non-zero index_width for it"
7541                )
7542                .into());
7543            }
7544            (None, _) => None,
7545        };
7546
7547        // --- absorbed MLA core over the latent plane ---
7548        let wk_b = Self::mla_split_operand(&mla.wk_b, "attn_k_b", il);
7549        let wv_b = Self::mla_split_operand(&mla.wv_b, "attn_v_b", il);
7550
7551        // MEMRA_MLA_TC_PREFILL door (default OFF; flag read per call — the rollback seam).
7552        // Engagement conditions, every one load-bearing:
7553        //   * a gathered selection exists — the door serves the DSA arm only; the dense
7554        //     absorbed arm (GLM-5.2, no indexer) keeps the f32 kernel it was gated on;
7555        //   * d_rope == 0 (NoPE) — the TC kernel treats the latent row as both K and V,
7556        //     which is only the whole truth when there is no rope plane;
7557        //   * kv_rank == 512 — the kernel's stamped head dim (glm5_next / GLM-5.2 class);
7558        //   * t >= 16 — prefill widths only. Decode (t == 1) and short resumes NEVER enter,
7559        //     which is what the decode byte-identity gate proves rather than assumes.
7560        // Anything else falls through to the unchanged f32 kernels below — behavior identical
7561        // to the flag being off.
7562        // A chain returning Ok(None) is a cuBLASLt shape DECLINE (announced once per shape);
7563        // the let-chain then simply does not match and the f32 kernels below serve the call.
7564        // glm5 TP composition guard (lane/glm5-composition): the TC prefill chain's gate
7565        // ran on the FULL-head geometry only; a head shard (any rank) declines it by name
7566        // and falls through to the f32 kernels below — behavior identical to the flag
7567        // being off for that layer, announced once. The composed door re-gates on the box
7568        // (real-artifact kv_rank 512 shapes; the rig fixtures are kv_rank 16 and never
7569        // reach this chain).
7570        // The announce shares the chain's OWN conjuncts (gathered + !portable_mma_gated),
7571        // so it can never blame TP for a decline the missing DSA gather or the MMA gate
7572        // caused (#82 review).
7573        if mla.tp_shard
7574            && gathered.is_some()
7575            && dr == 0
7576            && r == 512
7577            && t >= 16
7578            && !crate::portable_mma_gated()
7579            && mla_tc_prefill_enabled()
7580        {
7581            static TP_TC_DECLINE: std::sync::Once = std::sync::Once::new();
7582            TP_TC_DECLINE.call_once(|| {
7583                eprintln!(
7584                    "[mla-tc-prefill] DECLINED on glm5-TP head shards: the door's gate ran \
7585                     on full-head geometry; shards ride the f32 prefill kernels until the \
7586                     TP composition gate lands (pin MEMRA_MLA_TC_PREFILL=0 to silence)"
7587                );
7588            });
7589        }
7590        if let Some((idx, slots)) = &gathered
7591            && dr == 0
7592            && r == 512
7593            && t >= 16
7594            && !rows_exact // verify-batch stays on the decode-exact classes (t <= 15 anyway)
7595            && !crate::portable_mma_gated()
7596            && !mla.tp_shard
7597            && mla_tc_prefill_enabled()
7598            && let Some(attn) = self.mla_tc_prefill_chain(
7599                e, wk_b, wv_b, &q_nope, latent, idx, *slots, t, t_kv, nh, dn, dv, r, g.scale,
7600            )?
7601        {
7602            return Ok(attn);
7603        }
7604
7605        let mut q_lat = e.uninit(t * nh * r)?;
7606        e.mla_absorb_q(&q_nope, wk_b, &mut q_lat, t, nh, dn, r)?;
7607        let mut o_lat = e.uninit(t * nh * r)?;
7608        match &gathered {
7609            Some((idx, slots)) => e.mla_attn_gathered(
7610                &q_lat, &q_pe, latent, idx, &mut o_lat, nh, r, dr, t, *slots, g.scale,
7611            )?,
7612            None => e.mla_attn_absorbed(
7613                &q_lat, &q_pe, latent, &mut o_lat, nh, r, dr, t, t_kv, g.scale,
7614            )?,
7615        }
7616        let mut attn = e.uninit(t * nh * dv)?;
7617        e.mla_decompress_v(&o_lat, wv_b, &mut attn, t, nh, dv, r)?;
7618
7619        Ok(attn)
7620    }
7621
7622    /// The MEMRA_MLA_TC_PREFILL chain: absorb and decompress as strided-batched bf16
7623    /// tensor-core GEMMs, attention as the gathered bf16 MMA kernel. Returns `Ok(None)` when
7624    /// cuBLASLt declines a GEMM shape (announced once per shape) so the caller falls back to
7625    /// the f32 kernels; every other failure is a hard error.
7626    ///
7627    /// FORM CHOICE, stated for the record (the dual-form MLA law): every fast engine runs
7628    /// MATERIALIZED (per-head MHA) attention at DENSE prefill and absorbed MQA at decode.
7629    /// glm5_next prefill is NOT dense: the DSA indexer caps every query at topk+tail rows and
7630    /// selects ONE list per query SHARED ACROSS ALL 64 HEADS. That shared list is what makes
7631    /// the ABSORBED form the GEMM-shaped one here — the head axis is the MMA m, the shared
7632    /// latent rows are one B operand per tile — while materializing K/V would give every head
7633    /// its own K plane and destroy exactly that sharing (back to per-(query,head) matvecs on
7634    /// the gathered walk). It is also FlashMLA's own sparse-prefill geometry (q 576/512 over
7635    /// gathered latent rows). Queries whose selection is trivial (visible <= topk: the lists
7636    /// ARE the full causal prefix, emitted by the selector itself) ride the SAME kernel with
7637    /// the identity gather — there is no separate dense program to gate.
7638    ///
7639    /// Transient cost per (layer, chunk) at the census shape (t=2313, t_kv=4626): bf16 q_lat
7640    /// 152 MB + bf16 latent window 4.7 MB + bf16 q_nope/o_lat copies ~230 MB — all freed with
7641    /// the call. (The materialized-K/V alternative would have been 4096 x 64 x (256+256) x 2B
7642    /// = 256 MB/layer-chunk of K/V ALONE, plus the per-head-K program cost above.) Weight
7643    /// bf16 converts (wk_b/wv_b, 8.4M elems each) run per call, ~50 us class; a resident
7644    /// mirror is a later diet, not correctness.
7645    #[allow(clippy::too_many_arguments)]
7646    fn mla_tc_prefill_chain(
7647        &self,
7648        e: &Engine,
7649        wk_b: &CudaSlice<f32>,
7650        wv_b: &CudaSlice<f32>,
7651        q_nope: &CudaSlice<f32>,
7652        latent: &CudaSlice<f32>,
7653        idx: &CudaSlice<i32>,
7654        width: usize,
7655        t: usize,
7656        t_kv: usize,
7657        nh: usize,
7658        dn: usize,
7659        dv: usize,
7660        r: usize,
7661        scale: f32,
7662    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7663        // Once-per-shape decline announce (the bf16_tc_gemm pattern): a door that quietly
7664        // stops engaging reads exactly like a door that never helped.
7665        fn declined(stage: &str, m: usize, n: usize, k: usize, batch: usize) {
7666            type ShapeSet = std::collections::HashSet<(usize, usize, usize, usize)>;
7667            static SAID: std::sync::Mutex<Option<ShapeSet>> = std::sync::Mutex::new(None);
7668            let mut g = SAID.lock().unwrap();
7669            if g.get_or_insert_with(std::collections::HashSet::new)
7670                .insert((m, n, k, batch))
7671            {
7672                eprintln!(
7673                    "[mla-tc-prefill] DECLINED at {stage} m={m} n={n} k={k} batch={batch} \
7674                     (no cuBLASLt heuristic) — this call falls back to the f32 MLA kernels"
7675                );
7676            }
7677        }
7678        // Weights and activations to bf16. The converts require n % 4 == 0; every operand here
7679        // is a multiple of the head dims (dn/dv/r all >= 16 and % 4 == 0 on the shapes the door
7680        // admits), asserted rather than assumed.
7681        for (name, n) in [
7682            ("wk_b", nh * r * dn),
7683            ("wv_b", nh * dv * r),
7684            ("q_nope", t * nh * dn),
7685            ("latent", t_kv * r),
7686        ] {
7687            debug_assert!(
7688                n.is_multiple_of(4),
7689                "mla-tc-prefill: {name} elems {n} % 4 != 0"
7690            );
7691            let _ = (name, n);
7692        }
7693        let wk_bf = e.f32_to_bf16(wk_b, nh * r * dn)?;
7694        let wv_bf = e.f32_to_bf16(wv_b, nh * dv * r)?;
7695        let qn_bf = e.f32_to_bf16(q_nope, t * nh * dn)?;
7696        // absorb: per head h, q_lat[:,h,:] [t, r] = q_nope[:,h,:] [t, dn] @ W_uk[h] [r, dn]^T.
7697        // wk_b is the conversion-split (h, l, p) plane, contiguous in p == the reduction axis:
7698        // per head it IS the [n=r, k=dn] row-major operand. bf16 out feeds the attention kernel.
7699        let mut q_lat_bf = e.alloc_u8_uninit(t * nh * r * 2)?;
7700        if !e.mla_bf16_gemm_sb_bf16out(
7701            &wk_bf,
7702            &qn_bf,
7703            &mut q_lat_bf,
7704            t,
7705            r,
7706            dn,
7707            nh * dn,
7708            dn,
7709            nh * r,
7710            r,
7711            nh,
7712        )? {
7713            declined("absorb", t, r, dn, nh);
7714            return Ok(None);
7715        }
7716        // The latent window rows 0..t_kv (this call's rows were appended above), bf16.
7717        let cache_bf = e.f32_to_bf16(latent, t_kv * r)?;
7718        let mut o_lat = e.uninit(t * nh * r)?;
7719        e.mla_attn_gathered_tc(
7720            &q_lat_bf, &cache_bf, idx, &mut o_lat, nh, r, t, width, scale,
7721        )?;
7722        // decompress: per head h, attn[:,h,:] [t, dv] = o_lat[:,h,:] [t, r] @ W_uv[h] [dv, r]^T.
7723        // wv_b is (h, j, l), contiguous in l == the reduction axis: per head [n=dv, k=r].
7724        let o_bf = e.f32_to_bf16(&o_lat, t * nh * r)?;
7725        let mut attn = e.uninit(t * nh * dv)?;
7726        if !e.mla_bf16_gemm_sb_f32out(
7727            &wv_bf,
7728            &o_bf,
7729            &mut attn,
7730            t,
7731            dv,
7732            r,
7733            nh * r,
7734            r,
7735            nh * dv,
7736            dv,
7737            nh,
7738        )? {
7739            declined("decompress", t, dv, r, nh);
7740            return Ok(None);
7741        }
7742        crate::MLA_TC_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7743        {
7744            static ANNOUNCED: std::sync::Once = std::sync::Once::new();
7745            ANNOUNCED.call_once(|| {
7746                eprintln!(
7747                    "[mla-tc-prefill] engaged: absorb/decompress = strided-batched bf16 TC \
7748                     GEMMs, attention = fa_mla_gathered_bf16 (t={t}, t_kv={t_kv}, nh={nh}, \
7749                     width={width}); dispatches counted in MLA_TC_PREFILL_DISPATCHES"
7750                );
7751            });
7752        }
7753        Ok(Some(attn))
7754    }
7755
7756    /// Layer-scoped wrapper: names the layer in any selection failure.
7757    #[allow(clippy::too_many_arguments)]
7758    fn mla_kpool_select(
7759        &self,
7760        e: &Engine,
7761        indexer: &crate::hybrid::MlaIndexer,
7762        h: &CudaSlice<f32>,
7763        q_resid: &CudaSlice<f32>,
7764        plane: IndexerPlanes<'_>,
7765        t: usize,
7766        slot: usize,
7767        il: usize,
7768        rows_exact: bool,
7769    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7770        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, rows_exact).map_err(
7771            |source| -> Box<dyn std::error::Error> {
7772                format!("layer {il}: DSA k-pool selection failed: {source}").into()
7773            },
7774        )
7775    }
7776
7777    /// DSA k-pool indexer: append this call's packed indexer state, then select the cache rows
7778    /// each query may attend. Returns the per-query position list and its width (`-1` padded).
7779    ///
7780    /// The program is `Glm5NextTextIndexer.forward`
7781    /// (research/glm53-flash-bringup-20260827/modular_glm5_next-ref.py:771), transcribed in
7782    /// `memra_reference::kpool_allowed_tokens`, which is this path's oracle:
7783    ///   1. `k = LayerNorm_affine(wk(x))` — LayerNorm WITH BIAS at eps 1e-5, NOT the model's
7784    ///      RMSNorm at `rms_norm_eps`; `gate = index_kpool_compress_gate(x)`. Both are cached.
7785    ///   2. Every COMPLETE pool of `pool` consecutive cached tokens collapses to one key by a
7786    ///      per-channel softmax over (gate + positional embedding).
7787    ///   3. `score[i][p] = sum_h relu(q[i][h] . pool_key[p] * d^-1/2) * weights_proj(x)[i][h] *
7788    ///      heads^-1/2`, with pools whose last token is invisible to the query masked out.
7789    ///   4. Top `top_k / pool` pools expand back to raw rows; the incomplete tail is appended raw.
7790    ///
7791    /// `q_resid` is `q_a_layernorm(q_a_proj(x))` — the SAME tensor the MLA query up-projection
7792    /// consumes, which is why the indexer is scored here rather than before the core.
7793    #[allow(clippy::too_many_arguments)]
7794    pub fn mla_kpool_indices(
7795        e: &Engine,
7796        indexer: &crate::hybrid::MlaIndexer,
7797        h: &CudaSlice<f32>,
7798        q_resid: &CudaSlice<f32>,
7799        plane: IndexerPlanes<'_>,
7800        t: usize,
7801        slot: usize,
7802    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7803        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, false)
7804    }
7805
7806    /// [`Self::mla_kpool_indices`] with the verify-batch matmul-class selector
7807    /// (lane/glm5-verify-batch): `rows_exact` routes the indexer's four projections
7808    /// through the decode-exact classes so each of the t rows is bit-identical to the
7809    /// t=1 decode program; `false` is the unchanged dispatch.
7810    #[allow(clippy::too_many_arguments)]
7811    pub fn mla_kpool_indices_ex(
7812        e: &Engine,
7813        indexer: &crate::hybrid::MlaIndexer,
7814        h: &CudaSlice<f32>,
7815        q_resid: &CudaSlice<f32>,
7816        plane: IndexerPlanes<'_>,
7817        t: usize,
7818        slot: usize,
7819        rows_exact: bool,
7820    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7821        let mm = |w: &crate::model::GpuTensor,
7822                  x: &CudaSlice<f32>|
7823         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7824            if rows_exact {
7825                e.matmul_rows_exact(w, x, t)
7826            } else {
7827                e.matmul(w, x, t)
7828            }
7829        };
7830        /// `nn.LayerNorm` default epsilon. The indexer's k_norm is a LayerNorm, so it does NOT
7831        /// take the model's `rms_norm_eps` (census: "eps 1e-5, NOT rms_norm_eps"); they coincide
7832        /// numerically on GLM-5.3-Flash and the constant keeps them from being coupled.
7833        const INDEX_NORM_EPS: f32 = 1e-5;
7834
7835        let ig = indexer.geom;
7836        let d = ig.head_dim;
7837        let t_kv = slot + t;
7838        let IndexerPlanes {
7839            state: plane,
7840            pool_keys: pool_key_plane,
7841            ready: pools_ready,
7842            state_ring_rows,
7843            capacity_tokens,
7844        } = plane;
7845
7846        // TAIL RING. The plane's rows are read EXACTLY ONCE, by the pool-key build of the pool
7847        // each row belongs to, so a ring of `ring` rows holds everything still live. The state
7848        // plan does not carry `pool`, so the allocator books physical rows and the EFFECTIVE ring
7849        // is rounded down here: a ring that is not a whole number of pools would split a pool
7850        // across the wrap. ONE POOL is the whole correctness floor (lane/glm53-ring-sizing): the
7851        // drain below serves any `t` from any ring at or above it, so `ring` never bounds a
7852        // prompt.
7853        let ring = if state_ring_rows == 0 {
7854            0
7855        } else {
7856            state_ring_rows / ig.pool * ig.pool
7857        };
7858        if state_ring_rows > 0 && ring == 0 {
7859            return Err(format!(
7860                "indexer tail ring of {state_ring_rows} rows cannot hold one pool of {}; \
7861                 raise MEMRA_DSA_INDEX_RING or set it to 0 for the flat plane",
7862                ig.pool
7863            )
7864            .into());
7865        }
7866        // RESIDENCY TRIPWIRE. `*pools_ready` counts pools whose keys were built over rows that are
7867        // now history. If the cache ever rewound past `slot` without clamping it (see
7868        // `LatentKvLayer::truncate_index_pool_keys`), those keys were built over rows this call is
7869        // about to overwrite — a silent wrong selection. Fail here instead.
7870        if *pools_ready > slot / ig.pool {
7871            return Err(format!(
7872                "resident k-pool key plane claims {} finished pools but the cache holds only {} \
7873                 complete pools before this call ({slot} rows / pool {}) — a rewind reduced the \
7874                 latent length without clamping index_pools_ready",
7875                *pools_ready,
7876                slot / ig.pool,
7877                ig.pool
7878            )
7879            .into());
7880        }
7881
7882        // 1. packed state rows [k | gate], appended at `slot` — the same [a|b] row shape the
7883        //    latent plane uses, so `mla_append_latent` packs it with no new kernel.
7884        let k_raw = mm(&indexer.wk, h)?;
7885        let mut k_norm = e.uninit(t * d)?;
7886        e.layer_norm_bias(
7887            &k_raw,
7888            indexer.k_norm_w.float_data(),
7889            indexer.k_norm_b.float_data(),
7890            &mut k_norm,
7891            d,
7892            t,
7893            INDEX_NORM_EPS,
7894        )?;
7895        let gate = mm(&indexer.kpool_gate, h)?;
7896
7897        // 2. pool keys over every COMPLETE pool in the cache — INCREMENTALLY. A pool's key is a
7898        //    function of its own `pool` state rows (append-only, never rewritten) and the constant
7899        //    `kpool_ape`, so it is final the instant the pool's last row lands. Only pools
7900        //    `[*pools_ready, n_pools)` are built; the rest are already resident and bit-identical
7901        //    to a rebuild. This turns the old O(t_kv * d) per-call pass into O(t * d).
7902        let n_pools = t_kv / ig.pool;
7903        let select_k = ig.select_k(n_pools);
7904        let width = ig.index_width(n_pools);
7905        // Sized to the SESSION's capacity, so a session that primes and then decodes never
7906        // reallocates (a fresh buffer would drop every resident key, and under the ring the rows
7907        // to rebuild them from are gone). `capacity_tokens` is that capacity; a pool covers
7908        // `ig.pool` tokens, so the key plane is `pool` times SHORTER than a flat state plane —
7909        // 32 f32 per token against 256. It is NOT read off `plane.len()` any more: once the state
7910        // plane is a ring, its length is one call's tail, not the context.
7911        // `.max(1)` keeps the slice non-null at t_kv < pool, where no complete pool exists yet.
7912        let capacity_pools = capacity_tokens / ig.pool;
7913        let need = (capacity_pools * d).max(n_pools * d).max(1);
7914        if pool_key_plane.as_ref().is_none_or(|k| k.len() < need) {
7915            *pool_key_plane = Some(e.uninit(need)?);
7916            *pools_ready = 0;
7917        }
7918        let pool_keys = pool_key_plane
7919            .as_mut()
7920            .expect("resident pool-key plane just allocated");
7921
7922        // THE DRAIN. The state plane is written by exactly one kernel and read by exactly one,
7923        // and a row's single read is the pool-key build of the pool that row belongs to. So the
7924        // rows that must be live at any instant are `[*pools_ready * pool, cur)`: everything
7925        // below has been read, everything above is not written yet, and the two kernels can be
7926        // interleaved in sub-ranges of the call instead of run once each over the whole call.
7927        //
7928        // That is what makes the ring size a WORKING-SET choice rather than a bound on `t`:
7929        // `index_ring_take` hands back how many rows fit before the ring must be drained, the
7930        // build drains it, and the loop continues. `k_norm`/`gate` are computed ONCE for the
7931        // whole call above and walked by source-row offset, so the values, their order, and the
7932        // ring addresses they land on are exactly what a single whole-call append produced.
7933        // A flat plane (`ring == 0`) takes the whole call in one iteration, byte for byte.
7934        let ape = indexer.kpool_ape.float_data();
7935        let mut cur = slot;
7936        let mut appended = 0usize;
7937        while appended < t {
7938            let take =
7939                crate::cache::index_ring_take(ring, ig.pool, *pools_ready, cur, t - appended)
7940                    .ok_or_else(|| -> Box<dyn std::error::Error> {
7941                        format!(
7942                            "indexer tail ring lapped: {ring} rows cannot hold the {} rows still \
7943                         owed to unbuilt pools at row {cur} (pools_ready {}, pool {}, slot \
7944                         {slot}, t {t}). The pool-key plane was reset or the cache rewound \
7945                         without clamping index_pools_ready, so rows this call must read were \
7946                         already overwritten. Raise MEMRA_DSA_INDEX_RING, or set \
7947                         MEMRA_DSA_INDEX_RING=0 for the flat plane",
7948                            cur.saturating_sub((*pools_ready).saturating_mul(ig.pool)),
7949                            *pools_ready,
7950                            ig.pool
7951                        )
7952                        .into()
7953                    })?;
7954            debug_assert!(take > 0 && appended + take <= t);
7955            e.mla_index_append(plane, &k_norm, &gate, appended, cur, take, d, d, ring)?;
7956            cur += take;
7957            appended += take;
7958            let ready_now = cur / ig.pool;
7959            e.mla_kpool_pool_keys(
7960                plane,
7961                ape,
7962                pool_keys,
7963                (*pools_ready).min(ready_now),
7964                ready_now,
7965                ig.pool,
7966                d,
7967                ring,
7968            )?;
7969            *pools_ready = ready_now;
7970        }
7971        debug_assert!(t == 0 || *pools_ready == n_pools);
7972        let pool_keys = &*pool_keys;
7973
7974        // 3. score + head mix, 4. top-k -> raw rows + tail.
7975        let q_index = mm(&indexer.wq_b, q_resid)?;
7976        let head_weights = mm(&indexer.weights_proj, h)?;
7977        let mut score = e.uninit((t * n_pools).max(1))?;
7978        e.mla_kpool_score(
7979            &q_index,
7980            pool_keys,
7981            &head_weights,
7982            &mut score,
7983            t,
7984            ig.heads,
7985            d,
7986            n_pools,
7987            ig.pool,
7988            slot,
7989            (d as f32).powf(-0.5),
7990            (ig.heads as f32).powf(-0.5),
7991        )?;
7992        let mut idx = e.uninit_i32(t * width)?;
7993        e.mla_kpool_select(
7994            &score,
7995            &mut idx,
7996            t,
7997            n_pools,
7998            ig.pool,
7999            select_k,
8000            width,
8001            slot,
8002            ig.always_select_tail,
8003        )?;
8004        Ok((idx, width))
8005    }
8006
8007    /// STATELESS MLA arm (`HybridModel::forward`): the latent plane lives for this call only,
8008    /// sized to the request. Same math as the cached arm — it is the same core with slot 0.
8009    pub fn mla_attn(
8010        &self,
8011        e: &Engine,
8012        mla: &crate::hybrid::MlaAttnLayer,
8013        h: &CudaSlice<f32>,
8014        pos_d: &CudaSlice<i32>,
8015        t: usize,
8016        il: usize,
8017    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8018        if mla.tp.is_some() {
8019            return Err(format!(
8020                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the stateless \
8021                 mixer path is unwired for a head shard"
8022            )
8023            .into());
8024        }
8025        let mut latent = e.uninit(t * mla.geom.latent_dim)?;
8026        let mut index_plane = match mla.index.as_ref() {
8027            Some(indexer) => Some(e.uninit(t * indexer.geom.state_width())?),
8028            None => None,
8029        };
8030        // No residency across calls here — the planes die with the call, so `ready` starts at 0 and
8031        // every pool is built exactly once, which is what the cached arm also does on its prime.
8032        let mut pool_keys = None;
8033        let mut pools_ready = 0usize;
8034        let planes = index_plane.as_mut().map(|state| IndexerPlanes {
8035            state,
8036            pool_keys: &mut pool_keys,
8037            ready: &mut pools_ready,
8038            // Per-call plane, sized to the request: no ring, capacity is the request itself.
8039            state_ring_rows: 0,
8040            capacity_tokens: t,
8041        });
8042        self.mla_attn_core(e, mla, h, pos_d, t, il, &mut latent, planes, 0, false)
8043    }
8044
8045    /// STATEFUL MLA arm (prime and T=1 decode): appends into the session's latent plane and
8046    /// attends the whole history. `cache.latent[il]` is allocated by the `LatentKvCache` arm of
8047    /// the cache allocator; a `None` here means the ModelPlan and the loaded mixer disagree.
8048    #[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
8049    pub fn mla_attn_cached(
8050        &self,
8051        e: &Engine,
8052        mla: &crate::hybrid::MlaAttnLayer,
8053        h: &CudaSlice<f32>,
8054        pos_d: &CudaSlice<i32>,
8055        t: usize,
8056        il: usize,
8057        cache: &mut Cache,
8058    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8059        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, false)
8060    }
8061
8062    /// [`Self::mla_attn_cached`] on the VERIFY-BATCH matmul classes (lane/glm5-verify-batch):
8063    /// the SAME core at t=K+1 rows with every internal projection routed decode-exact
8064    /// (`matmul_rows_exact`), so row r of the batched call is bit-identical to the t=1
8065    /// `mla_attn_cached` call the per-row verify walk makes at position pos0+r. Causality
8066    /// within the batch is per-query by construction: the kpool selection masks pools
8067    /// invisible to each query and appends each query's OWN raw tail
8068    /// (`first_pos + t + 1`), and the gathered attention walks each query's own idx list
8069    /// (-1 padding arithmetic-invariant). Held by `glm5_tparallel_verify_gpu` gates 1+2
8070    /// running the batched arm. ONLY the glm5 verify-batch walk calls this.
8071    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
8072    pub fn mla_attn_cached_rows_exact(
8073        &self,
8074        e: &Engine,
8075        mla: &crate::hybrid::MlaAttnLayer,
8076        h: &CudaSlice<f32>,
8077        pos_d: &CudaSlice<i32>,
8078        t: usize,
8079        il: usize,
8080        cache: &mut Cache,
8081    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8082        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, true)
8083    }
8084
8085    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
8086    fn mla_attn_cached_inner(
8087        &self,
8088        e: &Engine,
8089        mla: &crate::hybrid::MlaAttnLayer,
8090        h: &CudaSlice<f32>,
8091        pos_d: &CudaSlice<i32>,
8092        t: usize,
8093        il: usize,
8094        cache: &mut Cache,
8095        rows_exact: bool,
8096    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8097        // glm5 TP fail-closed choke point, covering BOTH plain entries (decode/prime AND
8098        // the verify-batch rows arm): a TP-sharded layer holds heads/2 and a per-rank
8099        // latent replica — running it on the plain path would compute a silently-halved
8100        // mixer against the wrong plane, so it refuses by name instead.
8101        if mla.tp.is_some() {
8102            return Err(format!(
8103                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the plain mixer \
8104                 path is unwired for a head shard — only the TP decode/prime walk may \
8105                 execute it (rows_exact={rows_exact})"
8106            )
8107            .into());
8108        }
8109        // Read before the layer borrow: it sizes the resident pool-key plane, which the ring'd
8110        // state plane's own length can no longer stand in for.
8111        let max_ctx = cache.max_ctx;
8112        let layer = cache.latent[il].as_mut().ok_or_else(|| {
8113            format!(
8114                "layer {il} is Mixer::Mla but the cache has no latent plane — the ModelPlan \
8115                 must declare StatePlan::LatentKvCache for it"
8116            )
8117        })?;
8118        let attn =
8119            self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?;
8120        // Verify-batch wo seam: the rows arm keeps its decode-exact output projection —
8121        // the wo dispatch moved here with the TP split, its routing did not change.
8122        if rows_exact {
8123            e.matmul_rows_exact(&mla.wo, &attn, t)
8124        } else {
8125            e.matmul(&mla.wo, &attn, t)
8126        }
8127    }
8128
8129    /// The stateful MLA call against ONE latent plane, up to (and excluding) the output
8130    /// projection. The plain path wraps it above (canonical plane + `wo`); the glm5 TP-2
8131    /// walk calls it once per rank (root shard on the canonical plane, peer shard on the
8132    /// replicated peer plane) and joins the halves through the column-parallel `wo`.
8133    #[allow(clippy::too_many_arguments)]
8134    pub(crate) fn mla_attn_cached_pre_wo(
8135        &self,
8136        e: &Engine,
8137        mla: &crate::hybrid::MlaAttnLayer,
8138        h: &CudaSlice<f32>,
8139        pos_d: &CudaSlice<i32>,
8140        t: usize,
8141        il: usize,
8142        layer: &mut memra_kv::LatentKvLayer,
8143        max_ctx: usize,
8144        rows_exact: bool,
8145    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8146        let slot = layer.len;
8147        let width = layer.width;
8148        assert_eq!(
8149            width, mla.geom.latent_dim,
8150            "layer {il}: cache latent width {width} != MlaGeom latent_dim {}",
8151            mla.geom.latent_dim
8152        );
8153        let capacity = layer.rows.len() / width;
8154        if slot + t > capacity {
8155            return Err(format!(
8156                "layer {il}: latent cache overflow — {slot} + {t} rows exceeds capacity {capacity}"
8157            )
8158            .into());
8159        }
8160        if mla.index.is_some() && layer.index_rows.is_none() {
8161            return Err(format!(
8162                "layer {il} loaded a DSA k-pool indexer but its latent cache carries no indexer \
8163                 state plane — StatePlan::LatentKvCache declared index_width 0 for a layer whose \
8164                 SparseIndexPlan is Own {{ kpool: Some(..) }}"
8165            )
8166            .into());
8167        }
8168        // Both planes are borrowed for the whole core call, so len bookkeeping happens after.
8169        // ONE `len` covers both: they are appended in the same call and must never drift.
8170        // The resident pool-key plane rides along: it is state that must SURVIVE the call, so the
8171        // core writes `ready` back through the borrow and it is restored with the buffers.
8172        let mut rows = std::mem::replace(&mut layer.rows, e.uninit(0)?);
8173        let mut index_rows = layer.index_rows.take();
8174        let mut pool_keys = layer.index_pool_keys.take();
8175        let mut pools_ready = layer.index_pools_ready;
8176        let index_ring_rows = layer.index_ring_rows.unwrap_or(0);
8177        let planes = index_rows.as_mut().map(|state| IndexerPlanes {
8178            state,
8179            pool_keys: &mut pool_keys,
8180            ready: &mut pools_ready,
8181            state_ring_rows: index_ring_rows,
8182            capacity_tokens: max_ctx,
8183        });
8184        let out =
8185            self.mla_attn_core_pre_wo(e, mla, h, pos_d, t, il, &mut rows, planes, slot, rows_exact);
8186        layer.rows = rows;
8187        layer.index_rows = index_rows;
8188        layer.index_pool_keys = pool_keys;
8189        // A FAILED core leaves `len` where it was, so the resident plane must go back too: it may
8190        // have advanced over pools built from rows a retry is about to rewrite with different
8191        // inputs. Clamping here (rather than letting the next call's tripwire fire) makes
8192        // retry-after-error correct instead of merely loud.
8193        layer.index_pools_ready = if out.is_ok() {
8194            pools_ready
8195        } else if let Some(indexer) = mla.index.as_ref() {
8196            pools_ready.min(layer.len / indexer.geom.pool)
8197        } else {
8198            pools_ready
8199        };
8200        let out = out?;
8201        // Resolve the layer's RESIDENT pool copy (the state plan does not carry `pool`; the
8202        // latent-plane snapshot path reads this field to address the tail ring). A nonzero
8203        // resident value that disagrees with the loaded geometry is corruption, not a race:
8204        // there is exactly one geometry per loaded layer.
8205        if let Some(indexer) = mla.index.as_ref() {
8206            let pool = indexer.geom.pool;
8207            if layer.index_pool != 0 && layer.index_pool != pool {
8208                return Err(format!(
8209                    "layer {il}: resident indexer pool {} != loaded geometry pool {pool}",
8210                    layer.index_pool,
8211                )
8212                .into());
8213            }
8214            layer.index_pool = pool;
8215        }
8216        layer.len = slot + t;
8217        let len_i32 = i32::try_from(layer.len).map_err(|_| "latent length exceeds i32 mirror")?;
8218        // Door H (`MEMRA_GLM5_HTOD_DIET`): the async `i32_set_k` launch instead of this
8219        // SYNCHRONIZING pageable 4-byte copy — 11 of these per round, one per MLA trunk layer.
8220        e.i32_mirror_store(&mut layer.len_d, len_i32)?;
8221        Ok(out)
8222    }
8223
8224    /// The glm5 TP MLA walk for one prime/decode call (`mla` is the ROOT head shard; its
8225    /// sidecar carries the peer shards + runtime). Replicated per-token work runs on EVERY
8226    /// rank from identical inputs (wq_a/wkv_a/indexer/k-pool selection — identical bytes by
8227    /// determinism on uniform hardware, gate-held); each rank attends with its heads over
8228    /// its OWN latent replica; the attention parts are gathered through the armed transport
8229    /// and each rank's COLUMN-parallel `wo` slice computes its slice of the output with the
8230    /// plain matvec kernel — no cross-rank arithmetic anywhere.
8231    #[allow(clippy::too_many_arguments)]
8232    pub(crate) fn mla_tp_attn_cached(
8233        &self,
8234        e: &Engine,
8235        mla: &crate::hybrid::MlaAttnLayer,
8236        h: &CudaSlice<f32>,
8237        pos_d: &CudaSlice<i32>,
8238        t: usize,
8239        il: usize,
8240        cache: &mut Cache,
8241        rows_exact: bool,
8242    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8243        let tp = mla
8244            .tp
8245            .as_ref()
8246            .ok_or("mla_tp_attn_cached called on an unsharded layer")?;
8247        let rt = &tp.rt;
8248        let ranks = tp.ranks();
8249        let g = mla.geom; // SHARD geometry: n_head = full/ranks
8250        let hl = g.n_head;
8251        let dv = g.d_v;
8252        let full_heads = tp.full_heads;
8253        let n_embd = tp.n_embd;
8254        let hh = n_embd / ranks;
8255        let max_ctx = cache.max_ctx;
8256
8257        // HOP 1 — fan-out of the mixer input and positions to every peer rank. Both move the
8258        // WHOLE buffer, exactly as the v1 arm did, so the transport arms move identical
8259        // byte ranges (lane/glm5-tp-transport).
8260        let hop = rt.hop(e);
8261        let h_peers = crate::tp_transport::fanout_f32(&hop, h, h.len())?;
8262        let pos_peers = crate::tp_transport::fanout_i32(&hop, pos_d, pos_d.len())?;
8263
8264        // Peer replica planes (lazily geometry-cloned from the canonical plane).
8265        {
8266            let canonical = cache.latent[il].as_ref().ok_or_else(|| {
8267                format!("layer {il}: glm5 TP MLA walk found no canonical latent plane")
8268            })?;
8269            crate::glm5_tp::ensure_mla_peer_latent(
8270                rt,
8271                canonical,
8272                &mut cache.glm5_tp_latent_peer[il],
8273            )?;
8274        }
8275
8276        // Peer passes first (each rank's heads over its replica), then root (canonical
8277        // plane unchanged) — v1's issue order at two ranks. `rows_exact` threads the
8278        // caller's matmul class through every rank: false = the prime/decode walk
8279        // (byte-for-byte the pre-composition arm), true = the spec x TP verify walk
8280        // (lane/glm5-composition) riding the same rows-exact classes as the unsharded
8281        // verify walk.
8282        let mut attn: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
8283        for r in 1..ranks {
8284            let layer = &mut cache.glm5_tp_latent_peer[il].as_mut().unwrap()[r - 1];
8285            attn[r] = Some(self.mla_attn_cached_pre_wo(
8286                &rt.peers[r - 1],
8287                &tp.peers[r - 1],
8288                &h_peers[r - 1],
8289                &pos_peers[r - 1],
8290                t,
8291                il,
8292                layer,
8293                max_ctx,
8294                rows_exact,
8295            )?);
8296        }
8297        attn[0] = {
8298            let layer = cache.latent[il].as_mut().unwrap();
8299            Some(self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?)
8300        };
8301
8302        // HOP 2 — gather the per-head parts into the FULL [t, full_heads*dv] layout on
8303        // every rank. `full_heads * dv == ranks * (hl * dv)` by the shard map.
8304        let part = hl * dv;
8305        debug_assert_eq!(full_heads * dv, ranks * part);
8306        let attn_refs: Vec<&CudaSlice<f32>> = attn
8307            .iter()
8308            .map(|a| a.as_ref().expect("filled above"))
8309            .collect();
8310        let fulls = crate::tp_transport::gather_parts(&hop, &attn_refs, t, part)?;
8311
8312        // Column-parallel wo slices + output concat (pure movement). The verify walk's
8313        // wo rides the rows-exact class, exactly like the unsharded verify walk's wo.
8314        let mut ys = Vec::with_capacity(ranks);
8315        if rows_exact {
8316            ys.push(e.matmul_rows_exact(&mla.wo, &fulls[0], t)?);
8317            for r in 1..ranks {
8318                ys.push(rt.peers[r - 1].matmul_rows_exact(&tp.peers[r - 1].wo, &fulls[r], t)?);
8319            }
8320        } else {
8321            ys.push(e.matmul(&mla.wo, &fulls[0], t)?);
8322            for r in 1..ranks {
8323                ys.push(rt.peers[r - 1].matmul(&tp.peers[r - 1].wo, &fulls[r], t)?);
8324            }
8325        }
8326        // HOP 3 — concat the column parts into the mixer output on ROOT.
8327        debug_assert_eq!(n_embd, ranks * hh);
8328        let y_refs: Vec<&CudaSlice<f32>> = ys.iter().collect();
8329        crate::tp_transport::concat_parts_on_root(&hop, &y_refs, t, hh)
8330    }
8331
8332    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
8333    pub fn linear_attn(
8334        &self,
8335        e: &Engine,
8336        la: &LinearAttnLayer,
8337        h: &CudaSlice<f32>,
8338        t: usize,
8339    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8340        let cfg = &self.cfg;
8341        let _n_embd = cfg.n_embd as usize;
8342        let geometry = la.geometry;
8343        let d_state = geometry.key_head_dim as usize;
8344        let num_k = geometry.key_heads as usize;
8345        let num_v = geometry.value_heads as usize;
8346        let d_conv = geometry.conv_kernel as usize;
8347        let head_k = d_state;
8348        let head_v = geometry.value_head_dim as usize;
8349        let key_dim = head_k * num_k; // 2048
8350        let value_dim = head_v * num_v; // 4096
8351        let conv_dim = key_dim * 2 + value_dim; // 8192
8352        let eps = cfg.rms_eps;
8353        let scale = 1.0 / (d_state as f32).sqrt();
8354
8355        // projections
8356        // grouped: one f16 activation convert feeds all four projections (matmul_group)
8357        let mut g4 = e.matmul_group(
8358            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8359            h,
8360            t,
8361        )?;
8362        let alpha = g4.pop().unwrap(); // [T, num_v]
8363        let beta_raw = g4.pop().unwrap(); // [T, num_v]
8364        let z = g4.pop().unwrap(); // [T, value_dim]
8365        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
8366
8367        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
8368        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
8369        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
8370        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
8371        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
8372        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
8373        let _ = (head_k, head_v);
8374        let mut q_g = e.uninit(d_state * num_v * t)?;
8375        let mut k_g = e.uninit(d_state * num_v * t)?;
8376        let mut v_g = e.uninit(d_state * num_v * t)?;
8377        e.ssm_conv1d_gdn(
8378            &qkv_mixed,
8379            la.ssm_conv1d.float_data(),
8380            &mut q_g,
8381            &mut k_g,
8382            &mut v_g,
8383            conv_dim,
8384            t,
8385            d_conv,
8386            d_state,
8387            num_v,
8388            num_k,
8389            key_dim,
8390        )?;
8391        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
8392        let mut q_l2 = e.uninit(d_state * num_v * t)?;
8393        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8394        let mut k_l2 = e.uninit(d_state * num_v * t)?;
8395        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8396        let v_gd = v_g;
8397
8398        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
8399        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
8400        let mut beta = e.uninit(t * num_v)?;
8401        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8402        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
8403        let mut g_log = e.uninit(t * num_v)?;
8404        e.gdn_glog(
8405            &alpha,
8406            la.ssm_dt.float_data(),
8407            la.ssm_a.float_data(),
8408            &mut g_log,
8409            num_v,
8410            t,
8411        )?;
8412
8413        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
8414        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
8415        let mut state_out = e.zeros(d_state * d_state * num_v)?;
8416        let mut o = e.uninit(d_state * num_v * t)?;
8417        e.gdn_scan_prefill(
8418            &q_l2,
8419            &k_l2,
8420            &v_gd,
8421            &g_log,
8422            &beta,
8423            None,
8424            None,
8425            &state_in,
8426            &mut state_out,
8427            &mut o,
8428            num_v,
8429            t,
8430            scale,
8431            num_v,
8432        )?;
8433
8434        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
8435        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
8436        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
8437        // o rows are (t*num_v+vh) too. Good.
8438        let mut gn = e.uninit(d_state * num_v * t)?;
8439        e.gated_rmsnorm(
8440            &o,
8441            la.ssm_norm.float_data(),
8442            &z,
8443            &mut gn,
8444            d_state,
8445            num_v * t,
8446            eps,
8447        )?;
8448
8449        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
8450        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
8451        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
8452        let out = e.matmul(&la.ssm_out, &gn, t)?;
8453        Ok(out)
8454    }
8455}
8456
8457impl HybridModel {
8458    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
8459    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
8460    ///
8461    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
8462    /// different 860160-byte block than the same expert of layer 7).
8463    ///
8464    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
8465    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
8466    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
8467    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
8468    pub fn moe_ffn_il(
8469        &self,
8470        e: &Engine,
8471        m: &MoeWeights,
8472        z: &CudaSlice<f32>,
8473        t: usize,
8474        il: u16,
8475    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8476        Self::moe_ffn_inner(
8477            e,
8478            m,
8479            z,
8480            None,
8481            t,
8482            &self.cfg,
8483            il,
8484            self.max_moe_block(),
8485            false,
8486            None,
8487            self.uses_sliding_gated_moe_program(),
8488            false,
8489        )
8490    }
8491
8492    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
8493    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
8494    pub fn moe_ffn_il_prefill(
8495        &self,
8496        e: &Engine,
8497        m: &MoeWeights,
8498        z: &CudaSlice<f32>,
8499        t: usize,
8500        il: u16,
8501    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8502        Self::moe_ffn_inner(
8503            e,
8504            m,
8505            z,
8506            None,
8507            t,
8508            &self.cfg,
8509            il,
8510            self.max_moe_block(),
8511            true,
8512            Some(&self.step_grouped_prefill),
8513            self.uses_sliding_gated_moe_program(),
8514            false,
8515        )
8516    }
8517
8518    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
8519    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
8520    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
8521    pub fn moe_ffn_il_zq8(
8522        &self,
8523        e: &Engine,
8524        m: &MoeWeights,
8525        z: &CudaSlice<f32>,
8526        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8527        t: usize,
8528        il: u16,
8529    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8530        Self::moe_ffn_inner(
8531            e,
8532            m,
8533            z,
8534            zq8,
8535            t,
8536            &self.cfg,
8537            il,
8538            self.max_moe_block(),
8539            false,
8540            None,
8541            self.uses_sliding_gated_moe_program(),
8542            false,
8543        )
8544    }
8545
8546    /// Verify-rows twin of [`Self::moe_ffn_il_zq8`] (lane/glm5-vrest): the SAME routing and
8547    /// dispatch decisions with the pairs-shaped batched routed-expert arm armed. Only the
8548    /// verify walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`, t>=2) calls this; every
8549    /// unqualified shape inside falls closed to the byte-identical sequential loop.
8550    pub(crate) fn moe_ffn_il_zq8_vrows(
8551        &self,
8552        e: &Engine,
8553        m: &MoeWeights,
8554        z: &CudaSlice<f32>,
8555        t: usize,
8556        il: u16,
8557    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8558        Self::moe_ffn_inner(
8559            e,
8560            m,
8561            z,
8562            None,
8563            t,
8564            &self.cfg,
8565            il,
8566            self.max_moe_block(),
8567            false,
8568            None,
8569            self.uses_sliding_gated_moe_program(),
8570            true,
8571        )
8572    }
8573
8574    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
8575    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
8576    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
8577    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
8578    ///
8579    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
8580    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
8581    pub(crate) fn moe_ffn(
8582        e: &Engine,
8583        m: &MoeWeights,
8584        z: &CudaSlice<f32>,
8585        t: usize,
8586        cfg: &ModelConfig,
8587        il: u16,
8588        max_block: usize,
8589    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8590        Self::moe_ffn_inner(
8591            e, m, z, None, t, cfg, il, max_block, false, None, false, false,
8592        )
8593    }
8594
8595    #[allow(clippy::too_many_arguments)]
8596    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
8597    pub(crate) fn moe_ffn_inner(
8598        e: &Engine,
8599        m: &MoeWeights,
8600        z: &CudaSlice<f32>,
8601        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8602        t: usize,
8603        cfg: &ModelConfig,
8604        il: u16,
8605        max_block: usize,
8606        prefill: bool,
8607        grouped_prefill: Option<&std::sync::Mutex<crate::hybrid::StepEpGroupedPrefill>>,
8608        sliding_gated_moe: bool,
8609        vrows: bool,
8610    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8611        let worker_io = crate::spill_pread::worker_enabled();
8612        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
8613        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
8614            e.with_moe_cache(max_block, |cache, _| {
8615                cache.begin_forward_epoch(il, t);
8616                if worker_io {
8617                    cache.begin_worker_scope();
8618                }
8619                Ok(())
8620            })?;
8621        }
8622        if let Some(ep) = &m.glm5_ep {
8623            // glm5 TP-2 EP walk (MEMRA_GLM5_TP): whole-expert halves, root router, slot-ordered
8624            // canonical combine. Every other arm of this function is unreachable for an
8625            // EP-armed layer by construction. `prefill` keys the EP grouped-prime arm
8626            // (MEMRA_GLM5_EP_GROUPED_PRIME) exactly as it keys the plain grouped arm below.
8627            return Self::moe_ffn_glm5_ep(e, m, ep, z, zq8, t, cfg, il, prefill);
8628        }
8629        if m.step_ep.is_some() || m.step_tp.is_some() {
8630            let moe = cfg
8631                .moe
8632                .as_ref()
8633                .ok_or("Step distributed execution requires MoE model metadata")?;
8634            let n_embd = cfg.n_embd as usize;
8635            let n_expert = moe.expert_count as usize;
8636            let n_used = moe.expert_used_count as usize;
8637            let sigmoid = cfg
8638                .sigmoid_router()
8639                .ok_or("Step distributed execution requires the Step sigmoid router")?;
8640            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
8641            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
8642            let grouped_prefill_requested = prefill && step_ep_grouped_prefill_enabled()?;
8643            if grouped_prefill_requested && !step_tp_prefill_enabled()? {
8644                return Err(
8645                    "MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into(),
8646                );
8647            }
8648            if grouped_prefill_requested && !step_grouped_prefill_shape(true, prefill, t) {
8649                return Err(format!(
8650                    "Step grouped prefill tokens {t} are outside the qualified {}..={} range",
8651                    PRIME_MIN_T,
8652                    crate::cache::PRIME_CHUNK_MAX_TOKENS,
8653                )
8654                .into());
8655            }
8656            let grouped_decode_shape = step_grouped_decode_shape(prefill, t);
8657            let grouped_prefill_shape =
8658                step_grouped_prefill_shape(grouped_prefill_requested, prefill, t);
8659            if let Some(ep) = m.step_ep.as_ref().filter(|ep| {
8660                ep.grouped_decode.is_some() && (grouped_decode_shape || grouped_prefill_shape)
8661            }) {
8662                let (selected, route_weights) =
8663                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8664                crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
8665                Self::trace_moe_routes(il, t, &selected, &route_weights)?;
8666                Self::trace_moe_input(e, il, t, n_embd, z)?;
8667                let selected = selected
8668                    .iter()
8669                    .map(|&expert| expert as usize)
8670                    .collect::<Vec<_>>();
8671
8672                // The narrow route readback above orders the owning-stage producer. The grouped
8673                // runtime then copies the resident root activation into its persistent rank inputs.
8674                e.stream().synchronize()?;
8675                let execute = |state: &mut crate::hybrid::StepEpGroupedDecode| {
8676                    state.projection.set_activation_limit(ep.activation_limit)?;
8677                    ep.runtime
8678                        .refresh_step_grouped_expert_parallel_gate_from_root_device(
8679                            ep.experts.e4m3()?,
8680                            &mut state.projection,
8681                            z,
8682                            t,
8683                            &selected,
8684                        )?;
8685                    ep.runtime.refresh_step_grouped_expert_parallel_combine(
8686                        &state.projection,
8687                        &mut state.combine,
8688                        &route_weights,
8689                    )?;
8690                    ep.runtime.execute_step_grouped_expert_parallel_gate(
8691                        ep.experts.e4m3()?,
8692                        &mut state.projection,
8693                    )?;
8694                    ep.runtime.execute_step_grouped_expert_parallel_combine(
8695                        &state.projection,
8696                        &mut state.combine,
8697                    )?;
8698                    let mut output = ep.runtime.copy_step_grouped_expert_parallel_combine_root(
8699                        &state.projection,
8700                        &state.combine,
8701                        e,
8702                    )?;
8703                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8704                    if prefill {
8705                        // A shared plan may be reused by the next layer on a different runtime
8706                        // stream. Complete the owning-stage copy before its source is overwritten.
8707                        e.stream().synchronize()?;
8708                    }
8709                    eprintln!(
8710                        "[step-tp-ep-grouped] execute layer={il} tokens={t} devices={:?} \
8711                         attention_layout=tensor-parallel expert_layout=expert-parallel \
8712                         expert_transport={} native_p2p=true route_control=host-narrow \
8713                         input=root-device projection_workspaces=persistent \
8714                         combine=root-device output=owning-stage-device \
8715                         prefill={prefill} batched_decode=false capacity={} \
8716                         performance_claim=false",
8717                        ep.devices,
8718                        ep.runtime.transport_label(),
8719                        state.projection.max_tokens(),
8720                    );
8721                    Ok::<_, Box<dyn std::error::Error>>(output)
8722                };
8723
8724                if grouped_prefill_shape {
8725                    let grouped_prefill = grouped_prefill
8726                        .ok_or("Step grouped prefill has no model-scoped executor")?;
8727                    let mut shared = grouped_prefill
8728                        .lock()
8729                        .map_err(|_| "Step grouped prefill state lock is poisoned")?;
8730                    let needs_prepare = shared.state.as_ref().is_none_or(|state| {
8731                        state.devices != ep.devices
8732                            || state.grouped.projection.max_tokens() < t
8733                            || state.grouped.projection.input_width() != n_embd
8734                            || state.grouped.projection.expert_width()
8735                                != moe.expert_ff_length as usize
8736                    });
8737                    if needs_prepare {
8738                        let seed_input = vec![0.0f32; n_embd];
8739                        let seed_selected = &selected[..n_used];
8740                        let seed_weights = &route_weights[..n_used];
8741                        let projection = ep
8742                            .runtime
8743                            .prepare_step_grouped_expert_parallel_gate_with_capacity(
8744                                ep.experts.e4m3()?,
8745                                &seed_input,
8746                                1,
8747                                seed_selected,
8748                                ep.activation_limit,
8749                                t,
8750                            )?;
8751                        let combine = ep.runtime.prepare_step_grouped_expert_parallel_combine(
8752                            &projection,
8753                            seed_weights,
8754                        )?;
8755                        shared.state = Some(crate::hybrid::StepEpGroupedPrefillState {
8756                            devices: ep.devices.clone(),
8757                            grouped: crate::hybrid::StepEpGroupedDecode {
8758                                projection,
8759                                combine,
8760                            },
8761                        });
8762                        eprintln!(
8763                            "[step-tp-ep-grouped-prefill] prepare capacity={t} devices={:?} \
8764                             shared_across_layers=true performance_claim=false",
8765                            ep.devices,
8766                        );
8767                    }
8768                    return execute(
8769                        &mut shared
8770                            .state
8771                            .as_mut()
8772                            .expect("Step grouped prefill state prepared above")
8773                            .grouped,
8774                    );
8775                }
8776
8777                let mut grouped = ep
8778                    .grouped_decode
8779                    .as_ref()
8780                    .expect("grouped decode presence checked above")
8781                    .lock()
8782                    .map_err(|_| "Step grouped decode state lock is poisoned")?;
8783                return execute(&mut grouped);
8784            }
8785            if grouped_prefill_shape {
8786                return Err(
8787                    "Step grouped prefill requires native-P2P expert-owner device arithmetic"
8788                        .into(),
8789                );
8790            }
8791            // MEMRA_STEP_TP_DEV_ROUTER=1 (t=1): route on device and feed the device-routed
8792            // expert program — the per-layer host logits readback (the last per-layer host
8793            // sync) disappears. Selection tie-breaking may differ from the host router:
8794            // numeric-class door, run-gen argmax gate + boot battery.
8795            // STEP TP2 GEMM PRIME (2026-08-27, TTFT lane): a prime chunk's routed MoE goes
8796            // through ONE grouped f16 GEMM per projection over the resident NVFP4 banks —
8797            // the per-token device routes below cost 240 s at m=4092 (measured), the grouped
8798            // lane's sizing rows run 170-270 TFLOP/s. Router selections come from the same
8799            // sigmoid host oracle the EP arm uses; shexp rides the canonical grouped add.
8800            // t>=16 alone keys the branch: the batch prime reaches here through moe_ffn_il,
8801            // whose `prefill` is FALSE (only the _prefill twin sets it), and no other step37
8802            // route runs t>=16 — verify walks t<=8, decode t=1. Requiring `prefill` made the
8803            // first gate arm skip this branch entirely and wake the generic f16g arm instead
8804            // (48 s + kq_gemm_sk rc=1001, 2026-08-27).
8805            if t >= 16
8806                && crate::step_gemm_prime_on()
8807                && let Some(tp) = &m.step_tp
8808                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8809            {
8810                // MEMRA_PRIME_PROF=1 sub-split of the moe bucket. The phase timer put
8811                // 1788 ms of a 3093 ms chunk here, but forcing the 32-row tile form (4x
8812                // more weight dequant) moved it only 5% — so the grouped GEMM is not
8813                // obviously what dominates. The router below is a HOST oracle: sigmoid +
8814                // top-8 over 288 experts for every one of 4096 tokens, per layer, which
8815                // is a D2H copy and a full pipeline drain 42 times per chunk. Attribute
8816                // it before optimizing the kernel it sits in front of.
8817                let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1") && t >= 16;
8818                let mut mt = std::time::Instant::now();
8819                let (selected, route_weights) =
8820                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8821                let sel_i32: Vec<i32> = selected.iter().map(|&x| x as i32).collect();
8822                let d_router = if mprof {
8823                    let _ = e.stream().synchronize();
8824                    let v = mt.elapsed().as_secs_f64() * 1e3;
8825                    mt = std::time::Instant::now();
8826                    v
8827                } else {
8828                    0.0
8829                };
8830                // MEMRA_MOE_DETERM=1: run the WHOLE grouped routine twice on identical
8831                // inputs and diff. The standalone harness cleared the grouped GEMM kernels
8832                // (8 invocations, both lanes, maxdiff 0.0 over 20.9M elements) but it does
8833                // not model the cross-device join/scatter or the o_proj-style reduction,
8834                // and the loader refuses both topologies (TP1, same-device) that would
8835                // isolate those by env. This tests the un-excluded region directly, in
8836                // the place it actually runs.
8837                //
8838                // The prime is nondeterministic: same prompt, one forward, temperature=0,
8839                // max_tokens=1, and the first token varies across reps. That blocks
8840                // MEMRA_PP_BF16's correctness receipt and invalidates every byte-identity
8841                // gate taken through the server. This probe also yields the jitter
8842                // MAGNITUDE, which any tolerance band needs.
8843                let mdet =
8844                    std::env::var("MEMRA_MOE_DETERM").as_deref() == Ok("1") && t >= 16 && il < 4;
8845                if mdet {
8846                    let a = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8847                        bank,
8848                        e,
8849                        z,
8850                        t,
8851                        &sel_i32,
8852                        &route_weights,
8853                        n_used,
8854                        tp.activation_limit,
8855                    )?;
8856                    let b = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8857                        bank,
8858                        e,
8859                        z,
8860                        t,
8861                        &sel_i32,
8862                        &route_weights,
8863                        n_used,
8864                        tp.activation_limit,
8865                    )?;
8866                    let (ha, hb) = (e.dtoh(&a)?, e.dtoh(&b)?);
8867                    let mut md = 0.0f32;
8868                    let mut ndiff = 0usize;
8869                    for (x, y) in ha.iter().zip(hb.iter()) {
8870                        let d = (x - y).abs();
8871                        if d > 0.0 {
8872                            ndiff += 1;
8873                        }
8874                        if d > md {
8875                            md = d;
8876                        }
8877                    }
8878                    eprintln!(
8879                        "[moe-determ] il={il} t={t} maxdiff={md:.3e} \
8880                                 differing={ndiff}/{} -> {}",
8881                        ha.len(),
8882                        if ndiff == 0 {
8883                            "IDENTICAL"
8884                        } else {
8885                            "NONDETERMINISTIC"
8886                        }
8887                    );
8888                }
8889                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8890                    bank,
8891                    e,
8892                    z,
8893                    t,
8894                    &sel_i32,
8895                    &route_weights,
8896                    n_used,
8897                    tp.activation_limit,
8898                )?;
8899                let d_gemm = if mprof {
8900                    let _ = e.stream().synchronize();
8901                    let v = mt.elapsed().as_secs_f64() * 1e3;
8902                    mt = std::time::Instant::now();
8903                    v
8904                } else {
8905                    0.0
8906                };
8907                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8908                if mprof {
8909                    let _ = e.stream().synchronize();
8910                    let d_shared = mt.elapsed().as_secs_f64() * 1e3;
8911                    // Per LAYER, not accumulated: the four trunk phases already carry the
8912                    // per-chunk totals, and one line per layer is what shows whether the
8913                    // cost is flat across layers or concentrated in a few.
8914                    eprintln!(
8915                        "[moe-prof] il={il} t={t} router={d_router:.1}ms \
8916                                 gemm={d_gemm:.1}ms shared={d_shared:.1}ms"
8917                    );
8918                }
8919                return Ok(output);
8920            }
8921            if t == 1
8922                && crate::tp::step_nvfp4_dev_routes_enabled()?
8923                && crate::tp::step_tp_dev_router_enabled()?
8924                && let Some(tp) = &m.step_tp
8925                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8926            {
8927                let (sf, route_norm) = sigmoid;
8928                // MEMRA_ROUTES_PRESTAGE=1: issue the rank input pull + quantize NOW,
8929                // before the router — the rank streams overlap the gemv+topk.
8930                // MEMRA_DEV1_ROUTER=1 rides the prestage hook: rank1 routes LOCALLY
8931                // from its own z copy (replicated deterministic router — identical
8932                // bits in, identical sel/w out) and starts its sweep without
8933                // waiting the root's sel broadcast.
8934                static D1_ROUTER: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8935                let d1_router = *D1_ROUTER
8936                    .get_or_init(|| std::env::var("MEMRA_DEV1_ROUTER").as_deref() == Ok("1"));
8937                if d1_router {
8938                    let (sf_h, rn_h) = sigmoid;
8939                    let n_ex = m.gate_exps.n_expert;
8940                    let act_ct = m.active_count();
8941                    let _ = tp.runtime.nvfp4_routes_prestage_with(
8942                        bank,
8943                        e,
8944                        z,
8945                        |rank1, in1, sel1, w1| {
8946                            let mut guard = DEV1_ROUTER_REPS
8947                                .lock()
8948                                .map_err(|_| "dev1 router replica lock")?;
8949                            let (reps, scratch) =
8950                                guard.get_or_insert_with(|| (Default::default(), None));
8951                            if !reps.contains_key(&il) {
8952                                use cudarc::driver::DevicePtr;
8953                                let (g1, p1, a1) = (
8954                                    rank1.htod(&vec![0.0f32; n_ex * n_embd])?,
8955                                    rank1.htod(&vec![0.0f32; n_ex])?,
8956                                    rank1.alloc_u8_uninit(n_ex)?,
8957                                );
8958                                for (src, dst_len, dst) in [
8959                                    (
8960                                        {
8961                                            let s = e.stream();
8962                                            let (p, _g) = m.gate_inp.float_data().device_ptr(&s);
8963                                            p
8964                                        },
8965                                        n_ex * n_embd * 4,
8966                                        {
8967                                            let s = rank1.stream();
8968                                            let (p, _g) = g1.device_ptr(&s);
8969                                            p
8970                                        },
8971                                    ),
8972                                    (
8973                                        {
8974                                            let s = e.stream();
8975                                            let (p, _g) = m.exp_probs_b_dev.device_ptr(&s);
8976                                            p
8977                                        },
8978                                        n_ex * 4,
8979                                        {
8980                                            let s = rank1.stream();
8981                                            let (p, _g) = p1.device_ptr(&s);
8982                                            p
8983                                        },
8984                                    ),
8985                                    (
8986                                        {
8987                                            let s = e.stream();
8988                                            let (p, _g) = m.active_experts_dev.device_ptr(&s);
8989                                            p
8990                                        },
8991                                        n_ex,
8992                                        {
8993                                            let s = rank1.stream();
8994                                            let (p, _g) = a1.device_ptr(&s);
8995                                            p
8996                                        },
8997                                    ),
8998                                ] {
8999                                    crate::tp::raw_copy_bytes(dst, src, dst_len, rank1)?;
9000                                }
9001                                rank1.stream().synchronize()?;
9002                                reps.insert(il, (g1, p1, a1));
9003                            }
9004                            if scratch.is_none() {
9005                                *scratch = Some(rank1.htod(&vec![0.0f32; n_ex])?);
9006                            }
9007                            let (g1, p1, a1) = reps.get(&il).expect("armed above");
9008                            let logits1 = scratch.as_mut().expect("armed above");
9009                            rank1.router_gemv_into(g1, in1, logits1, n_embd, n_ex, 1)?;
9010                            rank1.moe_router_sigmoid_topk_into(
9011                                logits1, 1, n_ex, n_used, act_ct, p1, a1, sf_h, rn_h, sel1, w1,
9012                            )?;
9013                            Ok(true)
9014                        },
9015                    )?;
9016                } else {
9017                    let _ = tp.runtime.nvfp4_routes_prestage(bank, e, z)?;
9018                }
9019                // Persistent selection buffers: the allocating topk built two fresh
9020                // slices per layer; sel/w land in process-static rows instead
9021                // (host-op diet — same kernel, same bytes).
9022                #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
9023                static SELW: std::sync::Mutex<Option<(usize, CudaSlice<i32>, CudaSlice<f32>)>> =
9024                    std::sync::Mutex::new(None);
9025                let mut selw = SELW.lock().map_err(|_| "selw lock poisoned")?;
9026                if selw.as_ref().is_none_or(|(d, ..)| *d != e.ctx().ordinal()) {
9027                    *selw = Some((
9028                        e.ctx().ordinal(),
9029                        e.htod_i32(&vec![0i32; n_used])?,
9030                        e.htod(&vec![0.0f32; n_used])?,
9031                    ));
9032                }
9033                let (_, sel_d, w_d) = selw.as_mut().expect("armed above");
9034                e.moe_router_sigmoid_topk_into(
9035                    &logits,
9036                    t,
9037                    n_expert,
9038                    n_used,
9039                    m.active_count(),
9040                    &m.exp_probs_b_dev,
9041                    &m.active_experts_dev,
9042                    sf,
9043                    route_norm,
9044                    sel_d,
9045                    w_d,
9046                )?;
9047                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
9048                // FAIL-CLOSED for the route taps: this walk keeps the selection
9049                // device-side, so `trace_moe_routes` (MEMRA_MOE_TRACE /
9050                // MEMRA_MOE_WEIGHT_TRACE) never sees its rows. Every other MoE walk is
9051                // either host-routed (the taps ride the existing readback) or forced to
9052                // the host-visible path by observation mode — this one is neither. A
9053                // trace that silently misses whole layers poisons any placement mint
9054                // built on it (LAW:coactivation-expert-placement measurement leg), so an
9055                // armed tap refuses by name instead of dropping rows.
9056                if std::env::var("MEMRA_MOE_TRACE").is_ok()
9057                    || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9058                {
9059                    return Err("MEMRA_MOE_TRACE/MEMRA_MOE_WEIGHT_TRACE cannot trace the \
9060                         device-routed step TP walk (selection never returns to host; \
9061                         tracing would add a new sync). Route through the host-router \
9062                         arm — refused rather than silently dropping rows"
9063                        .into());
9064                }
9065                // MEMRA_SHEXP_OVERLAP=1: issue the shared expert from the routes
9066                // PREJOIN hook so it executes while the peer rank drains its sweep
9067                // (fills dev0's join wait); apply adds the identical values after.
9068                static SHEXP_OV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9069                let shexp_ov = *SHEXP_OV
9070                    .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
9071                // MEMRA_SHEXP_DEV1=1 (supersedes the dev0 overlap): the shared
9072                // expert runs on rank1 — the idle device — same kernels, same
9073                // split program, down row root-resident: bit-identical.
9074                static SHEXP_D1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9075                let shexp_d1 = *SHEXP_D1
9076                    .get_or_init(|| std::env::var("MEMRA_SHEXP_DEV1").as_deref() == Ok("1"))
9077                    && tp.runtime.rank_engine(1).is_some();
9078                // MOE TAIL FUSION M1 (MEMRA_TAIL_ADD3=0 reverts): pre-arm the
9079                // overlap ws + ones row and hand their RAW pointers to the routed
9080                // run — the join add folds the shexp apply into one launch.
9081                static TAIL3: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9082                let tail3 =
9083                    *TAIL3.get_or_init(|| std::env::var("MEMRA_TAIL_ADD3").as_deref() != Ok("0"));
9084                let mut ov_issued = false;
9085                let mut d1_issued = false;
9086                let mut tail_folded = false;
9087                let mut output = if shexp_d1 {
9088                    let rank1 = tp.runtime.rank_engine(1).expect("checked above");
9089                    tp.runtime
9090                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin(
9091                            bank,
9092                            e,
9093                            z,
9094                            sel_d,
9095                            w_d,
9096                            n_used,
9097                            tp.activation_limit,
9098                            || {
9099                                d1_issued =
9100                                    Self::shexp_dev1_issue(e, rank1, m, z, cfg, il, n_embd)?;
9101                                Ok(())
9102                            },
9103                        )?
9104                } else if shexp_ov {
9105                    // Raw sh/ones pointers for the fused tail (persistent statics;
9106                    // pointers stable, no lock held across the routed call). The
9107                    // sh CONTENT is written by the prejoin-issued kernels earlier
9108                    // on e's stream — stream order covers the fused add.
9109                    let post_add = if tail3 {
9110                        Self::shexp_overlap_tail_ptrs(e, m, cfg, n_embd)?
9111                    } else {
9112                        None
9113                    };
9114                    let used_post = post_add.is_some();
9115                    let out = tp
9116                        .runtime
9117                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3(
9118                            bank,
9119                            e,
9120                            z,
9121                            sel_d,
9122                            w_d,
9123                            n_used,
9124                            tp.activation_limit,
9125                            || {
9126                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9127                                Ok(())
9128                            },
9129                            post_add,
9130                        )?;
9131                    // ov_issued false with post_add armed = an early-return arm
9132                    // (the GRAPH door) skipped the prejoin AND ignored post_add —
9133                    // fall through to the normal shexp add (battery v22 receipt:
9134                    // the strict error here failed every graph-door boot).
9135                    if used_post && ov_issued {
9136                        tail_folded = true; // apply folded into the join add
9137                    }
9138                    out
9139                } else {
9140                    tp.runtime.run_tensor_parallel_routes_nvfp4_device_routed(
9141                        bank,
9142                        e,
9143                        z,
9144                        sel_d,
9145                        w_d,
9146                        n_used,
9147                        tp.activation_limit,
9148                    )?
9149                };
9150                if output.len() != t * n_embd {
9151                    return Err(format!(
9152                        "Step tp routed output has {} values, expected {t}x{n_embd}",
9153                        output.len()
9154                    )
9155                    .into());
9156                }
9157                if tail_folded {
9158                    // shexp already folded into the join add (MOE TAIL FUSION M1)
9159                } else if d1_issued {
9160                    Self::shexp_dev1_apply(e, &mut output, n_embd)?;
9161                } else if ov_issued {
9162                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9163                } else {
9164                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9165                }
9166                static DR_LOGGED: std::sync::atomic::AtomicU64 =
9167                    std::sync::atomic::AtomicU64::new(0);
9168                let layer_bit = 1u64 << (il as u64 % 64);
9169                if DR_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9170                    == 0
9171                {
9172                    eprintln!(
9173                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
9174                                 expert_transport={} native_p2p={} router=device \
9175                                 activation=host-canonical accumulation=host-canonical \
9176                                 output=e-device io=device performance_claim=false \
9177                                 (logged once per layer)",
9178                        tp.devices,
9179                        tp.runtime.transport_label(),
9180                        tp.runtime.native_p2p(),
9181                    );
9182                }
9183                return Ok(output);
9184            }
9185            let automatic_ep_device_router = crate::tp::parallel_ep_device_router_enabled()?;
9186            let automatic_ep_q8_act = crate::tp::parallel_ep_q8_act_enabled()?;
9187            let automatic_ep_q8_scope = crate::tp::parallel_ep_q8_scope()?;
9188            crate::tp::parallel_ep_q8_gu_paired_enabled(
9189                automatic_ep_q8_act,
9190                automatic_ep_q8_scope,
9191            )?;
9192            let automatic_ep_q8_active =
9193                automatic_ep_q8_act && t <= crate::tp::NVFP4_EP_Q8_BATCH_CAP;
9194            if automatic_ep_q8_scope.is_some() && !automatic_ep_q8_act {
9195                return Err(
9196                    "MEMRA_PARALLEL_EP_Q8_SCOPE requires MEMRA_PARALLEL_EP_Q8_ACT=1".into(),
9197                );
9198            }
9199            if automatic_ep_q8_act && !automatic_ep_device_router {
9200                return Err(
9201                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires MEMRA_PARALLEL_EP_DEVICE_ROUTER=1".into(),
9202                );
9203            }
9204            if automatic_ep_q8_act && m.step_ep.as_ref().is_none_or(|ep| !ep.nvfp4_device_routes) {
9205                return Err(
9206                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires automatic W4A16 whole-expert EP".into(),
9207                );
9208            }
9209            if t <= crate::tp::NVFP4_EP_DEVICE_ROUTER_BATCH_CAP
9210                && automatic_ep_device_router
9211                && let Some(ep) = &m.step_ep
9212                && ep.nvfp4_device_routes
9213            {
9214                let bank = match &ep.experts {
9215                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9216                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
9217                        return Err("W4A16 device-routed EP reached an E4M3 expert bank".into());
9218                    }
9219                };
9220                let pairs = t
9221                    .checked_mul(n_used)
9222                    .ok_or("W4A16 device-routed EP pair count overflow")?;
9223                let capacity = crate::tp::NVFP4_EP_DEVICE_BATCH_CAP * n_used;
9224                /// Per-device persistent route scratch: device ordinal -> (armed capacity in
9225                /// pairs, selected-expert rows, route-weight rows). Named because the nested
9226                /// form is unreadable at this depth, not to hide it.
9227                type EpSelwByDevice =
9228                    std::collections::HashMap<usize, (usize, CudaSlice<i32>, CudaSlice<f32>)>;
9229                static EP_SELW: std::sync::Mutex<Option<EpSelwByDevice>> =
9230                    std::sync::Mutex::new(None);
9231                let mut selw = EP_SELW
9232                    .lock()
9233                    .map_err(|_| "automatic EP device-router workspace lock poisoned")?;
9234                let device = e.ctx().ordinal();
9235                let workspaces = selw.get_or_insert_with(Default::default);
9236                if workspaces
9237                    .get(&device)
9238                    .is_none_or(|(cap, ..)| *cap < capacity)
9239                {
9240                    workspaces.insert(
9241                        device,
9242                        (
9243                            capacity,
9244                            e.htod_i32(&vec![0i32; capacity])?,
9245                            e.htod(&vec![0.0f32; capacity])?,
9246                        ),
9247                    );
9248                }
9249                let (_, sel_d, w_d) = workspaces.get_mut(&device).expect("armed above");
9250                let (sf, route_norm) = sigmoid;
9251                e.moe_router_sigmoid_topk_into(
9252                    &logits,
9253                    t,
9254                    n_expert,
9255                    n_used,
9256                    m.active_count(),
9257                    &m.exp_probs_b_dev,
9258                    &m.active_experts_dev,
9259                    sf,
9260                    route_norm,
9261                    sel_d,
9262                    w_d,
9263                )?;
9264                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
9265                static SHEXP_OV_AUTO: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9266                let shexp_ov = t == 1
9267                    && *SHEXP_OV_AUTO
9268                        .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
9269                let mut ov_issued = false;
9270                let mut output = if shexp_ov {
9271                    ep.runtime
9272                        .run_routed_experts_nvfp4_w4a16_device_routed_prejoin(
9273                            bank,
9274                            e,
9275                            z,
9276                            sel_d,
9277                            w_d,
9278                            t,
9279                            n_used,
9280                            ep.activation_limit,
9281                            || {
9282                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9283                                Ok(())
9284                            },
9285                        )?
9286                } else {
9287                    ep.runtime.run_routed_experts_nvfp4_w4a16_device_routed(
9288                        bank,
9289                        e,
9290                        z,
9291                        sel_d,
9292                        w_d,
9293                        t,
9294                        n_used,
9295                        ep.activation_limit,
9296                    )?
9297                };
9298                if output.len() != t * n_embd {
9299                    return Err(format!(
9300                        "W4A16 device-routed EP output has {} values, expected \
9301                         {t}x{n_embd}={}",
9302                        output.len(),
9303                        t * n_embd,
9304                    )
9305                    .into());
9306                }
9307                if ov_issued {
9308                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9309                } else {
9310                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9311                }
9312                static DEVICE_ROUTER_LOGGED: std::sync::atomic::AtomicU64 =
9313                    std::sync::atomic::AtomicU64::new(0);
9314                let layer_bit = 1u64 << (il as u64 % 64);
9315                if DEVICE_ROUTER_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9316                    & layer_bit
9317                    == 0
9318                {
9319                    eprintln!(
9320                        "[parallel-ep] execute layer={il} tokens={t} devices={:?} \
9321                         router=device expert_transport={} native_p2p={} \
9322                         activation=bf16-rounded accumulation={} output=e-device \
9323                         performance_claim=false (logged once per layer)",
9324                        ep.devices,
9325                        ep.runtime.transport_label(),
9326                        ep.runtime.native_p2p(),
9327                        if automatic_ep_q8_active {
9328                            "token-slot-order-q8"
9329                        } else {
9330                            "token-slot-order"
9331                        },
9332                    );
9333                }
9334                debug_assert!(pairs <= capacity);
9335                return Ok(output);
9336            }
9337            // MEMRA_STEP_TP_TIMING=1: cumulative cost of the host routing seam (the dtoh here
9338            // drains every e-stream op queued since the layer's FFN entry, so this bills the
9339            // router matmul + glue too — the decode-bucket ffn residue decomposes here).
9340            static ROUTE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9341            static ROUTE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9342            let route_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
9343            let route_started = route_timing.then(std::time::Instant::now);
9344            let (selected, route_weights, input) = Self::moe_route_sigmoid_with_input(
9345                e,
9346                &logits,
9347                z,
9348                t,
9349                n_embd,
9350                n_expert,
9351                n_used,
9352                m.exp_probs_b.as_deref(),
9353                sigmoid,
9354                m.active_experts.as_deref(),
9355            )?;
9356            if let Some(started) = route_started {
9357                use std::sync::atomic::Ordering;
9358                let ns = ROUTE_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
9359                    + started.elapsed().as_nanos() as u64;
9360                let calls = ROUTE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
9361                if calls.is_multiple_of(430) {
9362                    eprintln!(
9363                        "[moe-route-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
9364                        ns as f64 / 1.0e6,
9365                        ns as f64 / calls as f64 / 1.0e3,
9366                    );
9367                }
9368            }
9369            crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
9370            Self::trace_moe_routes(il, t, &selected, &route_weights)?;
9371            Self::trace_moe_input(e, il, t, n_embd, z)?;
9372            let selected = selected
9373                .iter()
9374                .map(|&expert| expert as usize)
9375                .collect::<Vec<_>>();
9376            if t <= crate::tp::NVFP4_EP_DEVICE_BATCH_CAP
9377                && let Some(ep) = &m.step_ep
9378                && ep.nvfp4_device_routes
9379            {
9380                let bank = match &ep.experts {
9381                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9382                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
9383                        return Err("W4A16 NVFP4 device EP reached an E4M3 expert bank".into());
9384                    }
9385                };
9386                let mut output = ep.runtime.run_routed_experts_nvfp4_w4a16_device_io(
9387                    bank,
9388                    e,
9389                    z,
9390                    t,
9391                    &selected,
9392                    &route_weights,
9393                    n_used,
9394                    ep.activation_limit,
9395                )?;
9396                if output.len() != t * n_embd {
9397                    return Err(format!(
9398                        "W4A16 NVFP4 EP routed output has {} values, expected \
9399                         {t}x{n_embd}={}",
9400                        output.len(),
9401                        t * n_embd,
9402                    )
9403                    .into());
9404                }
9405                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9406                static W4A16_EP_LOGGED: std::sync::atomic::AtomicU64 =
9407                    std::sync::atomic::AtomicU64::new(0);
9408                let layer_bit = 1u64 << (il as u64 % 64);
9409                if W4A16_EP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9410                    & layer_bit
9411                    == 0
9412                {
9413                    eprintln!(
9414                        "[step-ep] execute layer={il} tokens={t} devices={:?} \
9415                         expert_transport={} native_p2p={} activation=bf16-rounded \
9416                        accumulation={} output=e-device \
9417                         performance_claim=false (logged once per layer)",
9418                        ep.devices,
9419                        ep.runtime.transport_label(),
9420                        ep.runtime.native_p2p(),
9421                        if t == 1 {
9422                            "owner-grouped-rank-order"
9423                        } else {
9424                            "token-slot-order"
9425                        },
9426                    );
9427                }
9428                return Ok(output);
9429            }
9430            // Device-IO routes (t=1): the layer input goes to the ranks as a device row and the
9431            // combined output comes back as an e-context row — no host round-trip, no host
9432            // stream sync. Program bytes identical to the host-IO twin (dtoh/htod and dtod
9433            // both preserve f32 bits), gated by greedy token identity.
9434            if t == 1
9435                && crate::tp::step_nvfp4_dev_routes_enabled()?
9436                && let Some(tp) = &m.step_tp
9437                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
9438            {
9439                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_device_io(
9440                    bank,
9441                    e,
9442                    z,
9443                    &selected,
9444                    &route_weights,
9445                    n_used,
9446                    tp.activation_limit,
9447                )?;
9448                if output.len() != t * n_embd {
9449                    return Err(format!(
9450                        "Step tp routed output has {} values, expected {t}x{n_embd}",
9451                        output.len()
9452                    )
9453                    .into());
9454                }
9455                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9456                static IO_LOGGED: std::sync::atomic::AtomicU64 =
9457                    std::sync::atomic::AtomicU64::new(0);
9458                let layer_bit = 1u64 << (il as u64 % 64);
9459                if IO_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9460                    == 0
9461                {
9462                    eprintln!(
9463                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
9464                                 expert_transport={} native_p2p={} activation=host-canonical \
9465                                 accumulation=host-canonical output=e-device io=device \
9466                                 performance_claim=false (logged once per layer)",
9467                        tp.devices,
9468                        tp.runtime.transport_label(),
9469                        tp.runtime.native_p2p(),
9470                    );
9471                }
9472                return Ok(output);
9473            }
9474            let (routed, mode, devices, transport, native_p2p) = if let Some(tp) = &m.step_tp {
9475                (
9476                    match &tp.experts {
9477                        crate::hybrid::StepTpExpertBank::E4m3(bank) => {
9478                            tp.runtime.run_tensor_parallel_routes(
9479                                bank,
9480                                &input,
9481                                t,
9482                                &selected,
9483                                &route_weights,
9484                                n_used,
9485                            )?
9486                        }
9487                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => {
9488                            if t == 1 && crate::tp::step_nvfp4_dev_routes_enabled()? {
9489                                tp.runtime.run_tensor_parallel_routes_nvfp4_device(
9490                                    bank,
9491                                    &input,
9492                                    &selected,
9493                                    &route_weights,
9494                                    n_used,
9495                                    tp.activation_limit,
9496                                )?
9497                            } else {
9498                                tp.runtime.run_tensor_parallel_routes_nvfp4(
9499                                    bank,
9500                                    &input,
9501                                    t,
9502                                    &selected,
9503                                    &route_weights,
9504                                    n_used,
9505                                    tp.activation_limit,
9506                                )?
9507                            }
9508                        }
9509                    },
9510                    "tp",
9511                    &tp.devices,
9512                    tp.runtime.transport_label(),
9513                    tp.runtime.native_p2p(),
9514                )
9515            } else {
9516                let ep = m
9517                    .step_ep
9518                    .as_ref()
9519                    .ok_or("Step distributed runtime has no EP or TP state")?;
9520                (
9521                    match &ep.experts {
9522                        crate::hybrid::StepEpExpertBank::E4m3(bank) => {
9523                            ep.runtime.run_routed_experts(
9524                                bank,
9525                                &input,
9526                                t,
9527                                &selected,
9528                                &route_weights,
9529                                n_used,
9530                                ep.activation_limit,
9531                            )?
9532                        }
9533                        crate::hybrid::StepEpExpertBank::Nvfp4(bank) => {
9534                            ep.runtime.run_routed_experts_nvfp4(
9535                                bank,
9536                                &input,
9537                                t,
9538                                &selected,
9539                                &route_weights,
9540                                n_used,
9541                                ep.activation_limit,
9542                            )?
9543                        }
9544                    },
9545                    if ep.configured_by_tp { "tp-ep" } else { "ep" },
9546                    &ep.devices,
9547                    ep.runtime.transport_label(),
9548                    ep.runtime.native_p2p(),
9549                )
9550            };
9551            if routed.len() != t * n_embd {
9552                return Err(format!(
9553                    "Step {mode} routed output has {} values, expected {t}x{n_embd}",
9554                    routed.len()
9555                )
9556                .into());
9557            }
9558            let mut output = e.htod(&routed)?;
9559            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9560            // Once per layer per process: the topology contract line is a boot receipt, not a
9561            // per-token trace — 4520 of these per 64-token run measured as real decode wall.
9562            static STEP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9563            let layer_bit = 1u64 << (il as u64 % 64);
9564            if STEP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9565                == 0
9566            {
9567                eprintln!(
9568                    "[step-{mode}] execute layer={il} tokens={t} devices={devices:?} \
9569                     expert_transport={transport} native_p2p={native_p2p} \
9570                     activation={} accumulation={} output={} \
9571                     performance_claim=false (logged once per layer)",
9572                    if let Some(ep) = &m.step_ep {
9573                        ep.runtime.expert_activation_label()
9574                    } else {
9575                        "host-canonical"
9576                    },
9577                    if let Some(ep) = &m.step_ep {
9578                        ep.runtime.expert_accumulation_label()
9579                    } else {
9580                        "host-canonical"
9581                    },
9582                    if let Some(ep) = &m.step_ep {
9583                        ep.runtime.expert_output_label()
9584                    } else {
9585                        "host-accumulated"
9586                    },
9587                );
9588                if let Some(ep) = &m.step_ep
9589                    && let Some(limit) = ep.activation_limit
9590                {
9591                    eprintln!(
9592                        "[step-ep-clamp] execute layer={il} tokens={t} routed_clamp={limit} \
9593                             formula=min-silu-times-clamped-up performance_claim=false"
9594                    );
9595                }
9596            }
9597            return Ok(output);
9598        }
9599        if Self::sigmoid_resident_dev_eligible(e, m, cfg, sliding_gated_moe) {
9600            let moe = cfg.moe.as_ref().unwrap();
9601            let n_expert = moe.expert_count as usize;
9602            let n_used = moe.expert_used_count as usize;
9603            let sigmoid = cfg.sigmoid_router().unwrap();
9604            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9605            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
9606            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
9607        }
9608        // GROUPED MoE PREFILL, sigmoid-router class (glm5_next), MEMRA_MOE_GROUPED_PREFILL
9609        // default ON since 2026-08-29 (owner-accepted flip; =0 rollback seam; receipts on the
9610        // flag helper). The prefill-gap attribution (research/glm53-flash-bringup-20260827/
9611        // prefill-gap-20260829/PREFILL-GAP.md §1.1) measured this arch prefilling every prompt
9612        // token through the decode program: 49 launches per token-layer, ~8.4M launches and
9613        // 4.76 GB of expert-weight VRAM re-reads per token across 42 layers per 4096-token
9614        // chunk, because every batched arm is predicate-denied for sigmoid-router archs. This
9615        // arm is the composition of qualified ingredients: the m-invariant router + sigmoid
9616        // host oracle (routing sel/w BIT-identical to the sequential arm by construction),
9617        // host token-sort by expert (the moe_align_block_size shape), one grouped NVFP4
9618        // tensor-core GEMM per projection (the step37 grouped-prime kernel class via
9619        // `moe_f16_grouped`, generalized to the single-device resident slab), the PRE-clamped
9620        // SwiGLU epilogue and the per-expert weight_scale_2 macro fold the fused-epilogue lane
9621        // gated for this family. Keyed on `prefill` (only the _prefill twin sets it) so decode,
9622        // spec verify and the exact-16 batched-decode tier keep their dispatch class, and on
9623        // `t > MOE_DEV_MAX_T` so t<=16 stays on the per-token program (grouped/pairs prefill
9624        // classes start at 17, same seam as the softmax pairs arm).
9625        // ENGAGEMENT RECEIPT: the announce below prints in BOTH arms (flag on and off), once
9626        // per process, so an A/B grep distinguishes engagement without the line itself being
9627        // an arm-local cost (the step37 engagement-receipt trap: prove the path RAN before
9628        // attributing a number to it).
9629        if prefill && t > MOE_DEV_MAX_T && cfg.sigmoid_router().is_some() && cfg.glm5.is_some() {
9630            // Once per process PER FLAG VALUE (bit 0 = off, bit 1 = on): a server boot prints
9631            // exactly one line, and a gate process that flips the flag shows both arms.
9632            static GPF_ANNOUNCED: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
9633            let enabled = moe_grouped_prefill_enabled();
9634            let bit = 1u8 << u8::from(enabled);
9635            if GPF_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
9636                eprintln!(
9637                    "[moe-grouped-prefill] flag={} t={t} il={il} (announce printed in both \
9638                     arms; engagement is the per-layer execute line + the dispatch counter)",
9639                    if enabled { "on" } else { "off" },
9640                );
9641            }
9642            if enabled
9643                && let Some(out) = Self::moe_ffn_grouped_prefill_sigmoid(e, m, z, t, cfg, il)?
9644            {
9645                return Ok(out);
9646            }
9647        }
9648        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
9649        // current caller into this research arm; the naked default stays on the established path.
9650        if t > 1 && moe_grouped_enabled(cfg, prefill) {
9651            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
9652            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
9653            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
9654            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
9655            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
9656            if std::env::var("MEMRA_MOE_GATE").is_ok() {
9657                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
9658                let g_host = e.dtoh(&grouped_out)?;
9659                let s_host = e.dtoh(&seq_out)?;
9660                let g_bytes: &[u8] = unsafe {
9661                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
9662                };
9663                let s_bytes: &[u8] = unsafe {
9664                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
9665                };
9666                if g_bytes == s_bytes {
9667                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
9668                } else {
9669                    let diffs = g_host
9670                        .iter()
9671                        .zip(s_host.iter())
9672                        .enumerate()
9673                        .filter(|(_, (a, b))| a != b)
9674                        .count();
9675                    let maxdiff = g_host
9676                        .iter()
9677                        .zip(s_host.iter())
9678                        .map(|(a, b)| (a - b).abs())
9679                        .fold(0.0f32, f32::max);
9680                    panic!(
9681                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
9682                        g_host.len()
9683                    );
9684                }
9685            }
9686            return Ok(grouped_out);
9687        }
9688        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block, vrows)
9689    }
9690
9691    fn sigmoid_resident_dev_eligible(
9692        e: &Engine,
9693        m: &MoeWeights,
9694        cfg: &ModelConfig,
9695        sliding_gated_moe: bool,
9696    ) -> bool {
9697        let Some(moe) = cfg.moe.as_ref() else {
9698            return false;
9699        };
9700        // Cached once per process: this predicate runs per MoE layer per decode step, and five
9701        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
9702        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9703        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
9704            std::env::var("MEMRA_MOE_STATS").is_ok()
9705                || std::env::var("MEMRA_MOE_TRACE").is_ok()
9706                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9707                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
9708                || std::env::var("MEMRA_MOE_GATE").is_ok()
9709        });
9710        let resident_layout_supported = m.dev_exps.as_ref().is_some_and(|dev| {
9711            if dev.dev != e.ctx().ordinal() {
9712                return false;
9713            }
9714            let q8 = moe_q8_enabled_for_model(cfg, m);
9715            let fp8 = dev.fp8_blk.is_some()
9716                && m.gate_exps.qtype == crate::QT_F8_E4M3_BLK
9717                && m.up_exps.qtype == crate::QT_F8_E4M3_BLK
9718                && m.down_exps.qtype == crate::QT_F8_E4M3_BLK;
9719            q8 || fp8
9720        });
9721        sliding_gated_moe
9722            && sigmoid_router_enabled()
9723            && moe_dev_enabled()
9724            && moe_slab_enabled()
9725            && !observation_mode
9726            && moe.expert_used_count <= 8
9727            && m.has_uniform_expert_layout()
9728            && m.gate_exps.macros.is_none()
9729            && m.up_exps.macros.is_none()
9730            && m.down_exps.macros.is_none()
9731            && !m.has_macros
9732            && resident_layout_supported
9733    }
9734
9735    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
9736    pub(crate) fn moe_ffn_sequential(
9737        e: &Engine,
9738        m: &MoeWeights,
9739        z: &CudaSlice<f32>,
9740        t: usize,
9741        cfg: &ModelConfig,
9742        il: u16,
9743        max_block: usize,
9744    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9745        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block, false)
9746    }
9747
9748    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
9749    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
9750    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
9751    fn moe_router_logits(
9752        e: &Engine,
9753        m: &MoeWeights,
9754        z: &CudaSlice<f32>,
9755        t: usize,
9756        cfg: &ModelConfig,
9757    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9758        if t < PRIME_MIN_T {
9759            // Decode and speculative verify use one fixed per-row reduction program.
9760            if crate::router_kernel_on() {
9761                e.router_gemv(
9762                    m.gate_inp.float_data(),
9763                    z,
9764                    cfg.n_embd as usize,
9765                    m.gate_exps.n_expert,
9766                    t,
9767                )
9768            } else {
9769                e.matmul_decode_exact(&m.gate_inp, z, t)
9770            }
9771        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
9772            e.router_gemv(
9773                m.gate_inp.float_data(),
9774                z,
9775                cfg.n_embd as usize,
9776                m.gate_exps.n_expert,
9777                t,
9778            )
9779        } else {
9780            e.matmul(&m.gate_inp, z, t)
9781        }
9782    }
9783
9784    /// Append the host-visible router selection for one layer/forward when calibration tracing is
9785    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
9786    /// trace is independent of the dispatch optimization selected for the forward.
9787    /// `MEMRA_MOE_WEIGHT_TRACE` is also the co-activation measurement input of
9788    /// LAW:coactivation-expert-placement (lane/glm5-ep-place: rows ride this existing host
9789    /// readback — zero new device syncs; `glm5-tp-gate` arm T holds the ON-identity +
9790    /// row-count bar on the glm5 walks).
9791    fn trace_moe_routes(
9792        il: u16,
9793        t: usize,
9794        sel_all: &[u32],
9795        weights: &[f32],
9796    ) -> Result<(), Box<dyn std::error::Error>> {
9797        use std::io::Write as _;
9798        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
9799            let mut f = std::fs::OpenOptions::new()
9800                .create(true)
9801                .append(true)
9802                .open(path)?;
9803            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
9804            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
9805        }
9806        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
9807            let mut f = std::fs::OpenOptions::new()
9808                .create(true)
9809                .append(true)
9810                .open(path)?;
9811            let pairs: Vec<String> = sel_all
9812                .iter()
9813                .zip(weights)
9814                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
9815                .collect();
9816            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
9817        }
9818        Ok(())
9819    }
9820
9821    #[allow(clippy::too_many_arguments)]
9822    fn trace_sigmoid_router_logits(
9823        e: &Engine,
9824        il: u16,
9825        t: usize,
9826        n_expert: usize,
9827        n_used: usize,
9828        logits: &CudaSlice<f32>,
9829        m: &MoeWeights,
9830        (scaling_factor, route_norm): (f32, bool),
9831    ) -> Result<(), Box<dyn std::error::Error>> {
9832        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
9833            return Ok(());
9834        }
9835        let logits = e.dtoh(logits)?;
9836        let active: Vec<u8> = m
9837            .active_experts
9838            .as_ref()
9839            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
9840            .unwrap_or_else(|| vec![1; n_expert]);
9841        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
9842        crate::sigrouter_contract::capture_served_logits(
9843            il as u32,
9844            t,
9845            n_expert,
9846            n_used,
9847            scaling_factor,
9848            route_norm,
9849            &active,
9850            &bias,
9851            &logits,
9852        )?;
9853        Ok(())
9854    }
9855
9856    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
9857    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
9858    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
9859    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
9860    fn trace_moe_input(
9861        e: &Engine,
9862        il: u16,
9863        t: usize,
9864        n_embd: usize,
9865        z: &CudaSlice<f32>,
9866    ) -> Result<(), Box<dyn std::error::Error>> {
9867        use std::io::Write as _;
9868        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
9869            return Ok(());
9870        };
9871        let values = active_matrix_values(z.len(), t, n_embd, "MoE input trace activation")?;
9872        let host = e.dtoh_view(&z.slice(0..values))?;
9873        let bytes = unsafe {
9874            std::slice::from_raw_parts(
9875                host.as_ptr().cast::<u8>(),
9876                host.len() * std::mem::size_of::<f32>(),
9877            )
9878        };
9879        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
9880        let mut state = state
9881            .lock()
9882            .map_err(|_| "MoE input trace writer lock is poisoned")?;
9883        if state.is_none() {
9884            let dir = std::path::PathBuf::from(&dir);
9885            std::fs::create_dir_all(&dir)?;
9886            let index = std::fs::OpenOptions::new()
9887                .create(true)
9888                .append(true)
9889                .open(dir.join("index.jsonl"))?;
9890            *state = Some(MoeInputTraceWriter {
9891                dir,
9892                index,
9893                payloads: std::collections::HashMap::new(),
9894            });
9895        }
9896        let writer = state.as_mut().unwrap();
9897        if writer.dir != std::path::Path::new(&dir) {
9898            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
9899        }
9900        let file_name = format!("layer-{il:03}.f32");
9901        if !writer.payloads.contains_key(&il) {
9902            let payload = std::fs::OpenOptions::new()
9903                .create(true)
9904                .append(true)
9905                .open(writer.dir.join(&file_name))?;
9906            let offset = payload.metadata()?.len();
9907            writer.payloads.insert(il, (payload, offset));
9908        }
9909        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
9910        let row_offset = *offset;
9911        payload.write_all(bytes)?;
9912        *offset += bytes.len() as u64;
9913        writeln!(
9914            writer.index,
9915            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
9916             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
9917             \"payload_bytes\":{}}}",
9918            bytes.len()
9919        )?;
9920        Ok(())
9921    }
9922
9923    #[allow(clippy::too_many_arguments)]
9924    #[allow(clippy::too_many_arguments)]
9925    // allow: the parameter list mirrors its moe_ffn_inner caller's dispatch contract
9926    pub(crate) fn moe_ffn_sequential_zq8(
9927        e: &Engine,
9928        m: &MoeWeights,
9929        z: &CudaSlice<f32>,
9930        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
9931        t: usize,
9932        cfg: &ModelConfig,
9933        il: u16,
9934        max_block: usize,
9935        vrows: bool,
9936    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9937        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
9938        let moe = cfg.moe.as_ref().unwrap();
9939        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
9940        let n_expert = moe.expert_count as usize; // 256
9941        let n_used = moe.expert_used_count as usize; // 8
9942        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
9943
9944        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
9945        debug_assert_eq!(m.gate_exps.in_f, n_embd);
9946        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
9947        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
9948        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
9949        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
9950
9951        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
9952        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
9953        let lim_exp = cfg.clamp_exp_at(il as u32);
9954        let lim_shexp = cfg.clamp_shexp_at(il as u32);
9955        let use_cache = Engine::moe_cache_enabled();
9956        let uniform_experts = m.has_uniform_expert_layout();
9957        let moe_q8 = uniform_experts && moe_q8_enabled_for_model(cfg, m);
9958        // Experimental secondary backend: complete experts already resident in the SLRU stay on
9959        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
9960        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
9961        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
9962        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
9963        // commands and CI have no llama.cpp or OpenMP dependency.
9964        let cpu_expert_requested = crate::cpu_experts::configured();
9965        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
9966            return Err(std::io::Error::other(
9967                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
9968            )
9969            .into());
9970        }
9971        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
9972        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
9973        // Those backends are each deterministic but are different numeric configurations, so a
9974        // later prefill eviction can change greedy output. Freeze after the first real prefill;
9975        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
9976        // staging below and cannot change backend assignment.
9977        let freeze_cpu_residency = cpu_expert_requested
9978            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
9979        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
9980            .ok()
9981            .and_then(|value| value.parse::<usize>().ok())
9982            .is_some_and(|tokens| tokens > 0);
9983        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
9984            e.freeze_moe_cache();
9985        }
9986        let cache_frozen = use_cache && e.moe_cache_frozen();
9987        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
9988
9989        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
9990        // cannot change logits, selected expert ids, or routing weights.
9991        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9992        if let Some(sig) = cfg.sigmoid_router() {
9993            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
9994        }
9995
9996        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
9997        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
9998        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
9999        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
10000        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
10001        // per-token host stall that dominated the 35B decode wall after stages 1+2.
10002        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
10003        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
10004        // only difference is where sel/w/pointers are READ from (device instead of params).
10005        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
10006        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
10007        // Any non-resident layer falls through to host routing + the gdec/sequential path.
10008        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
10009        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
10010        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
10011        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
10012        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
10013        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
10014        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
10015        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
10016        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
10017        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
10018        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
10019        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
10020        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
10021        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
10022        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
10023        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
10024        // now rides the dev loop below (same kernels per token as decode); pairs serves real
10025        // prefill (t >= 16, where spec never verifies).
10026        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
10027        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
10028        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
10029        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
10030        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
10031        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
10032        // ride the macro-aware sequential/staged paths below or every expert output is off by
10033        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
10034        let no_exp_macros = m.gate_exps.macros.is_none()
10035            && m.up_exps.macros.is_none()
10036            && m.down_exps.macros.is_none();
10037        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
10038        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
10039        // so it cannot even see the per-layer limit.
10040        if cfg.sigmoid_router().is_none()
10041            && cfg.m3.is_none()
10042            && cfg.hy3.is_none()
10043            && !cfg.swiglu_clamped_at(il as u32)
10044            && no_exp_macros
10045            // > MOE_DEV_MAX_T, not >= PRIME_MIN_T: t==16 is a decode width under the
10046            // exact-16 tier and rides the dev per-token program (see MOE_DEV_MAX_T);
10047            // pairs serves real prefill from 17 up.
10048            && t > MOE_DEV_MAX_T
10049            && m.dev_exps.is_some()
10050            && moe_q8_enabled_for_model(cfg, m)
10051            && std::env::var("MEMRA_MOE_PAIRS")
10052                .map(|v| v != "0")
10053                .unwrap_or(true)
10054            && std::env::var("MEMRA_MOE_STATS").is_err()
10055        {
10056            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
10057        }
10058
10059        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
10060        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
10061        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
10062        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
10063        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
10064        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
10065        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
10066        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
10067        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
10068        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
10069        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
10070        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
10071        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
10072        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
10073        // Keyed off sigmoid_router() so arch #4 is denied by construction.
10074        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
10075        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
10076        let dev_ok = uniform_experts
10077            && cfg.sigmoid_router().is_none()
10078            && cfg.m3.is_none()
10079            && cfg.hy3.is_none()
10080            && !cfg.swiglu_clamped_at(il as u32);
10081        // Observation modes must route through the host-visible selection below. Otherwise a fully
10082        // resident layer returns through device dispatch before its trace/stats row is recorded,
10083        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
10084        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
10085            || std::env::var("MEMRA_MOE_TRACE").is_ok()
10086            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
10087            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
10088        if dev_ok
10089            && t <= MOE_DEV_MAX_T
10090            && m.dev_exps.is_some()
10091            && n_used <= 8
10092            && moe_dev_enabled()
10093            && !observe_routes
10094        {
10095            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
10096        }
10097        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
10098            let row_ok = e.with_moe_cache(max_block, |c, eng| {
10099                if moe_prewarm_enabled() {
10100                    c.prewarm_layer(il, m, eng)?;
10101                }
10102                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
10103            })?;
10104            if row_ok {
10105                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
10106            }
10107        }
10108
10109        // SLAB-LOCAL RESIDENT ARM bases, hoisted above the router (lane/glm5-moe-loc door D):
10110        // whether the layer can run the DEVICE vrows table build decides HOW it routes, and
10111        // that has to be settled before the router runs. Pure immutable pointer reads with no
10112        // side effects, so the hoist changes nothing for any other arm; the full rationale for
10113        // the arm itself is at the `slab_fused_may_fire` predicate below.
10114        let slab_local = m
10115            .dev_exps
10116            .as_ref()
10117            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
10118        let slab_bases = slab_local.map(|d| {
10119            use cudarc::driver::DevicePtr;
10120            let s = e.stream();
10121            let (pg, _g0) = d.gate.device_ptr(&s);
10122            let (pu, _g1) = d.up.device_ptr(&s);
10123            let (pd, _g2) = d.down.device_ptr(&s);
10124            (pg, pu, pd)
10125        });
10126        // DOOR D (`MEMRA_MOE_VROWS_DEV_TABLES`, default OFF): route WITHOUT the pinned sel/w
10127        // readback and build the pair's pointer/scale tables on device instead. On the serving
10128        // shape the host table build is the selection's ONLY consumer, and it costs a full
10129        // `cuStreamSynchronize` + 2 DtoH + 2 pageable HtoD + 2 host Vecs per MoE layer-call —
10130        // 42 device-wide drains, 84 DtoH and 84 HtoD per ship round (the decode-gap
10131        // attribution's "43 cuStreamSynchronize/token ... the per-layer router-admission sync
10132        // structure", and 44.6% of the unattributed 71.6 HtoD calls/token).
10133        //
10134        // The extra conjuncts beyond `vrows_fires` (asserted equal at the dispatch) are exactly
10135        // the host-visible consumers of `sel_all` between here and there, each of which would
10136        // silently read an empty selection: `moesd::record_host_routes`, `hidden_trace`,
10137        // `MEMRA_MOE_TRACE`/`MEMRA_MOE_STATS`/the other `observe_routes` modes. Plus
10138        // `sigmoid_router_enabled()`, because `MEMRA_SIG_ROUTER=0` is a full-logit HOST oracle
10139        // with no device selection to read. `promote_worker_h2d` needs no conjunct: it requires
10140        // t == 1 and this arm requires t >= 2. Any miss falls closed to the host readback.
10141        let vrows_dev = vrows
10142            && t >= 2
10143            && crate::moe_vrows_dev_tables_on()
10144            && slab_bases.is_some()
10145            && moe_q8
10146            && uniform_experts
10147            && n_used <= 8
10148            && cfg.sigmoid_router().is_some()
10149            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10150            && !cpu_hybrid
10151            && sigmoid_router_enabled()
10152            && !observe_routes
10153            && !memra_reference::hidden_trace::enabled()
10154            && !crate::moesd::capture_active();
10155        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
10156        // Door D adds a fourth arm returning the router's DEVICE sel/w with no readback.
10157        let mut sel_dev: Option<(CudaSlice<i32>, CudaSlice<f32>)> = None;
10158        let (sel_all, w_all, routed_cpu_input) = if vrows_dev {
10159            let (sf, route_norm) = cfg
10160                .sigmoid_router()
10161                .expect("vrows_dev carries cfg.sigmoid_router().is_some()");
10162            sel_dev = Some(e.moe_router_sigmoid_topk(
10163                &logits,
10164                t,
10165                n_expert,
10166                n_used,
10167                m.active_count(),
10168                &m.exp_probs_b_dev,
10169                &m.active_experts_dev,
10170                sf,
10171                route_norm,
10172            )?);
10173            crate::MOE_VROWS_ROUTER_SYNCS_AVOIDED
10174                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10175            (Vec::new(), Vec::new(), None)
10176        } else if let Some(sig) = cfg.sigmoid_router() {
10177            if cpu_hybrid {
10178                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
10179                    e,
10180                    &logits,
10181                    z,
10182                    t,
10183                    n_embd,
10184                    n_expert,
10185                    n_used,
10186                    m.exp_probs_b.as_deref(),
10187                    sig,
10188                    m.active_experts.as_deref(),
10189                )?;
10190                (sel, w, Some(input))
10191            } else {
10192                let (sel, w) =
10193                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
10194                (sel, w, None)
10195            }
10196        } else {
10197            let (sel, w) =
10198                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
10199            (sel, w, None)
10200        };
10201        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
10202        if memra_reference::hidden_trace::enabled() {
10203            memra_reference::hidden_trace::emit_last_row(
10204                "router",
10205                il as i64,
10206                t,
10207                n_expert,
10208                &e.dtoh(&logits)?,
10209            );
10210            let last = (t - 1) * n_used;
10211            let mut route = Vec::with_capacity(n_used * 2);
10212            for slot in 0..n_used {
10213                route.push(sel_all[last + slot] as f32);
10214                route.push(w_all[last + slot]);
10215            }
10216            memra_reference::hidden_trace::emit_last_row("route", il as i64, 1, n_used * 2, &route);
10217        }
10218
10219        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
10220        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
10221        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
10222        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
10223        Self::trace_moe_input(e, il, t, n_embd, z)?;
10224
10225        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
10226        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
10227        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
10228        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
10229        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
10230        // wait for each pending block, so later copies can overlap the earlier expert kernels while
10231        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
10232        // T=1; batched forwards can have token-local consumers still in flight between selections.
10233        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
10234        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
10235        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
10236        let worker_disk_prefetch =
10237            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
10238        let promote_worker_h2d =
10239            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
10240        if promote_worker_h2d {
10241            let mut selected_blocks = Vec::with_capacity(n_used * 3);
10242            for &ex in sel_all.iter().take(n_used) {
10243                let ex = ex as u16;
10244                selected_blocks.extend([
10245                    BlockId::new(il, PROJ_GATE, ex),
10246                    BlockId::new(il, PROJ_UP, ex),
10247                    BlockId::new(il, PROJ_DOWN, ex),
10248                ]);
10249            }
10250            for &ex in sel_all.iter().take(n_used) {
10251                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
10252            }
10253            e.with_moe_cache(max_block, |cache, eng| {
10254                cache.promote_worker_reads_at_safe_boundary(
10255                    &selected_blocks,
10256                    &selected_blocks,
10257                    eng,
10258                )?;
10259                Ok(())
10260            })?;
10261        }
10262
10263        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
10264        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
10265        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
10266            let mut cnt = vec![0u32; n_expert];
10267            for &s in sel_all.iter() {
10268                cnt[s as usize] += 1;
10269            }
10270            let total = sel_all.len() as f64;
10271            let mut h = 0.0f64;
10272            let mut active = 0usize;
10273            for &c in &cnt {
10274                if c > 0 {
10275                    active += 1;
10276                    let p = c as f64 / total;
10277                    h -= p * p.log2();
10278                }
10279            }
10280            let maxc = cnt.iter().copied().max().unwrap_or(0);
10281            println!(
10282                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
10283                il,
10284                t,
10285                sel_all.len(),
10286                active,
10287                n_expert,
10288                h,
10289                (n_expert as f64).log2(),
10290                total / active.max(1) as f64,
10291                maxc
10292            );
10293        }
10294
10295        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
10296        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
10297        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
10298        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
10299        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
10300        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
10301        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
10302        // zeroed-then-accumulated exactly as before (fallback).
10303        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
10304        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
10305        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
10306        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
10307        let gdec_may_fire = uniform_experts
10308            && use_cache
10309            && n_used <= 8
10310            && gdec_enabled()
10311            && !cfg.swiglu_clamped_at(il as u32);
10312        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
10313        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
10314        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
10315        // archs the slabs were uploaded but never read, and every expert went through the
10316        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
10317        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
10318        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
10319        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
10320        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
10321        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
10322        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
10323        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
10324        // strictly worse than staging); under PP-2 without the prime walker this admits
10325        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
10326        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
10327        // `slab_local` / `slab_bases` are bound ABOVE the router: door D
10328        // (`MEMRA_MOE_VROWS_DEV_TABLES`) has to pre-decide the routing arm, and this is the
10329        // predicate it needs. Pure immutable pointer reads, so the hoist is behaviour-neutral.
10330        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
10331        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
10332        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
10333        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
10334        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
10335        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
10336        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
10337        // all-resident tokens, staged loop for misses), which is a dispatch-class
10338        // comparison, not a provenance one.
10339        let slab_fused_may_fire = slab_bases.is_some()
10340            && n_used <= 8
10341            && gdec_enabled()
10342            && !cfg.swiglu_clamped_at(il as u32)
10343            && cfg.m3.is_none()
10344            && no_exp_macros
10345            && moe_q8;
10346        // FUSED MoE EPILOGUE (lane/glm53-epilogue 2026-08-28, MEMRA_MOE_FUSED_EPI default OFF).
10347        // The arm glm5_next is denied by every other predicate in this function. It runs the SAME
10348        // launch pair shape as gdec — one gate/up kernel, one down/FMA kernel per token — but
10349        // with the three things this arch actually needs, none of which any existing fused
10350        // epilogue has:
10351        //   * the SIGMOID noaux_tc router's sel/w (host-routed above; the fused softmax router
10352        //     `moe_router_topk` that pairs/dev use would pick different experts — the M3
10353        //     gate-MISMATCH 74602-vs-92 lesson);
10354        //   * the PRE-clamped SwiGLU epilogue `silu(min(g,l)) * clamp(u,±l)` — step35's POST form
10355        //     is a different, plausible-but-wrong program (`fused_post_limit`);
10356        //   * the per-expert NVFP4 `weight_scale_2` macro fold, gate/up through the kernel's
10357        //     gs/us and down through the routing weight, exactly as `ffn_act_lim` + `axpy_into`
10358        //     do it in the sequential loop.
10359        // UNLIKE gdec it does NOT require the layer to be already-resident: it ADMITS the
10360        // 3*n_used selected blocks through the same `dispatch_source` the sequential loop uses
10361        // (hit = no copy, miss = the identical H2D into a slot) and only then collects the fixed
10362        // slot addresses. The staged bytes are unchanged — the §B.3 provenance property — so the
10363        // arm engages at any miss rate instead of gdec's P(all resident). It needs the cache to
10364        // hold 3*n_used blocks at once; `moe_fused_epi_token_q8` returns false when it cannot and
10365        // the token falls through to the sequential loop below.
10366        // TWO PROVENANCES, ONE LAUNCH PATH (slab arm added 2026-08-28). The SLRU arm below keys
10367        // on `slab_local.is_none()`; the SLAB arm keys on the slab existing. They differ ONLY in
10368        // where the eight expert pointers come from and both call `moe_fused_epi_launch`, so the
10369        // macro fold, the clamp and the kernel pair cannot drift apart between them.
10370        //
10371        // The slab arm is not an optimization, it is the arm that matters. Full two-card expert
10372        // residency makes `dev_exps` present on every stage engine, which makes `slab_local`
10373        // `Some`, which under the original predicate DENIED the fused epilogue outright — the
10374        // measured A/B would have read 0 dispatches and looked like "no effect". The residency
10375        // config is the serving config now, so the slab provenance is the one the product runs.
10376        //
10377        // It is also the SIMPLER arm: a slab holds every expert by construction, so there is no
10378        // admission, no eviction, no pass-2 re-verification and no fall-through. The SLRU arm's
10379        // capacity floor and re-check exist only because admission can move a slot.
10380        let fused_epi_common = n_used <= 8
10381            && moe_q8
10382            && cfg.m3.is_none()
10383            && cfg.sigmoid_router().is_some()
10384            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10385            && moe_fused_epi_enabled();
10386        let fused_epi_may_fire = fused_epi_common
10387            && uniform_experts
10388            && use_cache
10389            && cache_dispatch
10390            && slab_local.is_none();
10391        let fused_epi_slab_may_fire = fused_epi_common && slab_bases.is_some();
10392        // VERIFY-ROWS BATCHED ROUTED-EXPERT ARM (lane/glm5-vrest, 2026-08-31; rides
10393        // `MEMRA_GLM5_VERIFY_BATCH` — only the verify walk's batched arm passes `vrows`).
10394        // ONE launch pair covers ALL t x n_used routed pairs (the fused-epilogue kernels'
10395        // verify-rows twins) instead of the per-(token,expert) loop's ~49 launches per
10396        // token-layer — the flip-reprice cell-2 vrest wall (9.46 ms/row marginal at K=3).
10397        // Bit identity per row vs the sequential chain is the bar and it is structural:
10398        // routing is the SAME host invocation above; per-pair dots are qmatvec_expert_q8's
10399        // g-strided order; the epilogue is swiglu_preclamped_mul_scaled_f32's expression
10400        // with the per-expert macro fold exactly where ffn_act_lim/axpy_into fold it; the
10401        // down accumulation is the slot-ordered __fmaf_rn chain (the gdec-gated class).
10402        // Gated by glm5_verify_batch_gpu (kernel pair vs sequential chain + swapped-pair
10403        // and dropped-macro reds) and the glm5_tparallel_verify_gpu walk battery on the
10404        // NVFP4+macro serving expert class. Fail-closed: any unqualified shape falls
10405        // through to the unchanged loop below. Same slab-only scope as the fused epilogue
10406        // (the serving config's provenance); n_used<=8 mirrors its cap.
10407        let vrows_fires = vrows
10408            && t >= 2
10409            && slab_bases.is_some()
10410            && moe_q8
10411            && uniform_experts
10412            && n_used <= 8
10413            && cfg.sigmoid_router().is_some()
10414            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10415            && !cpu_hybrid;
10416        // moe_out memset elision: EVERY full-row-overwrite arm (gdec, slab fused, fused epilogue,
10417        // verify-rows) allocates uninit; a token that falls through to any accumulating loop
10418        // zeroes its own row. The fused epilogue's `moe_down8_fma_q8` fully overwrites `dst[o]`,
10419        // same as gdec's; `moe_down8_fma_q8_rows` fully overwrites every row.
10420        let mut moe_out = if gdec_may_fire
10421            || slab_fused_may_fire
10422            || fused_epi_may_fire
10423            || fused_epi_slab_may_fire
10424            || vrows_fires
10425        {
10426            e.uninit(t * n_embd)?
10427        } else {
10428            e.zeros(t * n_embd)?
10429        };
10430        // The router readback above already established a host boundary. Copy each small-t hidden
10431        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
10432        let cpu_input = if cpu_hybrid {
10433            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
10434        } else {
10435            None
10436        };
10437
10438        // Door D's predicate was evaluated at the router, BEFORE `vrows_fires` existed. If the
10439        // two ever disagreed, the layer would have skipped its readback and then dispatched a
10440        // per-row arm holding an EMPTY host selection — a silent wrong-answer class. The two
10441        // predicates share every conjunct by construction; assert it rather than hope.
10442        if vrows_dev && !vrows_fires {
10443            return Err(
10444                "MEMRA_MOE_VROWS_DEV_TABLES routed device-only but the verify-rows arm did not \
10445                 fire: the door-D and vrows_fires predicates disagree"
10446                    .into(),
10447            );
10448        }
10449        if vrows_fires {
10450            let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10451                return Err(
10452                    "verify-rows MoE arm fired without a live PRE clamp: the predicate and \
10453                     the dispatch disagree"
10454                        .into(),
10455                );
10456            };
10457            let bases = slab_bases.expect("vrows_fires carries slab_bases.is_some()");
10458            let sel = match sel_dev.as_ref() {
10459                Some((si, sw)) => VrowsSel::Dev(si, sw),
10460                None => VrowsSel::Host(&sel_all, &w_all),
10461            };
10462            Self::moe_vrows_pairs_q8(
10463                e,
10464                m,
10465                z,
10466                sel,
10467                il,
10468                bases,
10469                t,
10470                n_embd,
10471                n_ff_exp,
10472                n_used,
10473                limit,
10474                &mut moe_out,
10475            )?;
10476            if memra_reference::hidden_trace::enabled() {
10477                memra_reference::hidden_trace::emit_last_row(
10478                    "routed",
10479                    il as i64,
10480                    t,
10481                    n_embd,
10482                    &e.dtoh(&moe_out)?,
10483                );
10484            }
10485            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
10486            return Ok(moe_out);
10487        }
10488
10489        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
10490        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
10491        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
10492        // measured ~123 memsets/token of the decode wall).
10493        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
10494        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
10495        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
10496        let mut scratch_g: Option<CudaSlice<u8>> = None;
10497        let mut scratch_u: Option<CudaSlice<u8>> = None;
10498        let mut scratch_d: Option<CudaSlice<u8>> = None;
10499        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
10500        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
10501
10502        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
10503        // the copy stream before launching the current expert's compute. Pending slots stay invisible
10504        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
10505        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
10506        let page_window = moe_page_prefetch_window();
10507
10508        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
10509        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
10510        for tok in 0..t {
10511            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
10512            let w = &w_all[tok * n_used..(tok + 1) * n_used];
10513            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
10514            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10515
10516            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
10517            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
10518            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
10519            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
10520            // memcpy, zero admission, so no slot can move under the collected pointers) — any
10521            // miss falls through to the sequential loop below, which admits as before. In steady
10522            // state on a fully-resident rig every token-layer takes the grouped path.
10523            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
10524            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
10525            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
10526            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
10527            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
10528            // per-expert macro-scales the fused kernels don't fold — those fall through too.
10529            let no_macros = m.gate_exps.macros.is_none()
10530                && m.up_exps.macros.is_none()
10531                && m.down_exps.macros.is_none();
10532            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
10533            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
10534            // with pointers computed from the resident slab base + ex*stride instead of
10535            // collected SLRU slot addresses. No cache lock, no residency predicate — the
10536            // slab holds every expert by construction, so this arm never falls through
10537            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
10538            // staging both die). Bit-identity class: pointer provenance only, the same
10539            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
10540            // slab exists it is strictly better (no lock, no miss).
10541            if slab_fused_may_fire {
10542                let (pg, pu, pd) = slab_bases.unwrap();
10543                let mut gp = [0u64; 8];
10544                let mut up = [0u64; 8];
10545                let mut dp = [0u64; 8];
10546                for (j, &ex) in sel.iter().enumerate() {
10547                    let ex = ex as usize;
10548                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10549                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10550                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10551                }
10552                let mut wv = [0f32; 8];
10553                wv[..n_used].copy_from_slice(w);
10554                if tok_q8.is_none() {
10555                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10556                }
10557                let (zq, zd) = tok_q8.as_ref().unwrap();
10558                let act = e.moe_gate_up_silu8_q8(
10559                    crate::WPtr8(gp),
10560                    crate::WPtr8(up),
10561                    zq,
10562                    zd,
10563                    n_embd,
10564                    n_ff_exp,
10565                    n_used,
10566                    m.gate_exps.qtype,
10567                    m.up_exps.qtype,
10568                    m.gate_exps.row_bytes,
10569                    m.up_exps.row_bytes,
10570                )?;
10571                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
10572                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10573                e.moe_down8_fma_q8(
10574                    crate::WPtr8(dp),
10575                    crate::F32x8(wv),
10576                    &aq2,
10577                    &ad2,
10578                    &mut dst,
10579                    n_ff_exp,
10580                    n_embd,
10581                    n_used,
10582                    m.down_exps.qtype,
10583                    m.down_exps.row_bytes,
10584                )?;
10585                continue;
10586            }
10587            // FUSED MoE EPILOGUE, SLAB PROVENANCE. Ordered first: when a local slab exists it is
10588            // strictly better than anything the SLRU can offer — every expert is present by
10589            // construction, so there is no admission, no eviction and no fall-through. This is
10590            // the arm the two-card residency serving config actually runs.
10591            if fused_epi_slab_may_fire {
10592                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10593                    return Err(
10594                        "fused MoE epilogue (slab) fired without a live PRE clamp: the \
10595                                predicate and the dispatch disagree"
10596                            .into(),
10597                    );
10598                };
10599                let (pg, pu, pd) = slab_bases.unwrap();
10600                let mut g = [0u64; 8];
10601                let mut u = [0u64; 8];
10602                let mut d = [0u64; 8];
10603                for (j, &ex) in sel.iter().enumerate() {
10604                    let ex = ex as usize;
10605                    g[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10606                    u[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10607                    d[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10608                }
10609                if tok_q8.is_none() {
10610                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10611                }
10612                let (zq, zd) = tok_q8.as_ref().unwrap();
10613                Self::moe_fused_epi_launch(
10614                    e,
10615                    m,
10616                    zq,
10617                    zd,
10618                    sel,
10619                    w,
10620                    g,
10621                    u,
10622                    d,
10623                    &mut moe_out,
10624                    tok,
10625                    n_embd,
10626                    n_ff_exp,
10627                    n_used,
10628                    limit,
10629                )?;
10630                continue;
10631            }
10632            // FUSED MoE EPILOGUE, SLRU PROVENANCE. Ordered before gdec (which this arch never
10633            // reaches anyway: `gdec_may_fire` carries `!swiglu_clamped_at`). A `false` return
10634            // means the cache could not hold 3*n_used blocks at once — the token falls through to
10635            // the sequential loop, which zeroes its own row below.
10636            if fused_epi_may_fire {
10637                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10638                    return Err(
10639                        "fused MoE epilogue fired without a live PRE clamp: the predicate and \
10640                         the dispatch disagree"
10641                            .into(),
10642                    );
10643                };
10644                if tok_q8.is_none() {
10645                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10646                }
10647                let (zq, zd) = tok_q8.as_ref().unwrap();
10648                if Self::moe_fused_epi_token_q8(
10649                    e,
10650                    m,
10651                    il,
10652                    max_block,
10653                    zq,
10654                    zd,
10655                    sel,
10656                    w,
10657                    &mut moe_out,
10658                    tok,
10659                    n_embd,
10660                    n_ff_exp,
10661                    n_used,
10662                    limit,
10663                )? {
10664                    continue;
10665                }
10666            }
10667            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
10668                if tok_q8.is_none() {
10669                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10670                }
10671                let (zq, zd) = tok_q8.as_ref().unwrap();
10672                if Self::moe_gdec_token_q8(
10673                    e,
10674                    m,
10675                    il,
10676                    max_block,
10677                    zq,
10678                    zd,
10679                    sel,
10680                    w,
10681                    &mut moe_out,
10682                    tok,
10683                    n_embd,
10684                    n_ff_exp,
10685                    n_used,
10686                )? {
10687                    continue;
10688                }
10689            } else if gdec_may_fire
10690                && cfg.m3.is_none()
10691                && no_macros
10692                && Self::moe_gdec_token(
10693                    e,
10694                    m,
10695                    il,
10696                    max_block,
10697                    &zt,
10698                    sel,
10699                    w,
10700                    &mut moe_out,
10701                    tok,
10702                    n_embd,
10703                    n_ff_exp,
10704                    n_used,
10705                )?
10706            {
10707                continue;
10708            }
10709
10710            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
10711            // slab pair could fire. This token fell through to a sequential axpy loop, which
10712            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
10713            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
10714            // has no fallible predicate), included for the allocation invariant's symmetry.
10715            if gdec_may_fire || slab_fused_may_fire || fused_epi_may_fire || fused_epi_slab_may_fire
10716            {
10717                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10718                e.memset_zeros_view(&mut row)?;
10719            }
10720
10721            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
10722            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
10723            // stall this path exists to remove, while mixing projections would require another
10724            // activation round-trip. Weight addresses remain valid until this worker is joined at
10725            // the bottom of the token scope.
10726            let mut cpu_mask = vec![false; sel.len()];
10727            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
10728                let gpu_resident = if use_cache {
10729                    e.with_moe_cache(max_block, |cache, _| {
10730                        Ok(sel
10731                            .iter()
10732                            .map(|&expert| {
10733                                let expert = expert as u16;
10734                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
10735                                    .into_iter()
10736                                    .filter(|&projection| {
10737                                        cache
10738                                            .resident(BlockId::new(il, projection, expert))
10739                                            .is_some()
10740                                    })
10741                                    .count()
10742                            })
10743                            .collect::<Vec<_>>())
10744                    })?
10745                } else {
10746                    vec![0; sel.len()]
10747                };
10748                let mut cpu_selected = Vec::new();
10749                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
10750                    if gpu_resident[index] != 3 {
10751                        cpu_mask[index] = true;
10752                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
10753                        let expert = expert as usize;
10754                        cpu_selected.push((expert, route_weight));
10755                    }
10756                }
10757                if crate::cpu_experts::predictor_enabled() {
10758                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
10759                    // from this layer's MoE input and prefetches predicted-and-missing
10760                    // experts into the companion RAM cache. Never blocks this thread.
10761                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10762                    crate::cpu_experts::predictor_submit(il, row);
10763                }
10764                if cpu_selected.is_empty() {
10765                    None
10766                } else {
10767                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10768                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
10769                        .map_err(std::io::Error::other)?;
10770                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
10771                }
10772            } else {
10773                None
10774            };
10775
10776            let worker_window = worker_disk_prefetch
10777                .then(worker_prefetch_window)
10778                .unwrap_or(0);
10779            for (j, &ex) in sel.iter().enumerate() {
10780                if cpu_mask[j] {
10781                    continue;
10782                }
10783                let ex = ex as usize;
10784                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
10785                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
10786                // fused form) and macro-carrying artifacts — still have their bytes in the
10787                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
10788                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
10789                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
10790                if let Some(d) = slab_local {
10791                    let gl = m.gate_exps.expert_layout(ex);
10792                    let ul = m.up_exps.expert_layout(ex);
10793                    let dl = m.down_exps.expert_layout(ex);
10794                    let (g0, u0, d0) = (
10795                        ex * m.gate_exps.expert_stride,
10796                        ex * m.up_exps.expert_stride,
10797                        ex * m.down_exps.expert_stride,
10798                    );
10799                    let (gate, up) = if moe_q8 {
10800                        if tok_q8.is_none() {
10801                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10802                        }
10803                        let (zq, zd) = tok_q8.as_ref().unwrap();
10804                        (
10805                            e.qmatvec_expert_q8(
10806                                &d.gate,
10807                                g0..g0 + gl.len,
10808                                zq,
10809                                zd,
10810                                1,
10811                                m.gate_exps.in_f,
10812                                m.gate_exps.out_f,
10813                                gl.qtype,
10814                                gl.row_bytes,
10815                            )?,
10816                            e.qmatvec_expert_q8(
10817                                &d.up,
10818                                u0..u0 + ul.len,
10819                                zq,
10820                                zd,
10821                                1,
10822                                m.up_exps.in_f,
10823                                m.up_exps.out_f,
10824                                ul.qtype,
10825                                ul.row_bytes,
10826                            )?,
10827                        )
10828                    } else {
10829                        (
10830                            m.qmatvec_view(
10831                                e,
10832                                &d.gate,
10833                                g0..g0 + gl.len,
10834                                &zt,
10835                                1,
10836                                m.gate_exps.in_f,
10837                                m.gate_exps.out_f,
10838                                gl.qtype,
10839                                gl.row_bytes,
10840                            )?,
10841                            m.qmatvec_view(
10842                                e,
10843                                &d.up,
10844                                u0..u0 + ul.len,
10845                                &zt,
10846                                1,
10847                                m.up_exps.in_f,
10848                                m.up_exps.out_f,
10849                                ul.qtype,
10850                                ul.row_bytes,
10851                            )?,
10852                        )
10853                    };
10854                    let mut act = e.uninit(n_ff_exp)?;
10855                    Self::ffn_act_lim(
10856                        e,
10857                        cfg,
10858                        &gate,
10859                        &up,
10860                        m.gate_exps.macro_scale(ex),
10861                        m.up_exps.macro_scale(ex),
10862                        lim_exp,
10863                        &mut act,
10864                        n_ff_exp,
10865                    )?;
10866                    let y = if moe_q8 {
10867                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10868                        e.qmatvec_expert_q8(
10869                            &d.down,
10870                            d0..d0 + dl.len,
10871                            &aq2,
10872                            &ad2,
10873                            1,
10874                            m.down_exps.in_f,
10875                            m.down_exps.out_f,
10876                            dl.qtype,
10877                            dl.row_bytes,
10878                        )?
10879                    } else {
10880                        let actv = act.slice(0..n_ff_exp);
10881                        m.qmatvec_view(
10882                            e,
10883                            &d.down,
10884                            d0..d0 + dl.len,
10885                            &actv,
10886                            1,
10887                            m.down_exps.in_f,
10888                            m.down_exps.out_f,
10889                            dl.qtype,
10890                            dl.row_bytes,
10891                        )?
10892                    };
10893                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10894                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10895                    continue;
10896                }
10897                for next in page_prefetch_positions(j, sel.len(), page_window) {
10898                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
10899                }
10900                let keep = [
10901                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
10902                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
10903                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
10904                ];
10905                if worker_disk_prefetch && worker_window > 0 {
10906                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
10907                        Self::moe_prefetch_disk_expert(
10908                            e,
10909                            il,
10910                            sel[next] as usize,
10911                            m,
10912                            max_block,
10913                            &keep,
10914                        )?;
10915                    }
10916                } else if cache_dispatch
10917                    && !cpu_hybrid
10918                    && moe_prefetch_enabled()
10919                    && j + 1 < sel.len()
10920                {
10921                    let next = sel[j + 1] as usize;
10922                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
10923                }
10924                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
10925                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
10926                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
10927                    // layouts stay on the metadata-aware f32 path.
10928                    if (gate_q8 || up_q8) && tok_q8.is_none() {
10929                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10930                    }
10931                    let gate = if gate_q8 {
10932                        let (zq, zd) = tok_q8.as_ref().unwrap();
10933                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
10934                    } else {
10935                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
10936                    };
10937                    let up = if up_q8 {
10938                        let (zq, zd) = tok_q8.as_ref().unwrap();
10939                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
10940                    } else {
10941                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
10942                    };
10943                    let mut act = e.uninit(n_ff_exp)?;
10944                    Self::ffn_act_lim(
10945                        e,
10946                        cfg,
10947                        &gate,
10948                        &up,
10949                        m.gate_exps.macro_scale(ex),
10950                        m.up_exps.macro_scale(ex),
10951                        lim_exp,
10952                        &mut act,
10953                        n_ff_exp,
10954                    )?;
10955                    let y = if down_q8 {
10956                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10957                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
10958                    } else {
10959                        let actv = act.slice(0..n_ff_exp);
10960                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
10961                    };
10962                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10963                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
10964                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10965                } else if cache_dispatch {
10966                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
10967                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
10968                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
10969                    // only difference between HIT and MISS is whether the memcpy_htod ran.
10970                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
10971                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
10972                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
10973                    Self::ffn_act_lim(
10974                        e,
10975                        cfg,
10976                        &gate,
10977                        &up,
10978                        m.gate_exps.macro_scale(ex),
10979                        m.up_exps.macro_scale(ex),
10980                        lim_exp,
10981                        &mut act,
10982                        n_ff_exp,
10983                    )?;
10984                    let actv = act.slice(0..n_ff_exp);
10985                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
10986                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10987                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
10988                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10989                } else if cache_frozen {
10990                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
10991                    // first prime. Reuse every fixed resident projection directly and stage only a
10992                    // true miss through the ordinary scratch slot. This preserves the established
10993                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
10994                    let gate = Self::moe_frozen_gemm(
10995                        e,
10996                        il,
10997                        PROJ_GATE,
10998                        ex,
10999                        m,
11000                        max_block,
11001                        &zt,
11002                        &mut scratch_g,
11003                        g_len,
11004                    )?;
11005                    let up = Self::moe_frozen_gemm(
11006                        e,
11007                        il,
11008                        PROJ_UP,
11009                        ex,
11010                        m,
11011                        max_block,
11012                        &zt,
11013                        &mut scratch_u,
11014                        u_len,
11015                    )?;
11016                    let mut act = e.uninit(n_ff_exp)?;
11017                    Self::ffn_act_lim(
11018                        e,
11019                        cfg,
11020                        &gate,
11021                        &up,
11022                        m.gate_exps.macro_scale(ex),
11023                        m.up_exps.macro_scale(ex),
11024                        lim_exp,
11025                        &mut act,
11026                        n_ff_exp,
11027                    )?;
11028                    let actv = act.slice(0..n_ff_exp);
11029                    let y = Self::moe_frozen_gemm(
11030                        e,
11031                        il,
11032                        PROJ_DOWN,
11033                        ex,
11034                        m,
11035                        max_block,
11036                        &actv,
11037                        &mut scratch_d,
11038                        d_len,
11039                    )?;
11040                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11041                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
11042                } else {
11043                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
11044                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
11045                    // fully overwrites the byte range the GEMM reads).
11046                    if scratch_g.is_none() {
11047                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
11048                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
11049                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
11050                    }
11051                    let (sg, su, sd) = (
11052                        scratch_g.as_mut().unwrap(),
11053                        scratch_u.as_mut().unwrap(),
11054                        scratch_d.as_mut().unwrap(),
11055                    );
11056                    let gl = m.gate_exps.expert_layout(ex);
11057                    let ul = m.up_exps.expert_layout(ex);
11058                    let dl = m.down_exps.expert_layout(ex);
11059                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
11060                    let gate = m.qmatvec_view(
11061                        e,
11062                        sg,
11063                        0..gl.len,
11064                        &zt,
11065                        1,
11066                        m.gate_exps.in_f,
11067                        m.gate_exps.out_f,
11068                        gl.qtype,
11069                        gl.row_bytes,
11070                    )?;
11071
11072                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
11073                    let up = m.qmatvec_view(
11074                        e,
11075                        su,
11076                        0..ul.len,
11077                        &zt,
11078                        1,
11079                        m.up_exps.in_f,
11080                        m.up_exps.out_f,
11081                        ul.qtype,
11082                        ul.row_bytes,
11083                    )?;
11084
11085                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
11086                    Self::ffn_act_lim(
11087                        e,
11088                        cfg,
11089                        &gate,
11090                        &up,
11091                        m.gate_exps.macro_scale(ex),
11092                        m.up_exps.macro_scale(ex),
11093                        lim_exp,
11094                        &mut act,
11095                        n_ff_exp,
11096                    )?;
11097
11098                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
11099                    let actv = act.slice(0..n_ff_exp);
11100                    let y = m.qmatvec_view(
11101                        e,
11102                        sd,
11103                        0..dl.len,
11104                        &actv,
11105                        1,
11106                        m.down_exps.in_f,
11107                        m.down_exps.out_f,
11108                        dl.qtype,
11109                        dl.row_bytes,
11110                    )?;
11111
11112                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11113                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
11114                }
11115            }
11116            if let Some(worker) = cpu_worker {
11117                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
11118                let cpu_output = e.htod(&cpu_output)?;
11119                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11120                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
11121            }
11122            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
11123                for (j, &ex) in sel.iter().enumerate() {
11124                    if cpu_mask[j] {
11125                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
11126                    }
11127                }
11128            }
11129        }
11130
11131        if memra_reference::hidden_trace::enabled() {
11132            memra_reference::hidden_trace::emit_last_row(
11133                "routed",
11134                il as i64,
11135                t,
11136                n_embd,
11137                &e.dtoh(&moe_out)?,
11138            );
11139        }
11140
11141        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11142
11143        Ok(moe_out)
11144    }
11145
11146    /// The glm5 TP-2 EP walk (`MEMRA_GLM5_TP`): one MoE layer's routed-expert FFN over
11147    /// whole-expert contiguous halves. The ROUTER is the unchanged root-side program
11148    /// (`moe_router_logits` + `moe_route_sigmoid_cfg` — bit-identical selection by
11149    /// construction); each rank computes its owned slots' UNWEIGHTED expert rows through
11150    /// the sequential per-expert program (gate/up qmatvec + `ffn_act_lim` + down qmatvec —
11151    /// per-expert-independent dots), the peer's rows return host-canonically, and root
11152    /// applies the slot-ordered `axpy` accumulation chain — the same rounded-operation
11153    /// sequence the plain sequential walk applies. The ROOT-owned shared expert then adds
11154    /// through the extracted `moe_shexp_add`, verbatim.
11155    ///
11156    /// Three transport arms behind ONE routing (lane/glm5-ep-diet; sel/w are shared so a
11157    /// dispatch change can never change selection):
11158    ///   * `MEMRA_GLM5_EP_GROUPED_PRIME` (prefill shapes only): per-rank grouped-GEMM prime
11159    ///     over the rank slabs — the plain grouped-prefill program split by ownership.
11160    ///     Falls closed to the arms below whenever the plain arm's conjuncts do not hold.
11161    ///   * `MEMRA_GLM5_EP_DIET`: the v1 walk's kernels and combine chain with dieted data
11162    ///     movement — one bulk fan-out, zero per-slot host round-trips, one combine launch.
11163    ///     Decode-byte-identical to v1 by construction.
11164    ///   * default: the v1 per-slot host-canonical walk, byte-for-byte.
11165    #[allow(clippy::too_many_arguments)]
11166    fn moe_ffn_glm5_ep(
11167        e: &Engine,
11168        m: &MoeWeights,
11169        ep: &crate::glm5_tp::Glm5EpExps,
11170        z: &CudaSlice<f32>,
11171        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11172        t: usize,
11173        cfg: &ModelConfig,
11174        il: u16,
11175        prefill: bool,
11176    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11177        let moe = cfg
11178            .moe
11179            .as_ref()
11180            .ok_or("glm5 EP execution requires MoE model metadata")?;
11181        let n_embd = cfg.n_embd as usize;
11182        let n_expert = moe.expert_count as usize;
11183        let n_used = moe.expert_used_count as usize;
11184        let n_ff_exp = moe.expert_ff_length as usize;
11185        let sig = cfg
11186            .sigmoid_router()
11187            .ok_or("glm5 EP execution requires the sigmoid router")?;
11188        let lim_exp = cfg.clamp_exp_at(il as u32);
11189        let lim_shexp = cfg.clamp_shexp_at(il as u32);
11190        if ep.slabs.iter().map(|s| s.n_experts).sum::<usize>() != n_expert {
11191            return Err(format!(
11192                "glm5 EP slabs cover {:?} experts, model declares {n_expert}",
11193                ep.slabs.iter().map(|s| s.n_experts).collect::<Vec<_>>()
11194            )
11195            .into());
11196        }
11197        let rt = &ep.rt;
11198        let ranks = ep.ranks();
11199
11200        // Root router, unchanged program (selection bit-identical to the sequential arm).
11201        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
11202        let (sel_all, w_all) =
11203            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
11204        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
11205        // The route trace taps ride this walk too (sel/w are already host-side here);
11206        // the EP walk must never be a blind spot for the co-activation measurement.
11207        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
11208
11209        // EP grouped prime (`MEMRA_GLM5_EP_GROUPED_PRIME`, default OFF): keyed exactly like
11210        // the plain grouped-prefill arm (`prefill` + t above the per-token tier) and honoring
11211        // the family rollback (`MEMRA_MOE_GROUPED_PREFILL=0` kills it too). The announce
11212        // prints once per process PER FLAG VALUE, in both arms, so an A/B grep distinguishes
11213        // engagement without the line being an arm-local cost.
11214        if prefill && t > MOE_DEV_MAX_T {
11215            static EPGP_ANNOUNCED: std::sync::atomic::AtomicU8 =
11216                std::sync::atomic::AtomicU8::new(0);
11217            let enabled = crate::ep_grouped_prime_on() && moe_grouped_prefill_enabled();
11218            let bit = 1u8 << u8::from(enabled);
11219            if EPGP_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
11220                eprintln!(
11221                    "[glm5-ep-grouped-prime] flag={} t={t} il={il} (announce printed in both \
11222                     arms; engagement is the dispatch counter + per-layer execute line)",
11223                    if enabled { "on" } else { "off" },
11224                );
11225            }
11226            if enabled
11227                && let Some(mut out) =
11228                    Self::moe_ffn_glm5_ep_grouped_prime(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?
11229            {
11230                Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11231                return Ok(out);
11232            }
11233        }
11234
11235        // spec x TP composition receipt (the #80 review's confirmed perf-shape finding):
11236        // at verify widths (1 < t < prefill) the EP walk PREEMPTS the batched vrows MoE
11237        // pair — a composed verify round pays the sequential per-(token,expert) walk per
11238        // MoE layer, re-inheriting the vrest wall the vrows lane removed on the unsharded
11239        // shape. Announced once so a composed-shape battery cannot mistake a flat
11240        // MOE_VROWS_DISPATCHES counter for a wiring bug; the EP-aware vrows arm is the
11241        // named lever (composition-20260901/box/CELLS.md).
11242        if t > 1 && !prefill {
11243            static EP_VERIFY_MARKED: std::sync::atomic::AtomicBool =
11244                std::sync::atomic::AtomicBool::new(false);
11245            if !EP_VERIFY_MARKED.swap(true, std::sync::atomic::Ordering::Relaxed) {
11246                eprintln!(
11247                    "[glm5-tp-ep] verify rows ride the SEQUENTIAL EP walk (t={t}): the \
11248                     batched vrows MoE pair is preempted by EP; the EP-aware vrows arm is \
11249                     the named lever performance_claim=false"
11250                );
11251            }
11252        }
11253        // EP dispatch diet (`MEMRA_GLM5_EP_DIET`, default OFF): same kernels, same combine
11254        // chain, restructured movement. Read per call — `=0`/unset restores the v1 walk.
11255        if crate::ep_diet_on() {
11256            let mut out = Self::moe_ffn_glm5_ep_diet(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?;
11257            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11258            return Ok(out);
11259        }
11260
11261        let mut moe_out = e.zeros(t * n_embd)?;
11262        use crate::tp_transport::TpTransport as TpXport;
11263        let hop = ep.rt.hop(e);
11264        // Peer replicas of `z`, one arm each (lane/glm5-tp-transport):
11265        //   host-canonical — v1's EXACT pattern: one draining `dtoh` of the whole block here,
11266        //     then one row `htod` per token PER PEER RANK inside the loop (at two ranks that
11267        //     is byte- and hop-identical to v1). Preserved hop-for-hop so
11268        //     `MEMRA_GLM5_TP_TRANSPORT=0` reproduces the banked v1 walk, not a faster cousin.
11269        //   peer-pull — ONE device copy of the whole `[t, n_embd]` block per peer rank; rows
11270        //     are sliced out of it. Same bytes on the peers, the host uploads and drains
11271        //     removed.
11272        let z_host = match hop.transport {
11273            TpXport::HostCanonical => Some(crate::tp_transport::host_stage_block(
11274                &hop,
11275                0,
11276                z,
11277                t * n_embd,
11278            )?),
11279            TpXport::PeerPull => None,
11280        };
11281        let z_peer_bulks = match hop.transport {
11282            TpXport::PeerPull => Some(crate::tp_transport::fanout_f32(&hop, z, t * n_embd)?),
11283            TpXport::HostCanonical => None,
11284        };
11285        for tok in 0..t {
11286            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11287            let w = &w_all[tok * n_used..(tok + 1) * n_used];
11288            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11289            // The host-canonical arm materializes this token's peer rows here (v1's
11290            // per-token `htod`, one per peer rank); the peer-pull arm slices them out of the
11291            // bulk blocks. The holders must outlive the views, hence the two-step.
11292            let z_peer_row_holders: Option<Vec<CudaSlice<f32>>> = match &z_host {
11293                Some(h) => {
11294                    let mut rows = Vec::with_capacity(ranks - 1);
11295                    for r in 1..ranks {
11296                        rows.push(crate::tp_transport::host_row_to(
11297                            &hop,
11298                            r,
11299                            &h[tok * n_embd..(tok + 1) * n_embd],
11300                        )?);
11301                    }
11302                    Some(rows)
11303                }
11304                None => None,
11305            };
11306            // Per slot, in ROUTER SLOT ORDER: compute the UNWEIGHTED expert row on its
11307            // owner, then fmaf-accumulate on root — the plain walk's exact chain.
11308            for (j, &ex) in sel.iter().enumerate() {
11309                let ex = ex as usize;
11310                let owner = ep.owner(ex);
11311                if owner != 0 {
11312                    // Engagement counter FIRST (a red skip still counts as ROUTED).
11313                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES
11314                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11315                    // Gate red arm: dropping the peers' slots MUST diverge — the
11316                    // non-vacuity proof that the peer ranks contribute real expert work.
11317                    if matches!(
11318                        crate::glm5_tp::gate_red(),
11319                        Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11320                    ) {
11321                        continue;
11322                    }
11323                }
11324                let zin_holder;
11325                let (dev, slab, zin) = if owner == 0 {
11326                    (e, &ep.slabs[0], &zt)
11327                } else {
11328                    zin_holder = match (&z_peer_row_holders, &z_peer_bulks) {
11329                        (Some(rows), _) => rows[owner - 1].slice(0..n_embd),
11330                        (None, Some(bulks)) => {
11331                            bulks[owner - 1].slice(tok * n_embd..(tok + 1) * n_embd)
11332                        }
11333                        (None, None) => {
11334                            return Err(
11335                                "glm5 EP: neither transport arm staged the peer activation".into(),
11336                            );
11337                        }
11338                    };
11339                    (
11340                        crate::glm5_tp::rank_engine(e, rt, owner),
11341                        &ep.slabs[owner],
11342                        &zin_holder,
11343                    )
11344                };
11345                // Placement indirection: the owner's slab packs its experts in
11346                // ascending-id order; `local_of` is the slot (identical to
11347                // `ex - first_expert` under the even split).
11348                let local = ep.local_of[ex] as usize;
11349                let gl = m.gate_exps.expert_stride;
11350                let ul = m.up_exps.expert_stride;
11351                let dl = m.down_exps.expert_stride;
11352                let gate = dev.qmatvec_view(
11353                    &slab.gate,
11354                    local * gl..(local + 1) * gl,
11355                    zin,
11356                    1,
11357                    m.gate_exps.in_f,
11358                    m.gate_exps.out_f,
11359                    m.gate_exps.qtype,
11360                    m.gate_exps.row_bytes,
11361                )?;
11362                let up = dev.qmatvec_view(
11363                    &slab.up,
11364                    local * ul..(local + 1) * ul,
11365                    zin,
11366                    1,
11367                    m.up_exps.in_f,
11368                    m.up_exps.out_f,
11369                    m.up_exps.qtype,
11370                    m.up_exps.row_bytes,
11371                )?;
11372                let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
11373                Self::ffn_act_lim(
11374                    dev,
11375                    cfg,
11376                    &gate,
11377                    &up,
11378                    m.gate_exps.macro_scale(ex),
11379                    m.up_exps.macro_scale(ex),
11380                    lim_exp,
11381                    &mut act,
11382                    n_ff_exp,
11383                )?;
11384                let actv = act.slice(0..n_ff_exp);
11385                let y = dev.qmatvec_view(
11386                    &slab.down,
11387                    local * dl..(local + 1) * dl,
11388                    &actv,
11389                    1,
11390                    m.down_exps.in_f,
11391                    m.down_exps.out_f,
11392                    m.down_exps.qtype,
11393                    m.down_exps.row_bytes,
11394                )?;
11395                // The owner's row returns through the armed transport; the root row stays
11396                // put. The slot-ordered axpy below is the ONE cross-rank arithmetic site and
11397                // it reproduces the sequential walk's accumulate chain operation for
11398                // operation — unchanged by which transport delivered the row.
11399                let y_root = if owner == 0 {
11400                    y
11401                } else {
11402                    crate::tp_transport::return_row_to_root(&hop, owner, &y, n_embd)?
11403                };
11404                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11405                e.axpy_into(
11406                    &y_root,
11407                    w[j] * m.down_exps.macro_scale(ex),
11408                    &mut dst,
11409                    n_embd,
11410                )?;
11411            }
11412        }
11413
11414        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11415        Ok(moe_out)
11416    }
11417
11418    /// The DIETED glm5 EP walk (`MEMRA_GLM5_EP_DIET`, lane/glm5-ep-diet): the v1 walk's
11419    /// per-slot expert kernels and its exact slot-ordered combine chain, with the data
11420    /// movement restructured in whole groups (the tp2-battery's measured 13-18 ms/token v1
11421    /// join+dispatch tax, attributed to per-token host fan-out x42 layers + ~4-5 sync
11422    /// peer-slot round-trips/layer + interleaved host-blocked issue):
11423    ///
11424    ///   1. ONE bulk peer z fan-out per layer-call ([t, n_embd] in one upload; SKIPPED
11425    ///      entirely when the call routed no peer-owned expert — the placement-map
11426    ///      multiplier: a single-rank layer-call moves zero activation bytes off root).
11427    ///   2. Peer-owned rows compute back-to-back on the peer stream into a compact block
11428    ///      (issue order cannot change bytes: every row is an independent per-expert
11429    ///      program; the combine order below is fixed by the id table, not by issue).
11430    ///   3. Root-owned rows compute on the root stream, un-blocked by peer returns.
11431    ///   4. ONE bulk peer return (peer DtoH + root HtoD of the compact block) replaces the
11432    ///      per-slot round-trip dribble.
11433    ///   5. ONE `moe_pairs_scatter` launch applies the per-token slot-ordered fmaf chain —
11434    ///      the kernel header carries the byte-identity contract vs the zeros +
11435    ///      sequential-`axpy_f32` chain this replaces, and the weights are the SAME host
11436    ///      fold (`w * macro_scale(ex)`) v1 passed per launch.
11437    ///
11438    /// BYTE-IDENTICAL to the v1 walk (and to plain, wherever v1 is) by construction: same
11439    /// kernels over the same bytes; copies (dtod / bulk DtoH+HtoD) preserve bits; the one
11440    /// arithmetic site keeps its exact chain. Transport stays HOST-CANONICAL: the two bulk
11441    /// hops of steps 1 and 4 are the named native-P2P swap points for the box arc (the
11442    /// `MEMRA_STEP_TP_BULK_P2P` precedent: peer copies 61,452 -> ~21/layer on step; the
11443    /// glm5 seam inherits `configure_native_p2p` but does NOT wire it on the rig — the
11444    /// same-device dual-context emulation has no real peer transport to qualify).
11445    #[allow(clippy::too_many_arguments)]
11446    fn moe_ffn_glm5_ep_diet(
11447        e: &Engine,
11448        m: &MoeWeights,
11449        ep: &crate::glm5_tp::Glm5EpExps,
11450        z: &CudaSlice<f32>,
11451        sel_all: &[u32],
11452        w_all: &[f32],
11453        t: usize,
11454        cfg: &ModelConfig,
11455        il: u16,
11456    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11457        use std::sync::atomic::Ordering;
11458        let moe = cfg
11459            .moe
11460            .as_ref()
11461            .ok_or("glm5 EP execution requires MoE model metadata")?;
11462        let n_embd = cfg.n_embd as usize;
11463        let n_used = moe.expert_used_count as usize;
11464        let n_ff_exp = moe.expert_ff_length as usize;
11465        let lim_exp = cfg.clamp_exp_at(il as u32);
11466        let rt = &ep.rt;
11467        let n_pairs = t * n_used;
11468        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11469            return Err("glm5 EP diet geometry".into());
11470        }
11471        let red_skip_peer = matches!(
11472            crate::glm5_tp::gate_red(),
11473            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11474        );
11475
11476        // Slab-position table: root-owned pairs pack the slab head in pair order; each peer
11477        // rank's pairs pack a contiguous tail segment (so every rank's bulk return is ONE
11478        // contiguous upload). `ids[p]` is pair p's slab row; the scatter walks ids in slot
11479        // order per token, which is what pins the combine chain to v1's regardless of
11480        // packing. At two ranks this is byte-for-byte the original head/tail split.
11481        let ranks = ep.ranks();
11482        let mut per_rank = vec![0usize; ranks];
11483        for &s in sel_all.iter().take(n_pairs) {
11484            let ex = s as usize;
11485            if ex >= ep.owner_of.len() {
11486                return Err(format!("glm5 EP diet: selection {ex} outside the bank").into());
11487            }
11488            per_rank[ep.owner(ex)] += 1;
11489        }
11490        let mut base = vec![0usize; ranks];
11491        for r in 1..ranks {
11492            base[r] = base[r - 1] + per_rank[r - 1];
11493        }
11494        let mut ids = vec![0i32; n_pairs];
11495        {
11496            let mut k = vec![0usize; ranks];
11497            for (p, id) in ids.iter_mut().enumerate() {
11498                let r = ep.owner(sel_all[p] as usize);
11499                *id = (base[r] + k[r]) as i32;
11500                k[r] += 1;
11501            }
11502        }
11503
11504        crate::glm5_tp::GLM5_EP_DIET_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11505        for r in 1..ranks {
11506            crate::glm5_tp::GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.fetch_add(
11507                if per_rank[r] > 0 {
11508                    (t - 1) as u64
11509                } else {
11510                    t as u64
11511                },
11512                Ordering::Relaxed,
11513            );
11514        }
11515        static EP_DIET_MARKED: std::sync::atomic::AtomicBool =
11516            std::sync::atomic::AtomicBool::new(false);
11517        if !EP_DIET_MARKED.swap(true, Ordering::Relaxed) {
11518            eprintln!(
11519                "[glm5-ep-diet] engaged: bulk fan-out + compact peer staging + single \
11520                 slot-ordered scatter combine; per-slot host round-trips removed \
11521                 transport={} performance_claim=false",
11522                ep.rt.transport.name(),
11523            );
11524        }
11525
11526        // Per-slot expert program, shared verbatim with the v1 walk (same kernels, same
11527        // argument order): gate/up qmatvec + ffn_act_lim + down qmatvec on the OWNING rank.
11528        let expert_row = |dev: &Engine,
11529                          slab: &crate::glm5_tp::EpRankSlab,
11530                          zin: &cudarc::driver::CudaView<f32>,
11531                          ex: usize|
11532         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11533            let local = ep.local_of[ex] as usize;
11534            let gl = m.gate_exps.expert_stride;
11535            let ul = m.up_exps.expert_stride;
11536            let dl = m.down_exps.expert_stride;
11537            let gate = dev.qmatvec_view(
11538                &slab.gate,
11539                local * gl..(local + 1) * gl,
11540                zin,
11541                1,
11542                m.gate_exps.in_f,
11543                m.gate_exps.out_f,
11544                m.gate_exps.qtype,
11545                m.gate_exps.row_bytes,
11546            )?;
11547            let up = dev.qmatvec_view(
11548                &slab.up,
11549                local * ul..(local + 1) * ul,
11550                zin,
11551                1,
11552                m.up_exps.in_f,
11553                m.up_exps.out_f,
11554                m.up_exps.qtype,
11555                m.up_exps.row_bytes,
11556            )?;
11557            let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
11558            Self::ffn_act_lim(
11559                dev,
11560                cfg,
11561                &gate,
11562                &up,
11563                m.gate_exps.macro_scale(ex),
11564                m.up_exps.macro_scale(ex),
11565                lim_exp,
11566                &mut act,
11567                n_ff_exp,
11568            )?;
11569            let actv = act.slice(0..n_ff_exp);
11570            dev.qmatvec_view(
11571                &slab.down,
11572                local * dl..(local + 1) * dl,
11573                &actv,
11574                1,
11575                m.down_exps.in_f,
11576                m.down_exps.out_f,
11577                m.down_exps.qtype,
11578                m.down_exps.row_bytes,
11579            )
11580        };
11581
11582        // Pass 1 — PEERS: one bulk fan-out per pair-owning rank, then every owned row
11583        // back-to-back on that rank's stream into its compact block. No host boundary until
11584        // the bulk returns.
11585        let hop = ep.rt.hop(e);
11586        let mut y_peer_blks: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11587        for r in 1..ranks {
11588            if per_rank[r] == 0 {
11589                continue;
11590            }
11591            let dev = crate::glm5_tp::rank_engine(e, rt, r);
11592            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only
11593            // (a rank with zero owned pairs moves zero activation bytes off root).
11594            let z_r = crate::tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11595            // Under the skip-peer-combine red the block stays ZERO for skipped rows (a red
11596            // must drop the peer contribution loudly, never multiply garbage into the chain).
11597            let mut blk = if red_skip_peer {
11598                dev.zeros(per_rank[r] * n_embd)?
11599            } else {
11600                dev.uninit(per_rank[r] * n_embd)?
11601            };
11602            let mut k = 0usize;
11603            for tok in 0..t {
11604                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11605                for &ex in sel.iter() {
11606                    let ex = ex as usize;
11607                    if ep.owner(ex) != r {
11608                        continue;
11609                    }
11610                    // Engagement counters FIRST (a red skip still counts as ROUTED).
11611                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11612                    crate::glm5_tp::GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED
11613                        .fetch_add(1, Ordering::Relaxed);
11614                    if red_skip_peer {
11615                        k += 1;
11616                        continue;
11617                    }
11618                    let zt_r = z_r.slice(tok * n_embd..(tok + 1) * n_embd);
11619                    let y = expert_row(dev, &ep.slabs[r], &zt_r, ex)?;
11620                    dev.copy_into(&mut blk, k * n_embd, &y, n_embd)?;
11621                    k += 1;
11622                }
11623            }
11624            y_peer_blks[r] = Some(blk);
11625        }
11626
11627        // Pass 2 — ROOT: every root-owned row into the slab head, never blocked on a peer
11628        // return (the v1 walk interleaved root issue behind per-slot peer syncs).
11629        let mut y_all = e.uninit(n_pairs * n_embd)?;
11630        {
11631            let mut k = 0usize;
11632            for tok in 0..t {
11633                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11634                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11635                for &ex in sel.iter() {
11636                    let ex = ex as usize;
11637                    if ep.owner(ex) != 0 {
11638                        continue;
11639                    }
11640                    let y = expert_row(e, &ep.slabs[0], &zt, ex)?;
11641                    e.copy_into(&mut y_all, k * n_embd, &y, n_embd)?;
11642                    k += 1;
11643                }
11644            }
11645        }
11646
11647        // SWAP POINT 2 (bulk returns) — Pass 3: ONE rank->root block move into each rank's
11648        // tail segment. On host-canonical each is the ONE draining peer sync of that rank's
11649        // layer-call share, exactly as before; on peer-pull each is one event-ordered device
11650        // copy and no host boundary at all.
11651        for r in 1..ranks {
11652            if let Some(blk) = &y_peer_blks[r] {
11653                crate::tp_transport::return_block_to_root(
11654                    &hop,
11655                    r,
11656                    blk,
11657                    &mut y_all,
11658                    base[r] * n_embd,
11659                    per_rank[r] * n_embd,
11660                )?;
11661                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11662            }
11663        }
11664
11665        // Pass 4 — ONE combine launch. Weights are v1's exact host fold, placed at slab
11666        // positions; the scatter walks each token's pairs in SLOT order (ids[p], p pair-major),
11667        // reproducing zeros + n_used sequential axpy_f32 per the kernel's bit contract.
11668        let mut wd = vec![0f32; n_pairs];
11669        for ((&id, &w), &s) in ids.iter().zip(w_all.iter()).zip(sel_all.iter()) {
11670            wd[id as usize] = w * m.down_exps.macro_scale(s as usize);
11671        }
11672        let pw = e.htod(&wd)?;
11673        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
11674        let toff_d = e.htod_i32(&toff)?;
11675        let ids_d = e.htod_i32(&ids)?;
11676        let mut moe_out = e.uninit(t * n_embd)?; // the scatter fully overwrites
11677        e.moe_pairs_scatter(&y_all, &pw, &toff_d, &ids_d, &mut moe_out, t, n_embd)?;
11678        Ok(moe_out)
11679    }
11680
11681    /// The EP GROUPED PRIME (`MEMRA_GLM5_EP_GROUPED_PRIME`, lane/glm5-ep-diet): the plain
11682    /// walk's grouped-prefill program (`moe_ffn_grouped_prefill_sigmoid`, default ON on the
11683    /// serving artifact — 85 -> 616-639 tok/s prefill in its box A/B) split by expert
11684    /// ownership. Per rank: expert-major CSR over the rank's OWNED (token, expert) pairs,
11685    /// one grouped f16 GEMM per projection over the rank's resident EP slab (pointer tables
11686    /// minted at arm time), the PRE-clamped SwiGLU epilogue, the per-expert macro folds,
11687    /// and the slot-ordered per-token scatter — all composed from the SAME Engine calls the
11688    /// plain arm makes, so per-expert GEMM bytes match the plain grouped arm's (grouping is
11689    /// per expert, and an expert's token rows all live on its owner). The ONE new
11690    /// reassociation is the per-token partial add (root chain + peer chain instead of one
11691    /// 8-term chain) — band-gated on minted NVFP4 slabs (`glm5_ep_diet_doors_gpu`), never
11692    /// claimed byte.
11693    ///
11694    /// Returns `Ok(None)` — fall closed to the sequential EP walk — whenever the plain
11695    /// grouped arm's own admission would (f16g-ineligible qtypes, bank/top-k shape, no
11696    /// sigmoid clamp form). The rig fixture's Q8_0 bank therefore ALWAYS falls closed;
11697    /// `glm5-tp-gate`'s grouped arm proves exactly that (dispatch counter pinned 0, walk
11698    /// bytes unchanged).
11699    #[allow(clippy::too_many_arguments)]
11700    fn moe_ffn_glm5_ep_grouped_prime(
11701        e: &Engine,
11702        m: &MoeWeights,
11703        ep: &crate::glm5_tp::Glm5EpExps,
11704        z: &CudaSlice<f32>,
11705        sel_all: &[u32],
11706        w_all: &[f32],
11707        t: usize,
11708        cfg: &ModelConfig,
11709        il: u16,
11710    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11711        use std::sync::atomic::Ordering;
11712        let moe = cfg
11713            .moe
11714            .as_ref()
11715            .ok_or("glm5 EP grouped prime requires MoE model metadata")?;
11716        let n_embd = cfg.n_embd as usize;
11717        let n_expert = moe.expert_count as usize;
11718        let n_used = moe.expert_used_count as usize;
11719        let n_ff_exp = moe.expert_ff_length as usize;
11720        // The plain grouped arm's admission, mirrored term for term (fall closed, never a
11721        // new admission class). MEMRA_MOE_GATE is the sequential byte-identity oracle; this
11722        // arm is a band class and must not shadow that comparison.
11723        if crate::moe_f16g_mode() == 0 || std::env::var("MEMRA_MOE_GATE").is_ok() {
11724            return Ok(None);
11725        }
11726        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
11727            && f16g_proj_ok(m.up_exps.qtype, n_embd)
11728            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
11729        {
11730            return Ok(None);
11731        }
11732        if n_expert > 512 || n_used == 0 || n_used > 8 {
11733            return Ok(None);
11734        }
11735        let lim_exp = cfg.clamp_exp_at(il as u32);
11736        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
11737            return Err(
11738                "EP grouped prime is qualified for the PRE-clamped SwiGLU form only; \
11739                 a POST-clamp layer must ride the sequential arm"
11740                    .into(),
11741            );
11742        }
11743        let n_pairs = t * n_used;
11744        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11745            return Err("EP grouped prime geometry".into());
11746        }
11747        let rt = &ep.rt;
11748        let red_skip_peer = matches!(
11749            crate::glm5_tp::gate_red(),
11750            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11751        );
11752
11753        // One rank's whole grouped program: CSR over OWNED pairs -> grouped gate/up GEMMs ->
11754        // macro folds -> PRE-clamped epilogue -> grouped down GEMM -> CSR->local permute ->
11755        // slot-ordered scatter into the rank partial [t, n_embd] (empty token windows write
11756        // 0.0 — the scatter fully overwrites, so partials add cleanly on root).
11757        let rank_pass = |dev: &Engine,
11758                         rank: u8,
11759                         ptr_row: &CudaSlice<u64>,
11760                         z_dev: &CudaSlice<f32>|
11761         -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11762            // Expert-major CSR restricted to this rank, local pair index l in ascending
11763            // global-pair order (so per-token slot order == ascending l).
11764            let mut buckets_l: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
11765            let mut local_tok = Vec::new(); // token of local pair l
11766            let mut local_ex = Vec::new(); // expert of local pair l (macro folds)
11767            let mut local_wd = Vec::new(); // v1's exact weight fold, at local positions
11768            let mut local_count_per_tok = vec![0i32; t];
11769            for p in 0..n_pairs {
11770                let ex = sel_all[p] as usize;
11771                if ex >= n_expert {
11772                    return Err(format!("EP grouped prime selection {ex} >= {n_expert}").into());
11773                }
11774                if ep.owner(ex) != rank as usize {
11775                    continue;
11776                }
11777                let l = local_tok.len() as i32;
11778                buckets_l[ex].push(l);
11779                let tok = p / n_used;
11780                local_tok.push(tok as i32);
11781                local_ex.push(ex);
11782                local_wd.push(w_all[p] * m.down_exps.macro_scale(ex));
11783                local_count_per_tok[tok] += 1;
11784            }
11785            let n_owned = local_tok.len();
11786            if n_owned == 0 {
11787                return Ok(None);
11788            }
11789            let mut ex_ids: Vec<i32> = Vec::new();
11790            let mut ex_off: Vec<i32> = vec![0];
11791            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_owned); // local l, CSR order
11792            let mut csr_tok: Vec<i32> = Vec::with_capacity(n_owned);
11793            for (e_id, b) in buckets_l.iter().enumerate() {
11794                if !b.is_empty() {
11795                    ex_ids.push(e_id as i32);
11796                    for &l in b {
11797                        ex_pairs.push(l);
11798                        csr_tok.push(local_tok[l as usize]);
11799                    }
11800                    ex_off.push(ex_pairs.len() as i32);
11801                }
11802            }
11803            let n_active = ex_ids.len();
11804            if n_active == 0 || n_active > 512 {
11805                return Err(format!("EP grouped prime n_active {n_active} outside 1..=512").into());
11806            }
11807
11808            let exi = dev.htod_i32(&ex_ids)?;
11809            let exo = dev.htod_i32(&ex_off)?;
11810            let exp_d = dev.htod_i32(&ex_pairs)?;
11811            let csr_tok_d = dev.htod_i32(&csr_tok)?;
11812
11813            // GATE/UP grouped GEMMs over the rank slab, CSR order end to end.
11814            let (z16, zs) = dev.moe_f16g_act(z_dev, Some(&csr_tok_d), n_embd, n_owned)?;
11815            let mut g = dev.moe_f16_grouped(
11816                ptr_row,
11817                0,
11818                n_expert,
11819                &exi,
11820                &ex_off,
11821                &exo,
11822                &z16,
11823                &zs,
11824                n_embd,
11825                n_ff_exp,
11826                n_active,
11827                n_owned,
11828                m.gate_exps.qtype,
11829                m.gate_exps.row_bytes,
11830            )?;
11831            if m.gate_exps.macros.is_some() {
11832                let mg: Vec<f32> = ex_pairs
11833                    .iter()
11834                    .map(|&l| m.gate_exps.macro_scale(local_ex[l as usize]))
11835                    .collect();
11836                let mg_d = dev.htod(&mg)?;
11837                dev.scale_rows(&mut g, &mg_d, n_ff_exp, n_owned)?;
11838            }
11839            let mut u = dev.moe_f16_grouped(
11840                ptr_row,
11841                1,
11842                n_expert,
11843                &exi,
11844                &ex_off,
11845                &exo,
11846                &z16,
11847                &zs,
11848                n_embd,
11849                n_ff_exp,
11850                n_active,
11851                n_owned,
11852                m.up_exps.qtype,
11853                m.up_exps.row_bytes,
11854            )?;
11855            if m.up_exps.macros.is_some() {
11856                let mu: Vec<f32> = ex_pairs
11857                    .iter()
11858                    .map(|&l| m.up_exps.macro_scale(local_ex[l as usize]))
11859                    .collect();
11860                let mu_d = dev.htod(&mu)?;
11861                dev.scale_rows(&mut u, &mu_d, n_ff_exp, n_owned)?;
11862            }
11863
11864            // Epilogue: PRE-clamped SwiGLU (POST refused above), plain-silu pair otherwise.
11865            let act = match lim_exp {
11866                Some(SwigluClamp::Pre(limit)) => {
11867                    let mut a = dev.uninit(n_owned * n_ff_exp)?;
11868                    dev.swiglu_preclamped_mul_scaled(
11869                        &g,
11870                        &u,
11871                        1.0,
11872                        1.0,
11873                        limit,
11874                        &mut a,
11875                        n_owned * n_ff_exp,
11876                    )?;
11877                    a
11878                }
11879                None => dev.moe_pairs_silu_mul(&g, &u, n_owned * n_ff_exp)?,
11880                Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
11881            };
11882
11883            // DOWN grouped GEMM, permute CSR -> local pair order, slot-ordered scatter.
11884            let (a16, a_s) = dev.moe_f16g_act(&act, None, n_ff_exp, n_owned)?;
11885            let d_csr = dev.moe_f16_grouped(
11886                ptr_row,
11887                2,
11888                n_expert,
11889                &exi,
11890                &ex_off,
11891                &exo,
11892                &a16,
11893                &a_s,
11894                n_ff_exp,
11895                n_embd,
11896                n_active,
11897                n_owned,
11898                m.down_exps.qtype,
11899                m.down_exps.row_bytes,
11900            )?;
11901            let y_local = dev.rows_permute(&d_csr, &exp_d, n_owned, n_embd)?;
11902            let mut toff: Vec<i32> = Vec::with_capacity(t + 1);
11903            let mut acc = 0i32;
11904            toff.push(0);
11905            for &c in &local_count_per_tok {
11906                acc += c;
11907                toff.push(acc);
11908            }
11909            let tids: Vec<i32> = (0..n_owned as i32).collect();
11910            let pw = dev.htod(&local_wd)?;
11911            let toff_d = dev.htod_i32(&toff)?;
11912            let tids_d = dev.htod_i32(&tids)?;
11913            let mut partial = dev.uninit(t * n_embd)?; // scatter fully overwrites
11914            dev.moe_pairs_scatter(&y_local, &pw, &toff_d, &tids_d, &mut partial, t, n_embd)?;
11915            Ok(Some(partial))
11916        };
11917
11918        // Peer passes first (their GEMMs overlap root's), each on its own runtime binding —
11919        // the grouped-MoE FFI follows the RUNTIME device, not cudarc's pushed context
11920        // (`bind_runtime_device`'s contract). Engagement counters count ROUTED peer pairs
11921        // before any red skip, exactly like the sequential walk.
11922        let ranks = ep.ranks();
11923        let n_peer_pairs = sel_all
11924            .iter()
11925            .take(n_pairs)
11926            .filter(|&&ex| ep.owner(ex as usize) != 0)
11927            .count() as u64;
11928        crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(n_peer_pairs, Ordering::Relaxed);
11929        let hop = ep.rt.hop(e);
11930        let mut peer_partials: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11931        for r in 1..ranks {
11932            let rank_owns_pairs = sel_all
11933                .iter()
11934                .take(n_pairs)
11935                .any(|&ex| ep.owner(ex as usize) == r);
11936            if !rank_owns_pairs {
11937                continue;
11938            }
11939            let dev = crate::glm5_tp::rank_engine(e, rt, r);
11940            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only.
11941            let z_r = crate::tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11942            dev.bind_runtime_device(dev.ctx().ordinal() as i32)?;
11943            let res = rank_pass(dev, r as u8, &ep.ptr_rows[r], &z_r);
11944            e.bind_runtime_device(e.ctx().ordinal() as i32)?;
11945            peer_partials[r] = res?;
11946        }
11947        let root_partial = rank_pass(e, 0, &ep.ptr_rows[0], z)?;
11948
11949        // Root combine: root partial + bulk-returned peer partials (SWAP POINT 2, the named
11950        // transport shape). One partial add per contributing rank — the same reassociation
11951        // class the two-rank arm band-gated (root chain + per-rank chains instead of one
11952        // 8-term chain), never claimed byte. The skip-peer-combine red drops every peer
11953        // partial AFTER counting — the loud non-vacuity arm.
11954        let mut out = match root_partial {
11955            Some(p) => p,
11956            None => e.zeros(t * n_embd)?,
11957        };
11958        for r in 1..ranks {
11959            if let Some(pp) = &peer_partials[r]
11960                && !red_skip_peer
11961            {
11962                let pp_root = crate::tp_transport::return_row_to_root(&hop, r, pp, t * n_embd)?;
11963                let mut dst = out.slice_mut(0..t * n_embd);
11964                e.axpy_into(&pp_root, 1.0, &mut dst, t * n_embd)?;
11965                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11966            }
11967        }
11968        crate::glm5_tp::GLM5_EP_GROUPED_PRIME_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11969        static EPGP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
11970        let layer_bit = 1u64 << (il as u64 % 64);
11971        if EPGP_LOGGED.fetch_or(layer_bit, Ordering::Relaxed) & layer_bit == 0 {
11972            eprintln!(
11973                "[glm5-ep-grouped-prime] execute layer={il} tokens={t} \
11974                 provenance=ep-rank-slabs router=sigmoid-host-oracle epilogue=pre-clamped \
11975                 combine=rank-partial-add transport={} performance_claim=false \
11976                 (logged once per layer)",
11977                hop.transport.name(),
11978            );
11979        }
11980        Ok(Some(out))
11981    }
11982
11983    /// Step 3 of the MoE body — SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z —
11984    /// qwen35moe only. OLMoE and most vanilla MoE have NO shared expert (the shexp tensors
11985    /// are absent / `None`); skip it then. gate_inp_shexp is OPTIONAL: qwen35moe gates the
11986    /// shared expert (sigmoid(gate_inp) x sh); MiniMax-M3 (DeepSeek-V3 class) has NO shexp
11987    /// gate — the shared expert adds directly. (Extracted verbatim from the sequential body
11988    /// so the glm5 EP-2 walk adds the ROOT-owned shared expert through the identical
11989    /// program.)
11990    #[allow(clippy::too_many_arguments)]
11991    fn moe_shexp_add(
11992        e: &Engine,
11993        m: &MoeWeights,
11994        z: &CudaSlice<f32>,
11995        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11996        t: usize,
11997        cfg: &ModelConfig,
11998        lim_shexp: Option<memra_gguf::config::SwigluClamp>,
11999        moe_out: &mut CudaSlice<f32>,
12000    ) -> Result<(), Box<dyn std::error::Error>> {
12001        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
12002            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
12003        {
12004            let n_embd = cfg.n_embd as usize;
12005            let n_ff_sh = gate_shexp.out_features(); // 512
12006            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
12007            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
12008            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
12009            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
12010            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
12011            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
12012            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
12013            let verify_t = t > 1 && t < PRIME_MIN_T;
12014            let (sg_gate, sg_up) = if t == 1 {
12015                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
12016            } else if verify_t {
12017                (
12018                    e.matmul_decode_exact(gate_shexp, z, t)?,
12019                    e.matmul_decode_exact(up_shexp, z, t)?,
12020                )
12021            } else {
12022                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
12023            };
12024            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
12025            Self::ffn_act_lim(
12026                e,
12027                cfg,
12028                &sg_gate,
12029                &sg_up,
12030                1.0,
12031                1.0,
12032                lim_shexp,
12033                &mut sa,
12034                t * n_ff_sh,
12035            )?;
12036            let sh = if verify_t {
12037                e.matmul_decode_exact(down_shexp, &sa, t)?
12038            } else {
12039                e.matmul(down_shexp, &sa, t)?
12040            }; // [T, n_embd]
12041
12042            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
12043            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
12044            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
12045            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
12046            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
12047            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
12048            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
12049            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
12050            // expert's contribution into every token's residual, so under cross-request
12051            // concat prefill a session's hidden state depended on its co-arrivals' token
12052            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
12053            // DOOR H (`MEMRA_GLM5_HTOD_DIET`): glm5 has no `ffn_gate_inp_shexp`, so this is the
12054            // LIVE arm on the serving artifact and it re-uploaded a constant `vec![1.0f32; t]`
12055            // on every MoE layer-call — 42 pageable HtoD per ship round (26.9% of the round's
12056            // 156). The resident ones buffer feeds the SAME `add_scaled_rows_f32` kernel the
12057            // same 1.0 values, so the arms are bit-identical.
12058            if m.gate_inp_shexp.is_none() && crate::htod_diet_on() {
12059                crate::HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12060                e.add_scaled_rows_ones(&sh, moe_out, n_embd, t)?;
12061                return Ok(());
12062            }
12063            let g = match &m.gate_inp_shexp {
12064                Some(gate_inp_shexp) => {
12065                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
12066                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
12067                    } else {
12068                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
12069                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
12070                        e.sigmoid(&gs, &mut g, t)?;
12071                        g
12072                    }
12073                }
12074                None => e.htod(&vec![1.0f32; t])?,
12075            };
12076            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
12077            e.add_scaled_rows(&sh, &g, moe_out, n_embd, t)?;
12078        }
12079        Ok(())
12080    }
12081
12082    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
12083    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
12084    pub fn stage1_h2d_per_token(&self) -> u64 {
12085        use crate::hybrid::Ffn;
12086        let n_used = self
12087            .cfg
12088            .moe
12089            .as_ref()
12090            .map(|m| m.expert_used_count as u64)
12091            .unwrap_or(0);
12092        let mut bytes = 0u64;
12093        for l in self.layers.iter() {
12094            if let Ffn::Moe(m) = &l.ffn {
12095                bytes += n_used
12096                    * (m.gate_exps.max_expert_bytes()
12097                        + m.up_exps.max_expert_bytes()
12098                        + m.down_exps.max_expert_bytes()) as u64;
12099            }
12100        }
12101        bytes
12102    }
12103
12104    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
12105    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
12106    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
12107    pub(crate) fn max_moe_block(&self) -> usize {
12108        use crate::hybrid::Ffn;
12109        let mut mx = 0usize;
12110        let mut scan = |ffn: &Ffn| {
12111            if let Ffn::Moe(m) = ffn {
12112                mx = mx
12113                    .max(m.gate_exps.max_expert_bytes())
12114                    .max(m.up_exps.max_expert_bytes())
12115                    .max(m.down_exps.max_expert_bytes());
12116            }
12117        };
12118        for l in self.layers.iter() {
12119            scan(&l.ffn);
12120        }
12121        if let Some(mtp) = self.mtp.as_ref() {
12122            scan(&mtp.ffn);
12123        }
12124        mx
12125    }
12126
12127    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
12128    /// but have no bytes and therefore consume no residency slot.
12129    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
12130        use crate::hybrid::Ffn;
12131        let mut sizes = Vec::new();
12132        let mut scan = |ffn: &Ffn| {
12133            let Ffn::Moe(m) = ffn else { return };
12134            for ex in 0..m.gate_exps.n_expert {
12135                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
12136                    continue;
12137                }
12138                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
12139                    let len = exps.expert_layout(ex).len;
12140                    if len > 0 {
12141                        sizes.push(len);
12142                    }
12143                }
12144            }
12145        };
12146        for layer in &self.layers {
12147            scan(&layer.ffn);
12148        }
12149        if let Some(mtp) = &self.mtp {
12150            scan(&mtp.ffn);
12151        }
12152        sizes
12153    }
12154
12155    /// Persist the frozen residency set so a later process can restage it directly and skip
12156    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
12157    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
12158    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
12159    /// post-freeze argmax gate still validates the serving assignment.
12160    pub fn save_cpu_expert_residency_profile(
12161        &self,
12162        e: &Engine,
12163        path: &std::path::Path,
12164    ) -> Result<(), Box<dyn std::error::Error>> {
12165        let Some(ids) = e.export_moe_residency() else {
12166            return Err("no MoE residency cache to persist".into());
12167        };
12168        let mut body = format!(
12169            "memra-freeze-profile v1 max_block={} blocks={}\n",
12170            self.max_moe_block(),
12171            ids.len()
12172        );
12173        for (layer, proj, ex) in &ids {
12174            body.push_str(&format!("{layer} {proj} {ex}\n"));
12175        }
12176        let tmp = path.with_extension("tmp");
12177        std::fs::write(&tmp, body)?;
12178        std::fs::rename(&tmp, path)?;
12179        println!(
12180            "[moe-cache] freeze profile saved: {} blocks -> {}",
12181            ids.len(),
12182            path.display()
12183        );
12184        Ok(())
12185    }
12186
12187    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
12188    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
12189    /// missing or its header does not match this model's slot geometry.
12190    pub fn restore_cpu_expert_residency_profile(
12191        &self,
12192        e: &Engine,
12193        path: &std::path::Path,
12194    ) -> Result<bool, Box<dyn std::error::Error>> {
12195        use crate::hybrid::Ffn;
12196        use crate::moe_cache::BlockId;
12197        let Ok(content) = std::fs::read_to_string(path) else {
12198            return Ok(false);
12199        };
12200        let mut lines = content.lines();
12201        let Some(header) = lines.next() else {
12202            return Ok(false);
12203        };
12204        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
12205        if !header.starts_with(&expected) {
12206            println!(
12207                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
12208                path.display()
12209            );
12210            return Ok(false);
12211        }
12212        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
12213            std::collections::HashMap::new();
12214        for line in lines {
12215            let mut fields = line.split_whitespace();
12216            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
12217            else {
12218                continue;
12219            };
12220            let (Ok(layer), Ok(proj), Ok(ex)) =
12221                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
12222            else {
12223                continue;
12224            };
12225            by_layer
12226                .entry(layer)
12227                .or_default()
12228                .push(BlockId::new(layer, proj, ex));
12229        }
12230        let requested: usize = by_layer.values().map(Vec::len).sum();
12231        if requested == 0 {
12232            return Ok(false);
12233        }
12234        let max_block = self.max_moe_block();
12235        let mut restaged = 0usize;
12236        let mut stage_layer =
12237            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
12238                let Ffn::Moe(m) = ffn else { return Ok(()) };
12239                let Some(ids) = by_layer.get(&layer_index) else {
12240                    return Ok(());
12241                };
12242                e.with_moe_cache(max_block, |cache, eng| {
12243                    for id in ids {
12244                        if cache.restage_block(*id, m, eng)? {
12245                            restaged += 1;
12246                        }
12247                    }
12248                    Ok(())
12249                })
12250            };
12251        for (index, layer) in self.layers.iter().enumerate() {
12252            stage_layer(index as u16, &layer.ffn)?;
12253        }
12254        if let Some(mtp) = self.mtp.as_ref() {
12255            stage_layer(u16::MAX, &mtp.ffn)?;
12256        }
12257        e.freeze_moe_cache();
12258        println!(
12259            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
12260            path.display()
12261        );
12262        Ok(true)
12263    }
12264
12265    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
12266    pub fn freeze_cpu_expert_residency(
12267        &self,
12268        e: &Engine,
12269    ) -> Result<(), Box<dyn std::error::Error>> {
12270        e.freeze_moe_cache();
12271        Ok(())
12272    }
12273
12274    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
12275    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
12276    /// the model's activation exactly.
12277    ///
12278    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
12279    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
12280    /// form for anything that can land on a clamped layer.
12281    pub fn ffn_act(
12282        e: &Engine,
12283        cfg: &ModelConfig,
12284        gate: &CudaSlice<f32>,
12285        up: &CudaSlice<f32>,
12286        act: &mut CudaSlice<f32>,
12287        n: usize,
12288    ) -> Result<(), Box<dyn std::error::Error>> {
12289        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
12290    }
12291
12292    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
12293    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
12294    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
12295    #[allow(clippy::too_many_arguments)]
12296    pub(crate) fn ffn_act_scaled(
12297        e: &Engine,
12298        cfg: &ModelConfig,
12299        gate: &CudaSlice<f32>,
12300        up: &CudaSlice<f32>,
12301        gs: f32,
12302        us: f32,
12303        act: &mut CudaSlice<f32>,
12304        n: usize,
12305    ) -> Result<(), Box<dyn std::error::Error>> {
12306        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
12307    }
12308
12309    /// ffn_act_scaled + a PER-LAYER clamped SwiGLU. `limit`:
12310    ///   * `None`          -> the unclamped dispatch (every arch with no live clamp).
12311    ///   * `Some(Post(l))` -> step35: `min(silu(gate*gs), l) * clamp(up*us, +-l)`
12312    ///     (llama-graph.cpp:2146/1751, non-DEEPSEEK4 branch).
12313    ///   * `Some(Pre(l))`  -> glm5_next: `silu(min(gate*gs, l)) * clamp(up*us, +-l)`.
12314    ///     Callers source it from `cfg.clamp_exp_at(il)` (routed experts) or `cfg.clamp_shexp_at(il)`
12315    ///     (shared expert / dense MLP) — on step35 the two arrays are SEPARATE and a layer can have
12316    ///     one without the other. The `> 1e-6` eps gate lives in the accessors, so a `Some` here is
12317    ///     already known live. The match is exhaustive so a new clamp form cannot default to either
12318    ///     existing one.
12319    #[allow(clippy::too_many_arguments)]
12320    pub(crate) fn ffn_act_lim(
12321        e: &Engine,
12322        cfg: &ModelConfig,
12323        gate: &CudaSlice<f32>,
12324        up: &CudaSlice<f32>,
12325        gs: f32,
12326        us: f32,
12327        limit: Option<SwigluClamp>,
12328        act: &mut CudaSlice<f32>,
12329        n: usize,
12330    ) -> Result<(), Box<dyn std::error::Error>> {
12331        if let Some(m3) = cfg.m3.as_ref() {
12332            debug_assert!(
12333                limit.is_none(),
12334                "m3 swigluoai and the step35/glm5_next clamps are different archs"
12335            );
12336            return e.swigluoai_mul_scaled(
12337                gate,
12338                up,
12339                gs,
12340                us,
12341                m3.swiglu_alpha,
12342                m3.swiglu_limit,
12343                act,
12344                n,
12345            );
12346        }
12347        match limit {
12348            Some(SwigluClamp::Post(l)) => {
12349                return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
12350            }
12351            Some(SwigluClamp::Pre(l)) => {
12352                return e.swiglu_preclamped_mul_scaled(gate, up, gs, us, l, act, n);
12353            }
12354            None => {}
12355        }
12356        if gs == 1.0 && us == 1.0 {
12357            return e.silu_mul(gate, up, act, n);
12358        }
12359        e.silu_mul_scaled(gate, up, gs, us, act, n)
12360    }
12361
12362    /// The bare POST limit for the fused kernels whose epilogue HARDCODES step35's form
12363    /// (`matvec_bf16_dual_silu` / `_rows`, qmatvec.cu:10676). `Ok` = the kernel may run;
12364    /// `Err(())` = glm5_next's PRE form, which has no fused twin, and the caller MUST return its
12365    /// not-handled value so the layer falls through to the unfused `ffn_act_lim` seam. Feeding a
12366    /// PRE limit to a POST epilogue compiles, runs, and returns plausible-but-wrong logits.
12367    fn fused_post_limit(lim: Option<SwigluClamp>) -> Result<Option<f32>, ()> {
12368        match lim {
12369            None => Ok(None),
12370            Some(SwigluClamp::Post(l)) => Ok(Some(l)),
12371            Some(SwigluClamp::Pre(_)) => Err(()),
12372        }
12373    }
12374
12375    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
12376    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
12377    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
12378    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
12379    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
12380    fn moe_route(
12381        e: &Engine,
12382        logits: &CudaSlice<f32>,
12383        t: usize,
12384        n_expert: usize,
12385        n_used: usize,
12386    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12387        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
12388    }
12389
12390    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
12391    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
12392    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
12393    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
12394    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
12395    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
12396    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
12397    #[allow(clippy::too_many_arguments)]
12398    fn moe_route_sigmoid_cfg(
12399        e: &Engine,
12400        logits: &CudaSlice<f32>,
12401        t: usize,
12402        n_expert: usize,
12403        n_used: usize,
12404        m: &MoeWeights,
12405        (sf, route_norm): (f32, bool),
12406    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12407        if sigmoid_router_enabled() {
12408            return e.moe_router_sigmoid_topk_host(
12409                logits,
12410                t,
12411                n_expert,
12412                n_used,
12413                m.active_count(),
12414                &m.exp_probs_b_dev,
12415                &m.active_experts_dev,
12416                sf,
12417                route_norm,
12418            );
12419        }
12420        let lg = e.dtoh(logits)?;
12421        Self::moe_route_sigmoid_host(
12422            &lg,
12423            t,
12424            n_expert,
12425            n_used,
12426            m.exp_probs_b.as_deref(),
12427            sf,
12428            route_norm,
12429            m.active_experts.as_deref(),
12430        )
12431    }
12432
12433    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
12434    /// the existing softmax device kernel has no mask input.
12435    #[allow(clippy::excessive_precision)] // allow: literal kept verbatim from the reference/measured value
12436    fn moe_route_cfg(
12437        e: &Engine,
12438        logits: &CudaSlice<f32>,
12439        t: usize,
12440        n_expert: usize,
12441        n_used: usize,
12442        active: Option<&[bool]>,
12443    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12444        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
12445        // rollback) via the single-sync pinned readback — softmax arch only.
12446        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
12447            return e.moe_router_topk_host(logits, t, n_expert, n_used);
12448        }
12449        // Host oracle (the §D bit-identity reference).
12450        let lg = e.dtoh(logits)?; // [T*n_expert] host
12451        let mut sel = vec![0u32; t * n_used];
12452        let mut w_out = vec![0f32; t * n_used];
12453        for tok in 0..t {
12454            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12455            // softmax over ALL n_expert (stable: subtract max)
12456            let maxl = row
12457                .iter()
12458                .enumerate()
12459                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
12460                .map(|(_, &x)| x)
12461                .fold(f32::NEG_INFINITY, f32::max);
12462            let mut probs = vec![0f32; n_expert];
12463            let mut den = 0f32;
12464            for i in 0..n_expert {
12465                if active.is_some_and(|mask| !mask[i]) {
12466                    continue;
12467                }
12468                let x = (row[i] - maxl).exp();
12469                probs[i] = x;
12470                den += x;
12471            }
12472            for p in probs.iter_mut() {
12473                *p /= den;
12474            }
12475            // stable DESC sort: prob DESC, ascending-index tiebreak.
12476            let mut idx: Vec<usize> = (0..n_expert)
12477                .filter(|&i| active.is_none_or(|mask| mask[i]))
12478                .collect();
12479            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
12480            let sl = &idx[..n_used];
12481            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
12482            let mut ws: f32 = wv.iter().sum();
12483            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
12484            for x in wv.iter_mut() {
12485                *x /= ws;
12486            }
12487            for j in 0..n_used {
12488                sel[tok * n_used + j] = sl[j] as u32;
12489                w_out[tok * n_used + j] = wv[j];
12490            }
12491        }
12492        Ok((sel, w_out))
12493    }
12494
12495    #[allow(clippy::too_many_arguments)]
12496    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
12497    fn moe_route_sigmoid_with_input(
12498        e: &Engine,
12499        logits: &CudaSlice<f32>,
12500        input: &CudaSlice<f32>,
12501        t: usize,
12502        in_features: usize,
12503        n_expert: usize,
12504        n_used: usize,
12505        bias: Option<&[f32]>,
12506        (sf, route_norm): (f32, bool),
12507        active: Option<&[bool]>,
12508    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12509        let logit_values =
12510            active_matrix_values(logits.len(), t, n_expert, "sigmoid router logits")?;
12511        let input_values =
12512            active_matrix_values(input.len(), t, in_features, "sigmoid router input")?;
12513        let (lg, input) = e.dtoh_pair_views(
12514            &logits.slice(0..logit_values),
12515            &input.slice(0..input_values),
12516        )?;
12517        let (sel, w) =
12518            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
12519        Ok((sel, w, input))
12520    }
12521
12522    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
12523    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
12524    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
12525    /// active mask, prebuilt projection descriptors) so no model reference escapes.
12526    pub fn start_moe_prefetch_predictor(
12527        &self,
12528        e: &Engine,
12529        cfg: &ModelConfig,
12530    ) -> Result<(), Box<dyn std::error::Error>> {
12531        use crate::hybrid::Ffn;
12532        let Some(sig) = cfg.sigmoid_router() else {
12533            return Err("prefetch predictor requires a sigmoid-router arch".into());
12534        };
12535        let resident: std::collections::HashSet<(u16, u8, u16)> = e
12536            .export_moe_residency()
12537            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
12538            .into_iter()
12539            .collect();
12540        let mut layers = Vec::new();
12541        for (index, layer) in self.layers.iter().enumerate() {
12542            let Ffn::Moe(m) = &layer.ffn else { continue };
12543            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
12544                continue;
12545            };
12546            let router = e.dtoh(data)?;
12547            let n_expert = m.gate_exps.n_expert;
12548            let n_embd = m.gate_exps.in_f;
12549            if router.len() != n_embd * n_expert {
12550                continue;
12551            }
12552            let build = |exps: &crate::model::HostExps| {
12553                (0..n_expert)
12554                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
12555                    .collect::<Vec<_>>()
12556            };
12557            layers.push((
12558                index as u16,
12559                crate::cpu_experts::PredictLayerInit {
12560                    router,
12561                    bias: m.exp_probs_b.clone(),
12562                    active: m.active_experts.clone(),
12563                    n_embd,
12564                    n_used: cfg
12565                        .moe
12566                        .as_ref()
12567                        .map(|moe| moe.expert_used_count as usize)
12568                        .ok_or("prefetch predictor requires MoE config")?,
12569                    sig,
12570                    weights_n_expert: n_expert,
12571                    gate: build(&m.gate_exps),
12572                    up: build(&m.up_exps),
12573                    down: build(&m.down_exps),
12574                },
12575            ));
12576        }
12577        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
12578    }
12579
12580    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
12581    /// selection math to the rollback runtime, applied to host-computed logits.
12582    #[allow(clippy::too_many_arguments)]
12583    pub fn moe_route_sigmoid_host_public(
12584        logits: &[f32],
12585        t: usize,
12586        n_expert: usize,
12587        n_used: usize,
12588        bias: Option<&[f32]>,
12589        sf: f32,
12590        route_norm: bool,
12591        active: Option<&[bool]>,
12592    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12593        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
12594    }
12595
12596    #[allow(clippy::too_many_arguments)]
12597    fn moe_route_sigmoid_host(
12598        lg: &[f32],
12599        t: usize,
12600        n_expert: usize,
12601        n_used: usize,
12602        bias: Option<&[f32]>,
12603        sf: f32,
12604        route_norm: bool,
12605        active: Option<&[bool]>,
12606    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12607        let active_count = active
12608            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
12609            .unwrap_or(n_expert);
12610        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
12611        if lg.len() != t * n_expert {
12612            return Err(format!(
12613                "sigmoid router logits length mismatch: got {}, expected {}",
12614                lg.len(),
12615                t * n_expert,
12616            )
12617            .into());
12618        }
12619        let mut sel = vec![0u32; t * n_used];
12620        let mut w_out = vec![0f32; t * n_used];
12621        for tok in 0..t {
12622            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12623            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
12624            // selection score = sigmoid + bias; weight = plain sigmoid.
12625            let selsc: Vec<f32> = match bias {
12626                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
12627                None => scores.clone(),
12628            };
12629            let mut idx: Vec<usize> = (0..n_expert)
12630                .filter(|&i| active.is_none_or(|mask| mask[i]))
12631                .collect();
12632            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
12633            let sl = &idx[..n_used];
12634            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
12635            if route_norm {
12636                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
12637                for x in wv.iter_mut() {
12638                    *x = *x / ws * sf;
12639                }
12640            } else {
12641                for x in wv.iter_mut() {
12642                    *x *= sf;
12643                }
12644            }
12645            for j in 0..n_used {
12646                sel[tok * n_used + j] = sl[j] as u32;
12647                w_out[tok * n_used + j] = wv[j];
12648            }
12649        }
12650        Ok((sel, w_out))
12651    }
12652
12653    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
12654    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
12655    /// macro-scaled experts, and observation modes are denied by the caller.
12656    #[allow(clippy::too_many_arguments)]
12657    fn moe_ffn_sigmoid_dev(
12658        e: &Engine,
12659        m: &MoeWeights,
12660        z: &CudaSlice<f32>,
12661        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
12662        logits: &CudaSlice<f32>,
12663        t: usize,
12664        cfg: &ModelConfig,
12665        il: u16,
12666        (scaling_factor, route_norm): (f32, bool),
12667    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12668        let moe = cfg.moe.as_ref().unwrap();
12669        let n_embd = cfg.n_embd as usize;
12670        let n_expert = moe.expert_count as usize;
12671        let n_used = moe.expert_used_count as usize;
12672        let n_ff_exp = moe.expert_ff_length as usize;
12673        let dev = m.dev_exps.as_ref().unwrap();
12674        debug_assert_eq!(dev.dev, e.ctx().ordinal());
12675        debug_assert!(m.has_uniform_expert_layout());
12676        debug_assert!(!m.has_macros);
12677
12678        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
12679            logits,
12680            t,
12681            n_expert,
12682            n_used,
12683            m.active_count(),
12684            &m.exp_probs_b_dev,
12685            &m.active_experts_dev,
12686            scaling_factor,
12687            route_norm,
12688        )?;
12689        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
12690        if let Some(fp8) = dev.fp8_blk.as_ref() {
12691            debug_assert_eq!(m.gate_exps.qtype, crate::QT_F8_E4M3_BLK);
12692            debug_assert_eq!(m.up_exps.qtype, crate::QT_F8_E4M3_BLK);
12693            debug_assert_eq!(m.down_exps.qtype, crate::QT_F8_E4M3_BLK);
12694            debug_assert_eq!(fp8.gate.rows, m.gate_exps.out_f.div_ceil(128));
12695            debug_assert_eq!(fp8.up.rows, m.up_exps.out_f.div_ceil(128));
12696            debug_assert_eq!(fp8.down.rows, m.down_exps.out_f.div_ceil(128));
12697
12698            // Official Step-3.7 FP8 uses dynamic per-token/per-128 E4M3
12699            // activations with block-128 E4M3 weights. This deliberately
12700            // simple resident reference is the correctness oracle for later
12701            // grouped and TP/EP execution. MEMRA_ST_E4M3=0 chooses the
12702            // load-time Q8 diagnostic representation, so one process never
12703            // crosses between numerical programs.
12704            let selected = e.dtoh_i32(&sel_d)?;
12705            let route_weights = e.dtoh(&w_d)?;
12706            let mut moe_out = e.zeros(t * n_embd)?;
12707            for tok in 0..t {
12708                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
12709                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12710                for j in 0..n_used {
12711                    let pair = tok * n_used + j;
12712                    let expert = selected[pair] as usize;
12713                    let gate = Self::moe_resident_fp8_e4m3(
12714                        e,
12715                        &m.gate_exps,
12716                        &dev.gate,
12717                        &fp8.gate,
12718                        expert,
12719                        &zt,
12720                        1,
12721                    )?;
12722                    let up = Self::moe_resident_fp8_e4m3(
12723                        e, &m.up_exps, &dev.up, &fp8.up, expert, &zt, 1,
12724                    )?;
12725                    let mut act = e.uninit(n_ff_exp)?;
12726                    Self::ffn_act_lim(
12727                        e,
12728                        cfg,
12729                        &gate,
12730                        &up,
12731                        1.0,
12732                        1.0,
12733                        cfg.clamp_exp_at(il as u32),
12734                        &mut act,
12735                        n_ff_exp,
12736                    )?;
12737                    let act = act.slice(0..n_ff_exp);
12738                    let down = Self::moe_resident_fp8_e4m3(
12739                        e,
12740                        &m.down_exps,
12741                        &dev.down,
12742                        &fp8.down,
12743                        expert,
12744                        &act,
12745                        1,
12746                    )?;
12747                    e.axpy_into(&down, route_weights[pair], &mut dst, n_embd)?;
12748                }
12749            }
12750            if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12751                eprintln!(
12752                    "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} \
12753                     native=fp8blk-w8a8-e4m3-reference clamp={}",
12754                    cfg.clamp_exp_at(il as u32).is_some(),
12755                );
12756            }
12757            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12758            return Ok(moe_out);
12759        }
12760        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
12761            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12762            (combined, combined)
12763        } else {
12764            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12765        };
12766        let (zq, zd) = match (t, zq8) {
12767            (1, Some((q, d))) => (q.clone(), d.clone()),
12768            _ => e.quantize_q8_1(z, t, n_embd)?,
12769        };
12770        let n_pairs = t * n_used;
12771        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
12772            // The final Step layers retain the established separate gate/up -> clamp -> down
12773            // arithmetic. Pair rows are derived from token position; selected expert ids and
12774            // routing weights remain the device router's buffers throughout.
12775            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
12776            let pair_tok_d = e.htod_i32(&pair_tok)?;
12777            let gate = e.moe_pairs_matvec_q8(
12778                &dev.ptr_row,
12779                0,
12780                &pair_tok_d,
12781                &sel_d,
12782                &zq,
12783                &zd,
12784                n_embd,
12785                n_ff_exp,
12786                n_expert,
12787                n_pairs,
12788                m.gate_exps.qtype,
12789                gate_row_bytes,
12790            )?;
12791            let up = e.moe_pairs_matvec_q8(
12792                &dev.ptr_row,
12793                1,
12794                &pair_tok_d,
12795                &sel_d,
12796                &zq,
12797                &zd,
12798                n_embd,
12799                n_ff_exp,
12800                n_expert,
12801                n_pairs,
12802                m.up_exps.qtype,
12803                up_row_bytes,
12804            )?;
12805            let mut act = e.uninit(n_pairs * n_ff_exp)?;
12806            Self::ffn_act_lim(
12807                e,
12808                cfg,
12809                &gate,
12810                &up,
12811                1.0,
12812                1.0,
12813                cfg.clamp_exp_at(il as u32),
12814                &mut act,
12815                n_pairs * n_ff_exp,
12816            )?;
12817            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12818            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
12819            let pair_self_d = e.htod_i32(&pair_self)?;
12820            let down = e.moe_pairs_matvec_q8(
12821                &dev.ptr_row,
12822                2,
12823                &pair_self_d,
12824                &sel_d,
12825                &aq2,
12826                &ad2,
12827                n_ff_exp,
12828                n_embd,
12829                n_expert,
12830                n_pairs,
12831                m.down_exps.qtype,
12832                m.down_exps.row_bytes,
12833            )?;
12834            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12835            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12836            let tok_off_d = e.htod_i32(&tok_off)?;
12837            let tok_ids_d = e.htod_i32(&tok_ids)?;
12838            let mut output = e.uninit(t * n_embd)?;
12839            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
12840            output
12841        } else {
12842            let act = e.moe_gate_up_silu8_dev_q8_rows(
12843                &dev.ptr_row,
12844                &sel_d,
12845                &zq,
12846                &zd,
12847                t,
12848                n_embd,
12849                n_ff_exp,
12850                n_used,
12851                n_expert,
12852                m.gate_exps.qtype,
12853                m.up_exps.qtype,
12854                gate_row_bytes,
12855                up_row_bytes,
12856                &m.dev_macros,
12857            )?;
12858            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12859            let mut output = e.uninit(t * n_embd)?;
12860            e.moe_down8_fma_dev_q8_rows_g(
12861                &dev.ptr_row,
12862                &sel_d,
12863                &w_d,
12864                &aq2,
12865                &ad2,
12866                &mut output,
12867                t,
12868                n_ff_exp,
12869                n_embd,
12870                n_used,
12871                n_expert,
12872                m.down_exps.qtype,
12873                m.down_exps.row_bytes,
12874            )?;
12875            output
12876        };
12877
12878        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12879            eprintln!(
12880                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
12881                cfg.clamp_exp_at(il as u32).is_some(),
12882                dev.gu_il,
12883            );
12884        }
12885        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12886        Ok(moe_out)
12887    }
12888
12889    #[allow(clippy::too_many_arguments)]
12890    fn moe_resident_fp8_e4m3(
12891        e: &Engine,
12892        exps: &crate::model::HostExps,
12893        bytes: &CudaSlice<u8>,
12894        scales: &crate::hybrid::DevExpertFp8ProjectionScales,
12895        expert: usize,
12896        x: &cudarc::driver::CudaView<f32>,
12897        m: usize,
12898    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12899        let layout = exps.expert_layout(expert);
12900        debug_assert_eq!(layout.qtype, crate::QT_F8_E4M3_BLK);
12901        debug_assert_eq!(scales.rows * scales.cols, scales.expert_stride);
12902        let byte_start = expert * exps.expert_stride;
12903        let scale_start = expert * scales.expert_stride;
12904        let weight = bytes.slice(byte_start..byte_start + layout.len);
12905        let scale = scales
12906            .scales
12907            .slice(scale_start..scale_start + scales.expert_stride);
12908        e.qmatvec_mmq_fp8_blk_view(&weight, &scale, x, m, exps.in_f, exps.out_f)
12909    }
12910
12911    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
12912    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
12913    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
12914    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
12915    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
12916    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
12917    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
12918    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
12919    fn moe_ffn_pairs(
12920        e: &Engine,
12921        m: &MoeWeights,
12922        z: &CudaSlice<f32>,
12923        logits: &CudaSlice<f32>,
12924        t: usize,
12925        cfg: &ModelConfig,
12926    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12927        let moe = cfg.moe.as_ref().unwrap();
12928        let n_embd = cfg.n_embd as usize;
12929        let n_expert = moe.expert_count as usize;
12930        let n_used = moe.expert_used_count as usize;
12931        let n_ff_exp = moe.expert_ff_length as usize;
12932        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
12933        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
12934        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
12935        // that forgets the gate fails loudly in debug instead of returning wrong logits.
12936        debug_assert!(
12937            !cfg.swiglu_clamped_anywhere(),
12938            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
12939        );
12940        let dev = m.dev_exps.as_ref().unwrap();
12941        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
12942        let (rbg_d, rbu_d) = if dev.gu_il {
12943            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12944            (sxx, sxx)
12945        } else {
12946            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12947        };
12948
12949        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
12950        let n_pairs = t * n_used;
12951        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
12952        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
12953        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
12954        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
12955        let pair_w: Vec<f32> = w_all.clone();
12956        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12957        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12958        let pt = e.htod_i32(&pair_tok)?;
12959        let px = e.htod_i32(&pair_ex)?;
12960        let pw = e.htod(&pair_w)?;
12961        let toff = e.htod_i32(&tok_off)?;
12962        let tids = e.htod_i32(&tok_ids)?;
12963
12964        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
12965        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
12966        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
12967        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
12968        for p in 0..n_pairs {
12969            by_ex[pair_ex[p] as usize].push(p as i32);
12970        }
12971        let mut ex_ids: Vec<i32> = Vec::new();
12972        let mut ex_off: Vec<i32> = vec![0];
12973        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
12974        for (ex, list) in by_ex.iter().enumerate() {
12975            if list.is_empty() {
12976                continue;
12977            }
12978            ex_ids.push(ex as i32);
12979            ex_pairs.extend_from_slice(list);
12980            ex_off.push(ex_pairs.len() as i32);
12981        }
12982        let n_active = ex_ids.len();
12983        let exi = e.htod_i32(&ex_ids)?;
12984        let exo = e.htod_i32(&ex_off)?;
12985        let exp_d = e.htod_i32(&ex_pairs)?;
12986        let _ = &px; // pair-major twin keeps it; em path uses CSR
12987
12988        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
12989        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
12990        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
12991        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
12992        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
12993        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
12994        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
12995        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
12996        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
12997        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
12998        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
12999        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
13000        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
13001        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
13002        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
13003        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
13004        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
13005        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
13006        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
13007        let mma_t = *MMA_T.get_or_init(|| {
13008            std::env::var("MEMRA_MOE_MMA_T")
13009                .ok()
13010                .and_then(|v| v.parse().ok())
13011                .unwrap_or(16)
13012        });
13013        let use_mma = std::env::var("MEMRA_MOE_MMA")
13014            .map(|v| v != "0")
13015            .unwrap_or(true)
13016            && t >= mma_t
13017            && q8_expert_dec_supported(m.gate_exps.qtype)
13018            && q8_expert_dec_supported(m.up_exps.qtype)
13019            && q8_expert_dec_supported(m.down_exps.qtype)
13020            && n_embd.is_multiple_of(256)
13021            && n_ff_exp.is_multiple_of(256);
13022        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
13023        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
13024        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
13025        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
13026        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
13027        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
13028        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
13029        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
13030        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
13031        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
13032        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
13033        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
13034        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
13035        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
13036        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
13037        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
13038            && q8_expert_dec_supported(m.up_exps.qtype)
13039            && q8_expert_dec_supported(m.down_exps.qtype)
13040            && n_embd.is_multiple_of(256)
13041            && n_ff_exp.is_multiple_of(256);
13042        let f16g_mode = crate::moe_f16g_mode();
13043        let f16g = f16g_mode != 0
13044            && t >= mma_t
13045            && (f16g_mode != 3 || !mma_capable)
13046            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
13047            && f16g_proj_ok(m.up_exps.qtype, n_embd)
13048            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
13049        if use_mma || f16g {
13050            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
13051            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
13052            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
13053            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
13054            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
13055            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
13056            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
13057            let y_down = if f16g {
13058                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
13059                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
13060                // permute at the very end back to pair-id order for the scatter.
13061                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
13062                let csr_tok_d = e.htod_i32(&csr_tok)?;
13063                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
13064                let g_csr = e.moe_f16_grouped(
13065                    &dev.ptr_row,
13066                    0,
13067                    n_expert,
13068                    &exi,
13069                    &ex_off,
13070                    &exo,
13071                    &z_f16,
13072                    &z_s,
13073                    n_embd,
13074                    n_ff_exp,
13075                    n_active,
13076                    n_pairs,
13077                    m.gate_exps.qtype,
13078                    rbg_d,
13079                )?;
13080                let u_csr = e.moe_f16_grouped(
13081                    &dev.ptr_row,
13082                    1,
13083                    n_expert,
13084                    &exi,
13085                    &ex_off,
13086                    &exo,
13087                    &z_f16,
13088                    &z_s,
13089                    n_embd,
13090                    n_ff_exp,
13091                    n_active,
13092                    n_pairs,
13093                    m.up_exps.qtype,
13094                    rbu_d,
13095                )?;
13096                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
13097                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
13098                let d_csr = e.moe_f16_grouped(
13099                    &dev.ptr_row,
13100                    2,
13101                    n_expert,
13102                    &exi,
13103                    &ex_off,
13104                    &exo,
13105                    &a_f16,
13106                    &a_s,
13107                    n_ff_exp,
13108                    n_embd,
13109                    n_active,
13110                    n_pairs,
13111                    m.down_exps.qtype,
13112                    m.down_exps.row_bytes,
13113                )?;
13114                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
13115            } else {
13116                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
13117                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
13118                let gate = e.mmq_iq_experts(
13119                    &dev.ptr_row,
13120                    0,
13121                    n_expert,
13122                    &exi,
13123                    &exo,
13124                    &exp_d,
13125                    &pt,
13126                    &z_scr,
13127                    n_embd,
13128                    n_ff_exp,
13129                    n_active,
13130                    n_pairs,
13131                    t,
13132                    m.gate_exps.qtype,
13133                    rbg_d,
13134                )?;
13135                let up = e.mmq_iq_experts(
13136                    &dev.ptr_row,
13137                    1,
13138                    n_expert,
13139                    &exi,
13140                    &exo,
13141                    &exp_d,
13142                    &pt,
13143                    &z_scr,
13144                    n_embd,
13145                    n_ff_exp,
13146                    n_active,
13147                    n_pairs,
13148                    t,
13149                    m.up_exps.qtype,
13150                    rbu_d,
13151                )?;
13152                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
13153                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
13154                // registers and writes ONLY the quantized scratch — the two-pass chain
13155                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
13156                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
13157                let a_scr = if crate::moe_fuse_actq_on() {
13158                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
13159                } else {
13160                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13161                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
13162                };
13163                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13164                let pself = e.htod_i32(&pair_self)?;
13165                e.mmq_iq_experts(
13166                    &dev.ptr_row,
13167                    2,
13168                    n_expert,
13169                    &exi,
13170                    &exo,
13171                    &exp_d,
13172                    &pself,
13173                    &a_scr,
13174                    n_ff_exp,
13175                    n_embd,
13176                    n_active,
13177                    n_pairs,
13178                    n_pairs,
13179                    m.down_exps.qtype,
13180                    m.down_exps.row_bytes,
13181                )?
13182            };
13183            let mut moe_out = e.uninit(t * n_embd)?;
13184            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13185            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13186                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13187            {
13188                let n_ff_sh = gate_shexp.out_features();
13189                let sg_gate = e.matmul(gate_shexp, z, t)?;
13190                let sg_up = e.matmul(up_shexp, z, t)?;
13191                let mut sa = e.uninit(t * n_ff_sh)?;
13192                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13193                let sh = e.matmul(down_shexp, &sa, t)?;
13194                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13195                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
13196                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
13197                // i.e. the one real prefill actually takes on a resident-expert MoE model,
13198                // so the concat-prime isolation fix has to land here as well.
13199                let g = match &m.gate_inp_shexp {
13200                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13201                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13202                    }
13203                    Some(gate_inp_shexp) => {
13204                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13205                        let mut g = e.uninit(t)?;
13206                        e.sigmoid(&gs, &mut g, t)?;
13207                        g
13208                    }
13209                    None => e.htod(&vec![1.0f32; t])?,
13210                };
13211                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13212            }
13213            return Ok(moe_out);
13214        }
13215
13216        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
13217        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
13218        let dec = std::env::var("MEMRA_MOE_DEC")
13219            .map(|v| v != "0")
13220            .unwrap_or(true);
13221        let matvec = |proj,
13222                      exi: &_,
13223                      exo: &_,
13224                      exp_d: &_,
13225                      pt: &_,
13226                      aq: &_,
13227                      ad: &_,
13228                      inf,
13229                      outf,
13230                      qtype,
13231                      rb|
13232         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13233            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
13234            let dec = dec && q8_expert_dec_supported(qtype);
13235            if dec {
13236                e.moe_pairs_matvec_q8_dec(
13237                    &dev.ptr_row,
13238                    proj,
13239                    exi,
13240                    exo,
13241                    exp_d,
13242                    pt,
13243                    aq,
13244                    ad,
13245                    inf,
13246                    outf,
13247                    n_expert,
13248                    n_active,
13249                    n_pairs,
13250                    qtype,
13251                    rb,
13252                )
13253            } else {
13254                e.moe_pairs_matvec_q8_em(
13255                    &dev.ptr_row,
13256                    proj,
13257                    exi,
13258                    exo,
13259                    exp_d,
13260                    pt,
13261                    aq,
13262                    ad,
13263                    inf,
13264                    outf,
13265                    n_expert,
13266                    n_active,
13267                    n_pairs,
13268                    qtype,
13269                    rb,
13270                )
13271            }
13272        };
13273        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13274        let gate = matvec(
13275            0,
13276            &exi,
13277            &exo,
13278            &exp_d,
13279            &pt,
13280            &zq,
13281            &zd,
13282            n_embd,
13283            n_ff_exp,
13284            m.gate_exps.qtype,
13285            rbg_d,
13286        )?;
13287        let up = matvec(
13288            1,
13289            &exi,
13290            &exo,
13291            &exp_d,
13292            &pt,
13293            &zq,
13294            &zd,
13295            n_embd,
13296            n_ff_exp,
13297            m.up_exps.qtype,
13298            rbu_d,
13299        )?;
13300        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13301        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13302        // down consumes PAIR-major activation rows: pair_tok = identity.
13303        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13304        let pself = e.htod_i32(&pair_self)?;
13305        let y_down = matvec(
13306            2,
13307            &exi,
13308            &exo,
13309            &exp_d,
13310            &pself,
13311            &aq2,
13312            &ad2,
13313            n_ff_exp,
13314            n_embd,
13315            m.down_exps.qtype,
13316            m.down_exps.row_bytes,
13317        )?;
13318        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
13319        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13320
13321        // SHARED EXPERT epilogue — same as the other paths.
13322        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
13323        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
13324        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13325            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13326        {
13327            let n_ff_sh = gate_shexp.out_features();
13328            // These decode-exact forms are required by the new Step resident arm. Keep the
13329            // established grouped shared-expert program for every other architecture: widening
13330            // this to Gemma changed its speculative acceptance despite green argmax gates.
13331            let step_exact = true;
13332            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
13333            let (sg_gate, sg_up) = if step_exact && t == 1 {
13334                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, None)?
13335            } else if verify_t {
13336                let mut fused = None;
13337                if crate::spec::spec_fused_t()
13338                    && (2..=4).contains(&t)
13339                    && e.uses_q8_1_fast(gate_shexp)
13340                    && e.uses_q8_1_fast(up_shexp)
13341                {
13342                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13343                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13344                }
13345                match fused {
13346                    Some(pair) => pair,
13347                    None => (
13348                        e.matmul_decode_exact(gate_shexp, z, t)?,
13349                        e.matmul_decode_exact(up_shexp, z, t)?,
13350                    ),
13351                }
13352            } else {
13353                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13354            };
13355            let mut sa = e.uninit(t * n_ff_sh)?;
13356            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13357            let sh = if verify_t {
13358                e.matmul_decode_exact(down_shexp, &sa, t)?
13359            } else {
13360                e.matmul(down_shexp, &sa, t)?
13361            };
13362            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13363            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
13364            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
13365            // dispatch choice cannot change bits.
13366            let g = match &m.gate_inp_shexp {
13367                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13368                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13369                }
13370                Some(gate_inp_shexp) => {
13371                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13372                    let mut g = e.uninit(t)?;
13373                    e.sigmoid(&gs, &mut g, t)?;
13374                    g
13375                }
13376                None => e.htod(&vec![1.0f32; t])?,
13377            };
13378            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13379        }
13380        Ok(moe_out)
13381    }
13382
13383    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
13384    #[allow(clippy::too_many_arguments)]
13385    #[allow(clippy::too_many_arguments)]
13386    fn moe_ffn_dev(
13387        e: &Engine,
13388        m: &MoeWeights,
13389        z: &CudaSlice<f32>,
13390        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
13391        logits: &CudaSlice<f32>,
13392        t: usize,
13393        cfg: &ModelConfig,
13394        il: u16,
13395        max_block: usize,
13396    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13397        let moe = cfg.moe.as_ref().unwrap();
13398        let n_embd = cfg.n_embd as usize;
13399        let n_expert = moe.expert_count as usize;
13400        let n_used = moe.expert_used_count as usize;
13401        let n_ff_exp = moe.expert_ff_length as usize;
13402        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
13403        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
13404        // clamped layers; assert both so a future caller that skips the gate fails loudly.
13405        debug_assert!(
13406            cfg.sigmoid_router().is_none(),
13407            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
13408        );
13409        debug_assert!(
13410            !cfg.swiglu_clamped_at(il as u32),
13411            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
13412        );
13413
13414        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
13415        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
13416        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
13417        // skipped entirely for macro-free experts (every k-quant GGUF).
13418        if m.has_macros {
13419            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
13420        }
13421
13422        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
13423        let mut moe_out = e.uninit(t * n_embd)?;
13424
13425        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
13426        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
13427        if let Some(dev) = m.dev_exps.as_ref() {
13428            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
13429            // the combined stride; up's base is offset in the ptr table. Down unchanged.
13430            let (rbg_d, rbu_d) = if dev.gu_il {
13431                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
13432                (sxx, sxx)
13433            } else {
13434                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
13435            };
13436            let q8 = moe_q8_enabled_for_model(cfg, m);
13437            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
13438            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
13439            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
13440            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
13441            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
13442            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
13443            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
13444            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
13445            let rows_arm = q8
13446                && t > 1
13447                && crate::spec::spec_m2()
13448                && n_ff_exp == 512
13449                && n_used <= 8
13450                && std::env::var("MEMRA_MOE_DEVQ8_GU")
13451                    .map(|v| v.is_empty() || v == "v")
13452                    .unwrap_or(true)
13453                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
13454                    .map(|v| v.is_empty() || v == "w8h2v")
13455                    .unwrap_or(true);
13456            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
13457            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
13458            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
13459            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
13460            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
13461            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
13462            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
13463            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
13464            let csr_mode = std::env::var("MEMRA_MOE_CSR")
13465                .ok()
13466                .and_then(|v| v.parse::<i32>().ok())
13467                .unwrap_or(1);
13468            // NVFP4 admission REVERTED 2026-08-21 (lane/samplat, decode-batch-gate2 find):
13469            // the csr_nvfp4 kernel drifts last-ULP vs the rows program (11041/32768 ACT
13470            // elements at t=8) and the drift is BATCH-COMPOSITION-DEPENDENT — gate2 (B=8 vs
13471            // isolated) FAILED on the ornith15 artifact, the one-numeric-program law's batch
13472            // axis. Three chain-pinning attempts did not close it (receipts,
13473            // research/samplat-20260821/); a source-verbatim per-pair helper form IS
13474            // bit-identical but loses the dedup win (-3% vs rows). NVFP4 stays on the rows
13475            // twins until a cached form passes gate2 + the =2 byte-compare at t=8. The
13476            // increment-1 qualification hole: =2 ran across run-spec (solo verify shapes),
13477            // never decode-batch-gate at B=8 on the MoE model itself.
13478            // MEMRA_MOE_CSR_NVFP4=1 (lane/orndecode, DIAGNOSTIC PROBE ONLY): re-admits NVFP4
13479            // to the CSR arm and widens it to the exact-16 decode widths, so gate2 B=12/16 +
13480            // the =2 byte-compare can re-adjudicate the cached form at the widths where the
13481            // serial dev loop hurts most (B=16 tick: 1280 launches/step). The v0.100.1
13482            // de-admission verdict above stands until those gates are GREEN on the MoE
13483            // artifact; this door must never default on.
13484            let csr_nvfp4_probe = std::env::var("MEMRA_MOE_CSR_NVFP4").as_deref() == Ok("1");
13485            let csr_qt = |qt: i32| {
13486                qt == crate::QT_IQ4_XS
13487                    || qt == crate::QT_IQ3_S
13488                    || (csr_nvfp4_probe && qt == crate::QT_NVFP4)
13489            };
13490            let csr_t_max = if csr_nvfp4_probe { MOE_DEV_MAX_T } else { 10 };
13491            let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
13492            let csr_arm = rows_arm
13493                && csr_mode > 0
13494                && t <= csr_t_max
13495                && csr_uniform
13496                && csr_qt(m.gate_exps.qtype)
13497                && csr_qt(m.up_exps.qtype)
13498                && csr_qt(m.down_exps.qtype);
13499            if csr_arm {
13500                if csr_mode == 2 {
13501                    static ENGAGED: std::sync::Once = std::sync::Once::new();
13502                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
13503                }
13504                let n_pairs = t * n_used;
13505                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13506                let act = e.moe_gate_up_silu8_dev_q8_csr(
13507                    &dev.ptr_row,
13508                    &sel_d,
13509                    &zq,
13510                    &zd,
13511                    n_pairs,
13512                    n_embd,
13513                    n_ff_exp,
13514                    n_used,
13515                    n_expert,
13516                    m.gate_exps.qtype,
13517                    m.up_exps.qtype,
13518                    rbg_d,
13519                    rbu_d,
13520                )?;
13521                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13522                // down stays on the _rows twin — BOTH CSR down variants measured negative
13523                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
13524                // 16-group rows have too little decode to amortize any dedup structure.
13525                e.moe_down8_fma_dev_q8_rows(
13526                    &dev.ptr_row,
13527                    &sel_d,
13528                    &w_d,
13529                    &aq2,
13530                    &ad2,
13531                    &mut moe_out,
13532                    t,
13533                    n_ff_exp,
13534                    n_embd,
13535                    n_used,
13536                    n_expert,
13537                    m.down_exps.qtype,
13538                    m.down_exps.row_bytes,
13539                )?;
13540                if csr_mode == 2 {
13541                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
13542                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
13543                        &dev.ptr_row,
13544                        &sel_d,
13545                        &zq,
13546                        &zd,
13547                        t,
13548                        n_embd,
13549                        n_ff_exp,
13550                        n_used,
13551                        n_expert,
13552                        m.gate_exps.qtype,
13553                        m.up_exps.qtype,
13554                        rbg_d,
13555                        rbu_d,
13556                        &m.dev_macros,
13557                    )?;
13558                    let mut out_r = e.uninit(t * n_embd)?;
13559                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
13560                    e.moe_down8_fma_dev_q8_rows(
13561                        &dev.ptr_row,
13562                        &sel_d,
13563                        &w_d,
13564                        &aq2r,
13565                        &ad2r,
13566                        &mut out_r,
13567                        t,
13568                        n_ff_exp,
13569                        n_embd,
13570                        n_used,
13571                        n_expert,
13572                        m.down_exps.qtype,
13573                        m.down_exps.row_bytes,
13574                    )?;
13575                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
13576                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
13577                    let ba = a1
13578                        .iter()
13579                        .zip(&a2)
13580                        .filter(|(x, y)| x.to_bits() != y.to_bits())
13581                        .count();
13582                    let bo = o1
13583                        .iter()
13584                        .zip(&o2)
13585                        .filter(|(x, y)| x.to_bits() != y.to_bits())
13586                        .count();
13587                    if ba + bo > 0 {
13588                        eprintln!(
13589                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
13590                            a1.len(),
13591                            o1.len()
13592                        );
13593                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
13594                        let sel_h = e.dtoh_i32(&sel_d)?;
13595                        let mut shown = 0;
13596                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
13597                            if x.to_bits() != y.to_bits() && shown < 4 {
13598                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
13599                                let ex = sel_h[p];
13600                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
13601                                eprintln!(
13602                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
13603                                );
13604                                shown += 1;
13605                            }
13606                        }
13607                        std::process::exit(3);
13608                    }
13609                }
13610            } else if rows_arm {
13611                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
13612                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
13613                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
13614                    use std::sync::atomic::{AtomicU64, Ordering};
13615                    static PAIRS: AtomicU64 = AtomicU64::new(0);
13616                    static UNIQ: AtomicU64 = AtomicU64::new(0);
13617                    static CALLS: AtomicU64 = AtomicU64::new(0);
13618                    let sel_h = e.dtoh_i32(&sel_d)?;
13619                    let mut u: Vec<i32> = sel_h.clone();
13620                    u.sort_unstable();
13621                    u.dedup();
13622                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
13623                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
13624                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
13625                    if c.is_multiple_of(480) {
13626                        let p = PAIRS.load(Ordering::Relaxed);
13627                        let q = UNIQ.load(Ordering::Relaxed);
13628                        eprintln!(
13629                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
13630                            q as f64 / p as f64
13631                        );
13632                    }
13633                }
13634                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13635                let act = e.moe_gate_up_silu8_dev_q8_rows(
13636                    &dev.ptr_row,
13637                    &sel_d,
13638                    &zq,
13639                    &zd,
13640                    t,
13641                    n_embd,
13642                    n_ff_exp,
13643                    n_used,
13644                    n_expert,
13645                    m.gate_exps.qtype,
13646                    m.up_exps.qtype,
13647                    rbg_d,
13648                    rbu_d,
13649                    &m.dev_macros,
13650                )?;
13651                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
13652                e.moe_down8_fma_dev_q8_rows(
13653                    &dev.ptr_row,
13654                    &sel_d,
13655                    &w_d,
13656                    &aq2,
13657                    &ad2,
13658                    &mut moe_out,
13659                    t,
13660                    n_ff_exp,
13661                    n_embd,
13662                    n_used,
13663                    n_expert,
13664                    m.down_exps.qtype,
13665                    m.down_exps.row_bytes,
13666                )?;
13667            } else {
13668                for tok in 0..t {
13669                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13670                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13671                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13672                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13673                    if q8 {
13674                        let (zq, zd) = match (t, zq8) {
13675                            (1, Some((q, d))) => (q.clone(), d.clone()),
13676                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
13677                        };
13678                        let act = e.moe_gate_up_silu8_dev_q8(
13679                            &dev.ptr_row,
13680                            &selt,
13681                            &zq,
13682                            &zd,
13683                            n_embd,
13684                            n_ff_exp,
13685                            n_used,
13686                            n_expert,
13687                            m.gate_exps.qtype,
13688                            m.up_exps.qtype,
13689                            rbg_d,
13690                            rbu_d,
13691                            &m.dev_macros,
13692                        )?;
13693                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13694                        e.moe_down8_fma_dev_q8(
13695                            &dev.ptr_row,
13696                            &selt,
13697                            &wt,
13698                            &aq2,
13699                            &ad2,
13700                            &mut dst,
13701                            n_ff_exp,
13702                            n_embd,
13703                            n_used,
13704                            n_expert,
13705                            m.down_exps.qtype,
13706                            m.down_exps.row_bytes,
13707                        )?;
13708                    } else {
13709                        let act = e.moe_gate_up_silu8_dev(
13710                            &dev.ptr_row,
13711                            &selt,
13712                            &zt,
13713                            n_embd,
13714                            n_ff_exp,
13715                            n_used,
13716                            n_expert,
13717                            m.gate_exps.qtype,
13718                            m.up_exps.qtype,
13719                            rbg_d,
13720                            rbu_d,
13721                            &m.dev_macros,
13722                        )?;
13723                        e.moe_down8_fma_dev(
13724                            &dev.ptr_row,
13725                            &selt,
13726                            &wt,
13727                            &act,
13728                            &mut dst,
13729                            n_ff_exp,
13730                            n_embd,
13731                            n_used,
13732                            n_expert,
13733                            m.down_exps.qtype,
13734                            m.down_exps.row_bytes,
13735                        )?;
13736                    }
13737                }
13738            }
13739        } else {
13740            // Launch under the cache lock: the row borrow lives as long as the closure, and the
13741            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
13742            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
13743            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
13744            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
13745            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
13746            let q8 = moe_q8_enabled_for_model(cfg, m);
13747            e.with_moe_cache(max_block, |c, eng| {
13748                let row = c
13749                    .layer_dev_row(il, n_expert, eng)?
13750                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
13751                for tok in 0..t {
13752                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13753                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13754                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13755                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13756                    if q8 {
13757                        let (zq, zd) = match (t, zq8) {
13758                            (1, Some((q, d))) => (q.clone(), d.clone()),
13759                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
13760                        };
13761                        let act = eng.moe_gate_up_silu8_dev_q8(
13762                            row,
13763                            &selt,
13764                            &zq,
13765                            &zd,
13766                            n_embd,
13767                            n_ff_exp,
13768                            n_used,
13769                            n_expert,
13770                            m.gate_exps.qtype,
13771                            m.up_exps.qtype,
13772                            m.gate_exps.row_bytes,
13773                            m.up_exps.row_bytes,
13774                            &m.dev_macros,
13775                        )?;
13776                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
13777                        eng.moe_down8_fma_dev_q8(
13778                            row,
13779                            &selt,
13780                            &wt,
13781                            &aq2,
13782                            &ad2,
13783                            &mut dst,
13784                            n_ff_exp,
13785                            n_embd,
13786                            n_used,
13787                            n_expert,
13788                            m.down_exps.qtype,
13789                            m.down_exps.row_bytes,
13790                        )?;
13791                    } else {
13792                        let act = eng.moe_gate_up_silu8_dev(
13793                            row,
13794                            &selt,
13795                            &zt,
13796                            n_embd,
13797                            n_ff_exp,
13798                            n_used,
13799                            n_expert,
13800                            m.gate_exps.qtype,
13801                            m.up_exps.qtype,
13802                            m.gate_exps.row_bytes,
13803                            m.up_exps.row_bytes,
13804                            &m.dev_macros,
13805                        )?;
13806                        eng.moe_down8_fma_dev(
13807                            row,
13808                            &selt,
13809                            &wt,
13810                            &act,
13811                            &mut dst,
13812                            n_ff_exp,
13813                            n_embd,
13814                            n_used,
13815                            n_expert,
13816                            m.down_exps.qtype,
13817                            m.down_exps.row_bytes,
13818                        )?;
13819                    }
13820                }
13821                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
13822                c.hits += (t * 3 * n_used) as u64;
13823                Ok(())
13824            })?;
13825        }
13826
13827        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
13828        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
13829        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
13830        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
13831        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13832            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13833        {
13834            let n_ff_sh = gate_shexp.out_features();
13835            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
13836            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
13837            let verify_t = t > 1 && t < PRIME_MIN_T;
13838            let (sg_gate, sg_up) = if t == 1 {
13839                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
13840            } else if verify_t {
13841                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
13842                // rides one shared quantize + one fused2 batched launch instead of two
13843                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
13844                let mut fused = None;
13845                if crate::spec::spec_fused_t()
13846                    && (2..=4).contains(&t)
13847                    && e.uses_q8_1_fast(gate_shexp)
13848                    && e.uses_q8_1_fast(up_shexp)
13849                {
13850                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13851                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13852                }
13853                match fused {
13854                    Some(pair) => pair,
13855                    None => (
13856                        e.matmul_decode_exact(gate_shexp, z, t)?,
13857                        e.matmul_decode_exact(up_shexp, z, t)?,
13858                    ),
13859                }
13860            } else {
13861                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13862            };
13863            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
13864            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13865            let sh = if verify_t {
13866                e.matmul_decode_exact(down_shexp, &sa, t)?
13867            } else {
13868                e.matmul(down_shexp, &sa, t)?
13869            };
13870            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13871            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
13872            // between the two arms; prefill keeps the batched cuBLASLt linear).
13873            let g = match &m.gate_inp_shexp {
13874                Some(gate_inp_shexp) => {
13875                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
13876                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
13877                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
13878                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13879                    } else {
13880                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13881                        let mut g = e.uninit(t)?;
13882                        e.sigmoid(&gs, &mut g, t)?;
13883                        g
13884                    }
13885                }
13886                None => e.htod(&vec![1.0f32; t])?,
13887            };
13888            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13889        }
13890
13891        Ok(moe_out)
13892    }
13893
13894    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
13895    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
13896    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
13897    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
13898    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
13899    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
13900    /// the collected raw pointers cannot move between collection and launch (single-threaded
13901    /// decode; the lock is held only for collection, launches are stream-ordered after any
13902    /// prior same-stream staging writes).
13903    #[allow(clippy::too_many_arguments)]
13904    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
13905    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
13906    #[allow(clippy::too_many_arguments)]
13907    fn moe_gdec_token_q8(
13908        e: &Engine,
13909        m: &MoeWeights,
13910        il: u16,
13911        max_block: usize,
13912        zq: &CudaSlice<i8>,
13913        zd: &CudaSlice<f32>,
13914        sel: &[u32],
13915        w: &[f32],
13916        moe_out: &mut CudaSlice<f32>,
13917        tok: usize,
13918        n_embd: usize,
13919        n_ff_exp: usize,
13920        n_used: usize,
13921    ) -> Result<bool, Box<dyn std::error::Error>> {
13922        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
13923        use cudarc::driver::DevicePtr;
13924        let ptrs = e.with_moe_cache(max_block, |c, eng| {
13925            let mut g = [0u64; 8];
13926            let mut u = [0u64; 8];
13927            let mut d = [0u64; 8];
13928            for (j, &ex) in sel.iter().enumerate() {
13929                let ex = ex as u16;
13930                let (Some(sg), Some(su), Some(sd)) = (
13931                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
13932                    c.resident(BlockId::new(il, PROJ_UP, ex)),
13933                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
13934                ) else {
13935                    return Ok(None);
13936                };
13937                let __s = eng.stream();
13938                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
13939                let (pu, _e1) = c.slot(su).device_ptr(&__s);
13940                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
13941                g[j] = pg;
13942                u[j] = pu;
13943                d[j] = pd;
13944            }
13945            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
13946                for &ex in sel {
13947                    let ex = ex as u16;
13948                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
13949                        c.note_profile_hit(BlockId::new(il, proj, ex));
13950                    }
13951                }
13952            }
13953            c.hits += (3 * n_used) as u64;
13954            Ok(Some((g, u, d)))
13955        })?;
13956        let Some((g, u, d)) = ptrs else {
13957            return Ok(false);
13958        };
13959        let mut wv = [0f32; 8];
13960        wv[..n_used].copy_from_slice(w);
13961        let act = e.moe_gate_up_silu8_q8(
13962            crate::WPtr8(g),
13963            crate::WPtr8(u),
13964            zq,
13965            zd,
13966            n_embd,
13967            n_ff_exp,
13968            n_used,
13969            m.gate_exps.qtype,
13970            m.up_exps.qtype,
13971            m.gate_exps.row_bytes,
13972            m.up_exps.row_bytes,
13973        )?;
13974        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
13975        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13976        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13977        e.moe_down8_fma_q8(
13978            crate::WPtr8(d),
13979            crate::F32x8(wv),
13980            &aq2,
13981            &ad2,
13982            &mut dst,
13983            n_ff_exp,
13984            n_embd,
13985            n_used,
13986            m.down_exps.qtype,
13987            m.down_exps.row_bytes,
13988        )?;
13989        Ok(true)
13990    }
13991
13992    /// glm5_next's fused MoE epilogue for ONE token-layer: the sigmoid router's already-selected
13993    /// `(sel, w)`, the PRE-clamped SwiGLU and the per-expert NVFP4 macro fold, in one launch
13994    /// pair. Returns `false` when the cache cannot hold this token's `3*n_used` blocks at once,
13995    /// in which case the caller must run the sequential loop (which zeroes its own row).
13996    ///
13997    /// WHY IT DOES NOT NEED A RESIDENT LAYER, unlike `moe_gdec_token_q8`. gdec collects pointers
13998    /// from blocks that are ALREADY resident and bails on the first miss, because a miss would
13999    /// mean an admission that could move a slot under the pointers it has already taken. This arm
14000    /// inverts the order: it ADMITS all `3*n_used` blocks first, through the same
14001    /// `dispatch_source` the sequential loop calls per projection (a hit copies nothing, a miss
14002    /// runs the identical `memcpy_htod` into a slot), and only then takes the addresses, in a
14003    /// second pass, with the cache lock still held. Nothing can move between the last admission
14004    /// and the pointer read, and the kernels are issued on the compute stream immediately after —
14005    /// the same in-order guarantee the sequential loop already relies on when it dispatches
14006    /// expert j+1 after launching expert j's kernels.
14007    ///
14008    /// The slot-capacity check is the fail-closed seam: `n_slots()` is a whole-cache bound, and
14009    /// with fewer than `3*n_used` slots an admission is guaranteed to evict one of this token's
14010    /// own blocks. The second pass re-reads `resident()` for every block rather than trusting the
14011    /// dispatch's return, so an eviction the capacity check did not predict falls through loudly
14012    /// to the sequential loop instead of dereferencing a reused slot.
14013    ///
14014    /// BIT-IDENTITY CLASS. Against the sequential loop this arm is a DISPATCH-class change, not a
14015    /// provenance one: the same block bytes and the same macro scales, but the gate/up dots are
14016    /// the fused kernel's warp reduction rather than `qmatvec_expert_q8`'s per-projection one, and
14017    /// the down accumulation is `moe_down8_fma_q8`'s slot-ordered `__fmaf_rn` chain rather than 8
14018    /// separate `axpy_into` calls. Those chains are the ones the gdec family documents as
14019    /// reproducing the sequential chain exactly; `tests/glm5_moe_epilogue_gpu.rs::the_two_arms_agree`
14020    /// measures the actual bit disagreement rather than asserting the claim.
14021    #[allow(clippy::too_many_arguments)]
14022    fn moe_fused_epi_token_q8(
14023        e: &Engine,
14024        m: &MoeWeights,
14025        il: u16,
14026        max_block: usize,
14027        zq: &CudaSlice<i8>,
14028        zd: &CudaSlice<f32>,
14029        sel: &[u32],
14030        w: &[f32],
14031        moe_out: &mut CudaSlice<f32>,
14032        tok: usize,
14033        n_embd: usize,
14034        n_ff_exp: usize,
14035        n_used: usize,
14036        limit: f32,
14037    ) -> Result<bool, Box<dyn std::error::Error>> {
14038        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14039        use cudarc::driver::DevicePtr;
14040        debug_assert!(
14041            limit > 1e-6,
14042            "the fused epilogue's kernel collapses every gate to silu(0) at limit 0"
14043        );
14044        debug_assert_eq!(sel.len(), n_used);
14045        debug_assert_eq!(w.len(), n_used);
14046
14047        let ptrs = e.with_moe_cache(max_block, |c, eng| {
14048            // Fail closed: below this bound an admission MUST evict one of this token's own
14049            // blocks, so there is no pointer set that stays valid.
14050            if c.n_slots() < 3 * n_used {
14051                return Ok(None);
14052            }
14053            // PASS 1 — admit. Identical dispatch to the sequential loop's `moe_cached_gemm_q8`,
14054            // projection for projection; only the GEMM is deferred.
14055            for &ex in sel.iter() {
14056                let ex_usize = ex as usize;
14057                for (proj, exps) in [
14058                    (PROJ_GATE, &m.gate_exps),
14059                    (PROJ_UP, &m.up_exps),
14060                    (PROJ_DOWN, &m.down_exps),
14061                ] {
14062                    let id = BlockId::new(il, proj, ex as u16);
14063                    let DispatchSlot::Resident(_) =
14064                        c.dispatch_source(id, exps.expert_source(ex_usize), eng)?;
14065                }
14066            }
14067            // PASS 2 — take the fixed slot addresses, with nothing left to admit.
14068            let mut g = [0u64; 8];
14069            let mut u = [0u64; 8];
14070            let mut d = [0u64; 8];
14071            for (j, &ex) in sel.iter().enumerate() {
14072                let ex = ex as u16;
14073                let (Some(sg), Some(su), Some(sd)) = (
14074                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
14075                    c.resident(BlockId::new(il, PROJ_UP, ex)),
14076                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
14077                ) else {
14078                    return Ok(None);
14079                };
14080                let __s = eng.stream();
14081                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
14082                let (pu, _e1) = c.slot(su).device_ptr(&__s);
14083                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
14084                g[j] = pg;
14085                u[j] = pu;
14086                d[j] = pd;
14087            }
14088            Ok(Some((g, u, d)))
14089        })?;
14090        let Some((g, u, d)) = ptrs else {
14091            return Ok(false);
14092        };
14093
14094        Self::moe_fused_epi_launch(
14095            e, m, zq, zd, sel, w, g, u, d, moe_out, tok, n_embd, n_ff_exp, n_used, limit,
14096        )?;
14097        Ok(true)
14098    }
14099
14100    /// The fused epilogue's ONLY launch path, shared by both provenances (SLRU slot addresses and
14101    /// device-resident slab base+stride). Everything that could differ semantically between them
14102    /// — the per-expert macro fold, the clamp, the kernel pair, the dispatch counter — lives here
14103    /// exactly once, so the two arms cannot drift into being different programs. The callers
14104    /// differ only in how they filled `g`/`u`/`d`.
14105    ///
14106    /// Per-expert macro scales in router slot order: gate/up ride the kernel's epilogue exactly
14107    /// where `ffn_act_lim`'s gs/us ride the unfused loop, and down folds into the routing weight
14108    /// exactly where `axpy_into`'s `w[j] * macro_scale(ex)` folds it. `macro_scale` answers 1.0
14109    /// for a macro-free bank, so a k-quant GGUF takes this path with no fold and no branch.
14110    #[allow(clippy::too_many_arguments)]
14111    fn moe_fused_epi_launch(
14112        e: &Engine,
14113        m: &MoeWeights,
14114        zq: &CudaSlice<i8>,
14115        zd: &CudaSlice<f32>,
14116        sel: &[u32],
14117        w: &[f32],
14118        g: [u64; 8],
14119        u: [u64; 8],
14120        d: [u64; 8],
14121        moe_out: &mut CudaSlice<f32>,
14122        tok: usize,
14123        n_embd: usize,
14124        n_ff_exp: usize,
14125        n_used: usize,
14126        limit: f32,
14127    ) -> Result<(), Box<dyn std::error::Error>> {
14128        let mut gs = [0f32; 8];
14129        let mut us = [0f32; 8];
14130        let mut wv = [0f32; 8];
14131        for (j, &ex) in sel.iter().enumerate() {
14132            let ex = ex as usize;
14133            gs[j] = m.gate_exps.macro_scale(ex);
14134            us[j] = m.up_exps.macro_scale(ex);
14135            wv[j] = w[j] * m.down_exps.macro_scale(ex);
14136        }
14137        let act = e.moe_gate_up_preclamp8_q8(
14138            crate::WPtr8(g),
14139            crate::WPtr8(u),
14140            zq,
14141            zd,
14142            crate::F32x8(gs),
14143            crate::F32x8(us),
14144            limit,
14145            n_embd,
14146            n_ff_exp,
14147            n_used,
14148            m.gate_exps.qtype,
14149            m.up_exps.qtype,
14150            m.gate_exps.row_bytes,
14151            m.up_exps.row_bytes,
14152        )?;
14153        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
14154        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
14155        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14156        e.moe_down8_fma_q8(
14157            crate::WPtr8(d),
14158            crate::F32x8(wv),
14159            &aq2,
14160            &ad2,
14161            &mut dst,
14162            n_ff_exp,
14163            n_embd,
14164            n_used,
14165            m.down_exps.qtype,
14166            m.down_exps.row_bytes,
14167        )?;
14168        crate::MOE_FUSED_EPI_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14169        Ok(())
14170    }
14171
14172    /// The verify-rows batched routed-expert program (lane/glm5-vrest): one layer-call's
14173    /// WHOLE t x n_used pair union through the fused-epilogue kernels' rows twins —
14174    /// one gate/up+preclamp launch, one pair-major activation quantize, one down+FMA
14175    /// launch — with pointers computed from the resident slab base + ex*stride (the
14176    /// sequential slab arm's exact arithmetic) and the macro folds landing exactly where
14177    /// `ffn_act_lim` / `axpy_into` land them. `moe_out` rows are FULLY overwritten.
14178    #[allow(clippy::too_many_arguments)]
14179    // allow: the parameter list mirrors its dispatch-arm caller's contract
14180    fn moe_vrows_pairs_q8(
14181        e: &Engine,
14182        m: &MoeWeights,
14183        z: &CudaSlice<f32>,
14184        sel: VrowsSel<'_>,
14185        il: u16,
14186        (pg, pu, pd): (u64, u64, u64),
14187        t: usize,
14188        n_embd: usize,
14189        n_ff_exp: usize,
14190        n_used: usize,
14191        limit: f32,
14192        moe_out: &mut CudaSlice<f32>,
14193    ) -> Result<(), Box<dyn std::error::Error>> {
14194        let n_pairs = t * n_used;
14195        // Door E (MEMRA_MOE_VROWS_DEDUP_ORDER, default OFF): the gate/up launch walks the pair
14196        // union EXPERT-MAJOR, reading the visit order from a FOURTH plane appended to the pointer
14197        // table (planes: gate | up | down | order). Carrying it in the existing table is what
14198        // makes the door free on the host arm — the order plane rides the single `htod_u64_into`
14199        // that was already uploading the pointers, so no new transfer and no new pool appear.
14200        // Door M (`MEMRA_MOE_VROWS_PACK`) refuses it in the launcher, so do not build the plane
14201        // when the refuted pack door is armed.
14202        let order_on = crate::moe_vrows_dedup_order_on() && !crate::moe_vrows_pack_on();
14203        let n_planes = if order_on { 4 } else { 3 };
14204        // Door W (MEMRA_GLM5_VERIFY_WS): the whole staging set — tables, token quantize,
14205        // act, pair quantize — draws from the verify workspace and recycles at the end of
14206        // the call (vws_* are alloc_uninit/plain-drop with the door off, so the OFF arm is
14207        // byte-for-byte the shipped program). Every buffer is fully overwritten before any
14208        // read by the SAME kernels (the sites' standing uninit contract).
14209        let mut ptrs_d = e.vws_uninit_u64(n_planes * n_pairs)?;
14210        let mut scl_d = e.vws_uninit(3 * n_pairs)?;
14211        // ONE launch path, TWO table provenances (the fused-epilogue arm's own discipline):
14212        // only the plane-major (gate | up | down) pointer/scale tables are built differently,
14213        // and door D's kernel evaluates the SAME terms as the host loop, so nothing downstream
14214        // can tell the arms apart. See the `moe_vrows_tables_from_sel` kernel comment for the
14215        // term-by-term bit-identity argument.
14216        match sel {
14217            VrowsSel::Host(sel_all, w_all) => {
14218                debug_assert_eq!(sel_all.len(), n_pairs);
14219                debug_assert_eq!(w_all.len(), n_pairs);
14220                let mut ptrs = vec![0u64; n_planes * n_pairs];
14221                let mut scl = vec![0f32; 3 * n_pairs];
14222                for (p, (&ex, &w)) in sel_all.iter().zip(w_all).enumerate() {
14223                    let ex = ex as usize;
14224                    ptrs[p] = pg + (ex * m.gate_exps.expert_stride) as u64;
14225                    ptrs[n_pairs + p] = pu + (ex * m.up_exps.expert_stride) as u64;
14226                    ptrs[2 * n_pairs + p] = pd + (ex * m.down_exps.expert_stride) as u64;
14227                    scl[p] = m.gate_exps.macro_scale(ex);
14228                    scl[n_pairs + p] = m.up_exps.macro_scale(ex);
14229                    // down-proj macro folds into the accumulate weight (1.0 for non-macro
14230                    // banks) — the axpy_into fold, verbatim.
14231                    scl[2 * n_pairs + p] = w * m.down_exps.macro_scale(ex);
14232                }
14233                if order_on {
14234                    // The order plane rides the SAME upload — the door adds no HtoD on this arm.
14235                    ptrs[3 * n_pairs..].copy_from_slice(&crate::vrows_expert_major_order(sel_all));
14236                    // The box receipt: the slab reads whose repeat visit this schedule places
14237                    // inside the reuse window (host arm only — see MOE_VROWS_SLAB_READS_AVOIDED).
14238                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14239                    crate::MOE_VROWS_SLAB_READS_AVOIDED
14240                        .fetch_add(visits - distinct, std::sync::atomic::Ordering::Relaxed);
14241                }
14242                e.htod_u64_into(&ptrs, &mut ptrs_d)?;
14243                e.htod_f32_into(&scl, &mut scl_d)?;
14244                // MEMRA_MOE_VROWS_DEDUP_STAT: size the ONLY remaining byte lever on this pair
14245                // (LANE.md §1 — it already runs at ~90% of theoretical DRAM peak, so the sole
14246                // way to cut it further is reading a shared expert slab once for the rows that
14247                // share it). `1 - distinct/visits` IS that lever; measuring it costs a bitset.
14248                if crate::moe_vrows_dedup_stat_on() {
14249                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14250                    debug_assert_eq!(visits, n_pairs as u64);
14251                    crate::MOE_VROWS_PAIR_VISITS
14252                        .fetch_add(visits, std::sync::atomic::Ordering::Relaxed);
14253                    crate::MOE_VROWS_PAIR_DISTINCT
14254                        .fetch_add(distinct, std::sync::atomic::Ordering::Relaxed);
14255                    crate::moe_vrows_dedup_report();
14256                }
14257            }
14258            VrowsSel::Dev(sel_d, selw_d) => {
14259                let macros = match (
14260                    m.gate_exps.macros.as_deref(),
14261                    m.up_exps.macros.as_deref(),
14262                    m.down_exps.macros.as_deref(),
14263                ) {
14264                    (Some(g), Some(u), Some(d)) => Some((g, u, d)),
14265                    // A partially-macro bank would need per-plane 1.0 defaults the kernel does
14266                    // not carry; the serving artifact's three planes are all present or all
14267                    // absent, so refuse rather than guess.
14268                    (None, None, None) => None,
14269                    _ => {
14270                        return Err("vrows device tables: expert macro planes are not uniform \
14271                                    across gate/up/down"
14272                            .into());
14273                    }
14274                };
14275                e.moe_vrows_tables_from_sel(
14276                    sel_d,
14277                    selw_d,
14278                    il,
14279                    macros,
14280                    (pg, pu, pd),
14281                    (
14282                        m.gate_exps.expert_stride,
14283                        m.up_exps.expert_stride,
14284                        m.down_exps.expert_stride,
14285                    ),
14286                    n_pairs,
14287                    &mut ptrs_d,
14288                    &mut scl_d,
14289                )?;
14290                if order_on {
14291                    // Door E on the device arm: one extra launch (the host arm gets the plane for
14292                    // free inside its existing upload). Bit-identical to the host's stable sort by
14293                    // (expert id, pair index) — gated directly against it in glm5_dedup_sched_gpu.
14294                    e.moe_vrows_order_from_sel(sel_d, n_pairs, &mut ptrs_d)?;
14295                }
14296                if crate::MOE_VROWS_DEV_TABLES_DISPATCHES
14297                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
14298                    == 0
14299                {
14300                    eprintln!(
14301                        "[moe-vrows-dev-tables] engaged: pointer/scale tables built on device \
14302                         from the router's own sel/w; the per-layer pinned readback and its \
14303                         cuStreamSynchronize are skipped (MEMRA_MOE_VROWS_DEV_TABLES=1)"
14304                    );
14305                }
14306            }
14307        }
14308        // Token rows quantized in one launch; per-row q8_1 bytes are position-independent
14309        // (the batched-MMVQ class), bit-gated against the per-token quantize_q8_1_view.
14310        let (mut zq, mut zd) = (
14311            e.vws_uninit_i8(t * n_embd)?,
14312            e.vws_uninit(t * (n_embd / 32))?,
14313        );
14314        e.quantize_q8_1_into(z, t, n_embd, &mut zq, &mut zd)?;
14315        let act = e.moe_gate_up_preclamp8_q8_rows(
14316            &ptrs_d,
14317            &scl_d,
14318            &zq,
14319            &zd,
14320            limit,
14321            n_embd,
14322            n_ff_exp,
14323            n_used,
14324            n_pairs,
14325            m.gate_exps.qtype,
14326            m.up_exps.qtype,
14327            m.gate_exps.row_bytes,
14328            m.up_exps.row_bytes,
14329        )?;
14330        // Pair-major activation quantize: [n_pairs, n_ff] rows in one launch.
14331        let (mut aq2, mut ad2) = (
14332            e.vws_uninit_i8(n_pairs * n_ff_exp)?,
14333            e.vws_uninit(n_pairs * (n_ff_exp / 32))?,
14334        );
14335        e.quantize_q8_1_into(&act, n_pairs, n_ff_exp, &mut aq2, &mut ad2)?;
14336        e.moe_down8_fma_q8_rows(
14337            &ptrs_d,
14338            &scl_d,
14339            &aq2,
14340            &ad2,
14341            moe_out,
14342            n_ff_exp,
14343            n_embd,
14344            n_used,
14345            n_pairs,
14346            m.down_exps.qtype,
14347            m.down_exps.row_bytes,
14348        )?;
14349        // Everything above is dead after the down launch (stream-ordered reuse is safe on
14350        // this engine's stream, the same guarantee the async free relies on).
14351        e.vws_recycle_u64(ptrs_d);
14352        e.vws_recycle(scl_d);
14353        e.vws_recycle_i8(zq);
14354        e.vws_recycle(zd);
14355        e.vws_recycle(act);
14356        e.vws_recycle_i8(aq2);
14357        e.vws_recycle(ad2);
14358        if crate::MOE_VROWS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
14359            eprintln!(
14360                "[glm5-vrows] verify MoE batched across rows: pairs={n_pairs} (t={t} x \
14361                 {n_used}), one gate/up+preclamp launch + one down/FMA launch per layer-call \
14362                 (rides MEMRA_GLM5_VERIFY_BATCH)"
14363            );
14364        }
14365        Ok(())
14366    }
14367
14368    /// GROUPED MoE PREFILL for the sigmoid-router glm5_next class (`MEMRA_MOE_GROUPED_PREFILL`,
14369    /// default ON since 2026-08-29, `=0` rollback). One call covers a whole prefill chunk's
14370    /// routed-expert FFN for one layer:
14371    ///
14372    ///   1. ROUTER: the SAME m-invariant `moe_router_logits` + `moe_route_sigmoid_cfg` host
14373    ///      oracle invocation the sequential arm makes, so selected experts and routing weights
14374    ///      are BIT-identical to the sequential arm by construction. Only the GEMM accumulation
14375    ///      order may move (the grouped GEMM is measured non-bit-stable,
14376    ///      `run_tensor_parallel_routes_nvfp4_prime_grouped`'s MEMRA_MOE_DETERM note), which is
14377    ///      why the acceptance gate is reference-band + routing-exactness, not byte identity.
14378    ///   2. TOKEN-SORT BY EXPERT: host counting sort of the (token, expert) pairs into an
14379    ///      expert-major CSR (vLLM's `moe_align_block_size` shape; same O(pairs) build the
14380    ///      softmax `moe_ffn_pairs` arm and the step37 grouped prime use).
14381    ///   3. ONE GROUPED TENSOR-CORE GEMM PER PROJECTION over the resident NVFP4 slab
14382    ///      (`moe_f16_grouped`, the sk single-kernel visitor with the NVFP4 direct tile
14383    ///      loaders; the step37 grouped-prime kernel class, 170-270 TFLOP/s on its lane's
14384    ///      sizing rows, generalized off the TP runtime to the single-device `dev_exps`
14385    ///      pointer-table provenance). Each expert's weights stream through tensor cores ONCE
14386    ///      per layer per chunk instead of once per (token, expert): at t=4096 that replaces
14387    ///      the sequential loop's 49 launches x 4096 tokens (~200k launches and ~113 MB x 4096
14388    ///      of expert VRAM re-reads per layer) with a ~15-launch chunk-wide program.
14389    ///   4. EPILOGUE: glm5_next's PRE-clamped SwiGLU `silu(min(g,l)) * clamp(u,±l)` with the
14390    ///      per-expert `weight_scale_2` macro fold: gate/up macros land BEFORE the nonlinearity
14391    ///      (`scale_rows` per CSR row; silu is nonlinear, so the fold cannot commute past it),
14392    ///      down macros fold into the scatter weight, exactly where the fused epilogue and the
14393    ///      sequential loop's `ffn_act_lim`/`axpy_into` put them.
14394    ///   5. SCATTER: permute CSR rows back to pair order, then the slot-ordered weighted
14395    ///      per-token accumulation (`moe_pairs_scatter`, the sequential-axpy accumulation
14396    ///      class). Shared expert rides the canonical clamp-aware grouped add.
14397    ///
14398    /// Returns `Ok(None)` (fail closed to the sequential arm) for every unqualified shape:
14399    /// no local resident slab, f16g door off, a projection the grouped kernel cannot walk, or
14400    /// an expert count past the sk visitor's group cap. Numeric class: f16-mirror activations
14401    /// (`moe_f16g_act` row-normalized f16), the class the softmax pairs f16g arm and the
14402    /// step37 grouped prime already serve prefill with; gated by
14403    /// `tests/glm5_moe_grouped_prefill_gpu.rs` against `memra_reference` at the fused-epilogue
14404    /// gate's tolerance class, plus the run-gen first-token argmax gate on real prompts.
14405    fn moe_ffn_grouped_prefill_sigmoid(
14406        e: &Engine,
14407        m: &MoeWeights,
14408        z: &CudaSlice<f32>,
14409        t: usize,
14410        cfg: &ModelConfig,
14411        il: u16,
14412    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14413        // Placement: the LOCAL resident slab only. The SLRU cannot serve a 4096-token chunk's
14414        // expert working set (a glm5 layer is 288 x 3 blocks against ~285 slots/layer on the
14415        // serving recipe), and a remote slab must never be dereferenced (m=1 peer reads are the
14416        // measured 34-150x class). Fail closed: the sequential arm stages as before.
14417        let Some(dev) = m
14418            .dev_exps
14419            .as_ref()
14420            .filter(|d| moe_slab_enabled() && d.dev == e.ctx().ordinal())
14421        else {
14422            return Ok(None);
14423        };
14424        if crate::moe_f16g_mode() == 0 {
14425            return Ok(None);
14426        }
14427        // MEMRA_MOE_GATE is the BYTE-identity oracle between sequential-class dispatches; this
14428        // arm is a different numeric class with its own reference-band gate, so it must not
14429        // shadow that comparison.
14430        if std::env::var("MEMRA_MOE_GATE").is_ok() {
14431            return Ok(None);
14432        }
14433        let moe = cfg
14434            .moe
14435            .as_ref()
14436            .ok_or("grouped sigmoid prefill requires MoE model metadata")?;
14437        let n_embd = cfg.n_embd as usize;
14438        let n_expert = moe.expert_count as usize;
14439        let n_used = moe.expert_used_count as usize;
14440        let n_ff_exp = moe.expert_ff_length as usize;
14441        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
14442            && f16g_proj_ok(m.up_exps.qtype, n_embd)
14443            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
14444        {
14445            return Ok(None);
14446        }
14447        // The sk visitor's direct-lane group cap (mirrors the grouped prime's guard). glm5's
14448        // 288 experts fit; a wider bank falls closed rather than erring mid-forward.
14449        if n_expert > 512 || n_used == 0 || n_used > 8 {
14450            return Ok(None);
14451        }
14452        let sigmoid = cfg
14453            .sigmoid_router()
14454            .ok_or("grouped sigmoid prefill requires the sigmoid router")?;
14455        // glm5_next carries the PRE form on every clamped layer (`clamp_exp_at`); a POST-form
14456        // arch reaching this arm is an unqualified semantic program, and the clamp-form law
14457        // says no dispatch site may pick a form by default. Err, not assert.
14458        let lim_exp = cfg.clamp_exp_at(il as u32);
14459        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
14460            return Err(
14461                "grouped sigmoid prefill is qualified for the PRE-clamped SwiGLU form only; \
14462                 a POST-clamp layer must ride the sequential arm"
14463                    .into(),
14464            );
14465        }
14466
14467        // 1. ROUTER. One selector shared with the sequential arm so changing dispatch cannot
14468        // change logits, selected expert ids, or routing weights (the routing-exactness half of
14469        // the acceptance gate holds by construction). The host readback here is the same one
14470        // the sequential arm performs; killing it is the L4 host-sync diet, not this arm.
14471        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
14472        Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
14473        let (sel_all, w_all) =
14474            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
14475        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
14476        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
14477        Self::trace_moe_input(e, il, t, n_embd, z)?;
14478
14479        let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
14480        let mut mt = std::time::Instant::now();
14481        let mut phase = |on: bool| -> f64 {
14482            if on {
14483                let _ = e.stream().synchronize();
14484                let v = mt.elapsed().as_secs_f64() * 1e3;
14485                mt = std::time::Instant::now();
14486                v
14487            } else {
14488                0.0
14489            }
14490        };
14491        let d_router = phase(mprof);
14492
14493        // 2. TOKEN-SORT BY EXPERT: expert-major CSR over the (token, expert) pairs.
14494        let n_pairs = t * n_used;
14495        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
14496            return Err("grouped sigmoid prefill geometry".into());
14497        }
14498        let mut buckets: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14499        for (p, &s_id) in sel_all.iter().take(n_pairs).enumerate() {
14500            let s_id = s_id as usize;
14501            if s_id >= n_expert {
14502                return Err(format!("grouped prefill selection {s_id} >= {n_expert}").into());
14503            }
14504            buckets[s_id].push(p as i32);
14505        }
14506        let mut ex_ids: Vec<i32> = Vec::new();
14507        let mut ex_off: Vec<i32> = vec![0];
14508        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
14509        for (e_id, b) in buckets.iter().enumerate() {
14510            if !b.is_empty() {
14511                ex_ids.push(e_id as i32);
14512                ex_pairs.extend_from_slice(b);
14513                ex_off.push(ex_pairs.len() as i32);
14514            }
14515        }
14516        let n_active = ex_ids.len();
14517        if n_active == 0 || n_active > 512 {
14518            return Err(format!("grouped prefill n_active {n_active} outside 1..=512").into());
14519        }
14520        let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
14521
14522        // 3. MACRO PLANES. Gate/up per-CSR-row scales (before silu); down folds into the
14523        // scatter weight; `macro_scale` answers 1.0 for a macro-free bank, so the fold is
14524        // skipped rather than launched as a no-op.
14525        let wd: Vec<f32> = (0..n_pairs)
14526            .map(|p| w_all[p] * m.down_exps.macro_scale(sel_all[p] as usize))
14527            .collect();
14528
14529        // Interleaved gate/up slab strides (see moe_ffn_pairs / moe_ffn_dev).
14530        let (rbg_d, rbu_d) = if dev.gu_il {
14531            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14532            (sxx, sxx)
14533        } else {
14534            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14535        };
14536
14537        let exi = e.htod_i32(&ex_ids)?;
14538        let exo = e.htod_i32(&ex_off)?;
14539        let exp_d = e.htod_i32(&ex_pairs)?;
14540        let csr_tok_d = e.htod_i32(&csr_tok)?;
14541        let pw = e.htod(&wd)?;
14542
14543        // 4. GATE/UP grouped GEMMs over the bank, CSR order end-to-end.
14544        let (z16, zs) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
14545        let mut g = e.moe_f16_grouped(
14546            &dev.ptr_row,
14547            0,
14548            n_expert,
14549            &exi,
14550            &ex_off,
14551            &exo,
14552            &z16,
14553            &zs,
14554            n_embd,
14555            n_ff_exp,
14556            n_active,
14557            n_pairs,
14558            m.gate_exps.qtype,
14559            rbg_d,
14560        )?;
14561        if m.gate_exps.macros.is_some() {
14562            let mg: Vec<f32> = ex_pairs
14563                .iter()
14564                .map(|&p| m.gate_exps.macro_scale(sel_all[p as usize] as usize))
14565                .collect();
14566            let mg_d = e.htod(&mg)?;
14567            e.scale_rows(&mut g, &mg_d, n_ff_exp, n_pairs)?;
14568        }
14569        let mut u = e.moe_f16_grouped(
14570            &dev.ptr_row,
14571            1,
14572            n_expert,
14573            &exi,
14574            &ex_off,
14575            &exo,
14576            &z16,
14577            &zs,
14578            n_embd,
14579            n_ff_exp,
14580            n_active,
14581            n_pairs,
14582            m.up_exps.qtype,
14583            rbu_d,
14584        )?;
14585        if m.up_exps.macros.is_some() {
14586            let mu: Vec<f32> = ex_pairs
14587                .iter()
14588                .map(|&p| m.up_exps.macro_scale(sel_all[p as usize] as usize))
14589                .collect();
14590            let mu_d = e.htod(&mu)?;
14591            e.scale_rows(&mut u, &mu_d, n_ff_exp, n_pairs)?;
14592        }
14593
14594        // 5. EPILOGUE: glm5_next's PRE-clamped SwiGLU (the POST form was refused above); a
14595        // config with no live limit takes the plain-silu pair kernel.
14596        let act = match lim_exp {
14597            Some(SwigluClamp::Pre(limit)) => {
14598                let mut a = e.uninit(n_pairs * n_ff_exp)?;
14599                // Scales are 1.0: the per-expert macros already landed via scale_rows (an
14600                // exact *1.0 inside the kernel keeps the value chain unchanged).
14601                e.swiglu_preclamped_mul_scaled(
14602                    &g,
14603                    &u,
14604                    1.0,
14605                    1.0,
14606                    limit,
14607                    &mut a,
14608                    n_pairs * n_ff_exp,
14609                )?;
14610                a
14611            }
14612            None => e.moe_pairs_silu_mul(&g, &u, n_pairs * n_ff_exp)?,
14613            Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
14614        };
14615        let d_gemm_gu = phase(mprof);
14616
14617        // 6. DOWN grouped GEMM (CSR order), permute back to pair order, weighted scatter.
14618        let (a16, a_s) = e.moe_f16g_act(&act, None, n_ff_exp, n_pairs)?;
14619        let d_csr = e.moe_f16_grouped(
14620            &dev.ptr_row,
14621            2,
14622            n_expert,
14623            &exi,
14624            &ex_off,
14625            &exo,
14626            &a16,
14627            &a_s,
14628            n_ff_exp,
14629            n_embd,
14630            n_active,
14631            n_pairs,
14632            m.down_exps.qtype,
14633            m.down_exps.row_bytes,
14634        )?;
14635        let y_pair = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
14636        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14637        let tids: Vec<i32> = (0..n_pairs as i32).collect();
14638        let toff_d = e.htod_i32(&toff)?;
14639        let tids_d = e.htod_i32(&tids)?;
14640        // The scatter fully overwrites every (token, col): slot-ordered accumulation over the
14641        // token's n_used pairs, the sequential-axpy class.
14642        let mut moe_out = e.uninit(t * n_embd)?;
14643        e.moe_pairs_scatter(&y_pair, &pw, &toff_d, &tids_d, &mut moe_out, t, n_embd)?;
14644        let d_down = phase(mprof);
14645
14646        // 7. SHARED EXPERT: the canonical clamp-aware grouped add (reads clamp_shexp_at and
14647        // the optional shexp gate; glm5_next has a live shared expert on every MoE layer).
14648        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
14649        if mprof {
14650            let d_shared = phase(true);
14651            eprintln!(
14652                "[moe-grouped-prefill-prof] il={il} t={t} router={d_router:.1}ms \
14653                 gemm_gu={d_gemm_gu:.1}ms down_scatter={d_down:.1}ms shared={d_shared:.1}ms"
14654            );
14655        }
14656
14657        crate::MOE_GROUPED_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14658        // Once per layer per process: the engagement receipt line (the A/B greps for it; the
14659        // both-arms flag announce lives at the dispatch site).
14660        static GPF_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
14661        let layer_bit = 1u64 << (il as u64 % 64);
14662        if GPF_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit == 0 {
14663            eprintln!(
14664                "[moe-grouped-prefill] execute layer={il} tokens={t} n_active={n_active} \
14665                 provenance=resident-slab router=sigmoid-host-oracle epilogue=pre-clamped \
14666                 macro_fold=gate-up-rows+down-weight performance_claim=false \
14667                 (logged once per layer)"
14668            );
14669        }
14670        Ok(Some(moe_out))
14671    }
14672
14673    #[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
14674    fn moe_gdec_token(
14675        e: &Engine,
14676        m: &MoeWeights,
14677        il: u16,
14678        max_block: usize,
14679        zt: &cudarc::driver::CudaView<f32>,
14680        sel: &[u32],
14681        w: &[f32],
14682        moe_out: &mut CudaSlice<f32>,
14683        tok: usize,
14684        n_embd: usize,
14685        n_ff_exp: usize,
14686        n_used: usize,
14687    ) -> Result<bool, Box<dyn std::error::Error>> {
14688        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14689        use cudarc::driver::DevicePtr;
14690        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
14691        let ptrs = e.with_moe_cache(max_block, |c, eng| {
14692            let mut g = [0u64; 8];
14693            let mut u = [0u64; 8];
14694            let mut d = [0u64; 8];
14695            for (j, &ex) in sel.iter().enumerate() {
14696                let ex = ex as u16;
14697                let (Some(sg), Some(su), Some(sd)) = (
14698                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
14699                    c.resident(BlockId::new(il, PROJ_UP, ex)),
14700                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
14701                ) else {
14702                    return Ok(None);
14703                };
14704                let __s = eng.stream();
14705                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
14706                let (pu, _e1) = c.slot(su).device_ptr(&__s);
14707                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
14708                g[j] = pg;
14709                u[j] = pu;
14710                d[j] = pd;
14711            }
14712            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
14713                for &ex in sel {
14714                    let ex = ex as u16;
14715                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
14716                        c.note_profile_hit(BlockId::new(il, proj, ex));
14717                    }
14718                }
14719            }
14720            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
14721            Ok(Some((g, u, d)))
14722        })?;
14723        let Some((g, u, d)) = ptrs else {
14724            return Ok(false);
14725        };
14726        let mut wv = [0f32; 8];
14727        wv[..n_used].copy_from_slice(w);
14728        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
14729        let act = e.moe_gate_up_silu8(
14730            crate::WPtr8(g),
14731            crate::WPtr8(u),
14732            zt,
14733            n_embd,
14734            n_ff_exp,
14735            n_used,
14736            m.gate_exps.qtype,
14737            m.up_exps.qtype,
14738            m.gate_exps.row_bytes,
14739            m.up_exps.row_bytes,
14740        )?;
14741        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14742        e.moe_down8_fma_into(
14743            crate::WPtr8(d),
14744            crate::F32x8(wv),
14745            &act,
14746            &mut dst,
14747            n_ff_exp,
14748            n_embd,
14749            n_used,
14750            m.down_exps.qtype,
14751            m.down_exps.row_bytes,
14752        )?;
14753        Ok(true)
14754    }
14755
14756    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
14757    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
14758    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
14759    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
14760    #[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
14761    fn moe_cached_gemm_q8(
14762        e: &Engine,
14763        il: u16,
14764        proj: u8,
14765        ex: usize,
14766        m: &MoeWeights,
14767        max_block: usize,
14768        aq: &CudaSlice<i8>,
14769        ad: &CudaSlice<f32>,
14770    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14771        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14772        let exps = match proj {
14773            PROJ_GATE => &m.gate_exps,
14774            PROJ_UP => &m.up_exps,
14775            _ => &m.down_exps,
14776        };
14777        let layout = exps.expert_layout(ex);
14778        let id = BlockId::new(il, proj, ex as u16);
14779        let source = exps.expert_source(ex);
14780        e.with_moe_cache(max_block, |c, eng| {
14781            let slot = c.dispatch_source(id, source, eng)?;
14782            let DispatchSlot::Resident(sl) = slot;
14783            let buf = c.slot(sl);
14784            eng.qmatvec_expert_q8(
14785                buf,
14786                0..layout.len,
14787                aq,
14788                ad,
14789                1,
14790                exps.in_f,
14791                exps.out_f,
14792                layout.qtype,
14793                layout.row_bytes,
14794            )
14795        })
14796    }
14797
14798    fn moe_cached_gemm(
14799        e: &Engine,
14800        il: u16,
14801        proj: u8,
14802        ex: usize,
14803        m: &MoeWeights,
14804        max_block: usize,
14805        x: &cudarc::driver::CudaView<f32>,
14806    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14807        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14808        let exps = match proj {
14809            PROJ_GATE => &m.gate_exps,
14810            PROJ_UP => &m.up_exps,
14811            _ => &m.down_exps,
14812        };
14813        let layout = exps.expert_layout(ex);
14814        let id = BlockId::new(il, proj, ex as u16);
14815        let source = exps.expert_source(ex);
14816        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
14817        e.with_moe_cache(max_block, |c, eng| {
14818            let slot = c.dispatch_source(id, source, eng)?;
14819            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
14820            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
14821            let DispatchSlot::Resident(sl) = slot;
14822            let buf = c.slot(sl);
14823            m.qmatvec_view(
14824                eng,
14825                buf,
14826                0..layout.len,
14827                x,
14828                1,
14829                exps.in_f,
14830                exps.out_f,
14831                layout.qtype,
14832                layout.row_bytes,
14833            )
14834        })
14835    }
14836
14837    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
14838    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
14839    /// so the current forward's backend assignment and output remain unchanged.
14840    fn moe_profile_admit_expert(
14841        e: &Engine,
14842        il: u16,
14843        ex: usize,
14844        m: &MoeWeights,
14845        max_block: usize,
14846    ) -> Result<(), Box<dyn std::error::Error>> {
14847        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14848        e.with_moe_cache(max_block, |cache, eng| {
14849            for (proj, exps) in [
14850                (PROJ_GATE, &m.gate_exps),
14851                (PROJ_UP, &m.up_exps),
14852                (PROJ_DOWN, &m.down_exps),
14853            ] {
14854                let id = BlockId::new(il, proj, ex as u16);
14855                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
14856            }
14857            Ok(())
14858        })
14859    }
14860
14861    /// Read a projection from the immutable residency set when present; otherwise use one
14862    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
14863    #[allow(clippy::too_many_arguments)]
14864    fn moe_frozen_gemm(
14865        e: &Engine,
14866        il: u16,
14867        proj: u8,
14868        ex: usize,
14869        m: &MoeWeights,
14870        max_block: usize,
14871        x: &cudarc::driver::CudaView<f32>,
14872        scratch: &mut Option<CudaSlice<u8>>,
14873        scratch_len: usize,
14874    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14875        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
14876        let exps = match proj {
14877            PROJ_GATE => &m.gate_exps,
14878            PROJ_UP => &m.up_exps,
14879            _ => &m.down_exps,
14880        };
14881        let layout = exps.expert_layout(ex);
14882        let id = BlockId::new(il, proj, ex as u16);
14883        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
14884            let Some(slot) = cache.resident(id) else {
14885                return Ok(None);
14886            };
14887            let buf = cache.slot(slot);
14888            Ok(Some(m.qmatvec_view(
14889                eng,
14890                buf,
14891                0..layout.len,
14892                x,
14893                1,
14894                exps.in_f,
14895                exps.out_f,
14896                layout.qtype,
14897                layout.row_bytes,
14898            )?))
14899        })? {
14900            return Ok(output);
14901        }
14902        if scratch.is_none() {
14903            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
14904        }
14905        let scratch = scratch.as_mut().unwrap();
14906        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
14907        m.qmatvec_view(
14908            e,
14909            scratch,
14910            0..layout.len,
14911            x,
14912            1,
14913            exps.in_f,
14914            exps.out_f,
14915            layout.qtype,
14916            layout.row_bytes,
14917        )
14918    }
14919
14920    fn moe_prefetch_expert(
14921        e: &Engine,
14922        il: u16,
14923        ex: usize,
14924        m: &MoeWeights,
14925        max_block: usize,
14926        keep: &[crate::moe_cache::BlockId],
14927    ) -> Result<(), Box<dyn std::error::Error>> {
14928        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14929        e.with_moe_cache(max_block, |c, eng| {
14930            for (proj, exps) in [
14931                (PROJ_GATE, &m.gate_exps),
14932                (PROJ_UP, &m.up_exps),
14933                (PROJ_DOWN, &m.down_exps),
14934            ] {
14935                let id = BlockId::new(il, proj, ex as u16);
14936                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
14937            }
14938            Ok(())
14939        })
14940    }
14941
14942    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
14943    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
14944    fn moe_prefetch_disk_expert(
14945        e: &Engine,
14946        il: u16,
14947        ex: usize,
14948        m: &MoeWeights,
14949        max_block: usize,
14950        keep: &[crate::moe_cache::BlockId],
14951    ) -> Result<(), Box<dyn std::error::Error>> {
14952        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14953        e.with_moe_cache(max_block, |c, eng| {
14954            for (proj, exps) in [
14955                (PROJ_GATE, &m.gate_exps),
14956                (PROJ_UP, &m.up_exps),
14957                (PROJ_DOWN, &m.down_exps),
14958            ] {
14959                let source = exps.expert_source(ex);
14960                if let crate::model::ExpertSource::Disk { .. } = &source {
14961                    let id = BlockId::new(il, proj, ex as u16);
14962                    let _ = c.prefetch_source(id, source, keep, eng)?;
14963                }
14964            }
14965            Ok(())
14966        })
14967    }
14968
14969    #[inline]
14970    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
14971        let _ = m.gate_exps.prefetch_expert_pages(ex);
14972        let _ = m.up_exps.prefetch_expert_pages(ex);
14973        let _ = m.down_exps.prefetch_expert_pages(ex);
14974    }
14975}
14976
14977// ================================================================================================
14978// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
14979//
14980// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
14981// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
14982// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
14983//
14984// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
14985// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
14986// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
14987// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
14988// identical to the per-token loop regardless of expert processing order.
14989//
14990// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
14991// ================================================================================================
14992
14993impl HybridModel {
14994    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
14995    /// sequential fused q8 program over the token axis; clamped layers use the separate
14996    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
14997    #[allow(clippy::too_many_arguments)]
14998    fn moe_ffn_grouped_resident_q8(
14999        e: &Engine,
15000        m: &MoeWeights,
15001        z: &CudaSlice<f32>,
15002        t: usize,
15003        cfg: &ModelConfig,
15004        il: u16,
15005        sel_all: &[u32],
15006        w_all: &[f32],
15007        table: &CudaSlice<u64>,
15008        gu_il: bool,
15009    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15010        let moe = cfg.moe.as_ref().unwrap();
15011        let n_embd = cfg.n_embd as usize;
15012        let n_expert = moe.expert_count as usize;
15013        let n_used = moe.expert_used_count as usize;
15014        let n_ff_exp = moe.expert_ff_length as usize;
15015        let n_pairs = t * n_used;
15016        debug_assert_eq!(sel_all.len(), n_pairs);
15017        debug_assert_eq!(w_all.len(), n_pairs);
15018        debug_assert!(
15019            m.gate_exps.macros.is_none()
15020                && m.up_exps.macros.is_none()
15021                && m.down_exps.macros.is_none(),
15022            "resident grouped q8 does not fold per-expert macro scales",
15023        );
15024
15025        // The rows twins run the resident sequential program verbatim on grid.z = token:
15026        // fused gate/up/SiLU per slot, batched activation quantization, then the original
15027        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
15028        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
15029        // never enter the softmax router.
15030        if !cfg.swiglu_clamped_at(il as u32) {
15031            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
15032            let sel_d = e.htod_i32(&sel)?;
15033            let w_d = e.htod(w_all)?;
15034            let (gate_row_bytes, up_row_bytes) = if gu_il {
15035                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
15036                (combined, combined)
15037            } else {
15038                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
15039            };
15040            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
15041            let act = e.moe_gate_up_silu8_dev_q8_rows(
15042                table,
15043                &sel_d,
15044                &zq,
15045                &zd,
15046                t,
15047                n_embd,
15048                n_ff_exp,
15049                n_used,
15050                n_expert,
15051                m.gate_exps.qtype,
15052                m.up_exps.qtype,
15053                gate_row_bytes,
15054                up_row_bytes,
15055                &m.dev_macros,
15056            )?;
15057            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
15058            let mut moe_out = e.uninit(t * n_embd)?;
15059            e.moe_down8_fma_dev_q8_rows_g(
15060                table,
15061                &sel_d,
15062                &w_d,
15063                &aq2,
15064                &ad2,
15065                &mut moe_out,
15066                t,
15067                n_ff_exp,
15068                n_embd,
15069                n_used,
15070                n_expert,
15071                m.down_exps.qtype,
15072                m.down_exps.row_bytes,
15073            )?;
15074
15075            if std::env::var("MEMRA_MOE_STATS").is_ok() {
15076                let mut counts = vec![0usize; n_expert];
15077                for &expert in sel_all {
15078                    counts[expert as usize] += 1;
15079                }
15080                let mut sizes: Vec<usize> =
15081                    counts.into_iter().filter(|&count| count != 0).collect();
15082                sizes.sort_unstable();
15083                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
15084                println!(
15085                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
15086                     m_e: min={} median={} mean={mean:.1} max={}",
15087                    sizes.len(),
15088                    n_expert,
15089                    sizes.first().copied().unwrap_or(0),
15090                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
15091                    sizes.last().copied().unwrap_or(0),
15092                );
15093            }
15094            return Ok(moe_out);
15095        }
15096
15097        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
15098        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
15099        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
15100        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
15101        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
15102        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
15103        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
15104
15105        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
15106        for (pair, &expert) in pair_ex.iter().enumerate() {
15107            by_expert[expert as usize].push(pair as i32);
15108        }
15109
15110        let pair_tok_d = e.htod_i32(&pair_tok)?;
15111        let pair_ex_d = e.htod_i32(&pair_ex)?;
15112        let pair_w_d = e.htod(w_all)?;
15113        let tok_off_d = e.htod_i32(&tok_off)?;
15114        let tok_ids_d = e.htod_i32(&tok_ids)?;
15115
15116        let matvec = |proj: i32,
15117                      pair_rows: &CudaSlice<i32>,
15118                      aq: &CudaSlice<i8>,
15119                      ad: &CudaSlice<f32>,
15120                      in_f: usize,
15121                      out_f: usize,
15122                      qtype: i32,
15123                      row_bytes: usize|
15124         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15125            e.moe_pairs_matvec_q8(
15126                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
15127                row_bytes,
15128            )
15129        };
15130
15131        let (gate_row_bytes, up_row_bytes) = if gu_il {
15132            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
15133            (combined, combined)
15134        } else {
15135            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
15136        };
15137        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
15138        let gate = matvec(
15139            0,
15140            &pair_tok_d,
15141            &zq,
15142            &zd,
15143            n_embd,
15144            n_ff_exp,
15145            m.gate_exps.qtype,
15146            gate_row_bytes,
15147        )?;
15148        let up = matvec(
15149            1,
15150            &pair_tok_d,
15151            &zq,
15152            &zd,
15153            n_embd,
15154            n_ff_exp,
15155            m.up_exps.qtype,
15156            up_row_bytes,
15157        )?;
15158        let mut act = e.uninit(n_pairs * n_ff_exp)?;
15159        Self::ffn_act_lim(
15160            e,
15161            cfg,
15162            &gate,
15163            &up,
15164            1.0,
15165            1.0,
15166            cfg.clamp_exp_at(il as u32),
15167            &mut act,
15168            n_pairs * n_ff_exp,
15169        )?;
15170        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
15171        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
15172        let pair_self_d = e.htod_i32(&pair_self)?;
15173        let down = matvec(
15174            2,
15175            &pair_self_d,
15176            &aq2,
15177            &ad2,
15178            n_ff_exp,
15179            n_embd,
15180            m.down_exps.qtype,
15181            m.down_exps.row_bytes,
15182        )?;
15183        let mut moe_out = e.uninit(t * n_embd)?;
15184        e.moe_pairs_scatter(
15185            &down,
15186            &pair_w_d,
15187            &tok_off_d,
15188            &tok_ids_d,
15189            &mut moe_out,
15190            t,
15191            n_embd,
15192        )?;
15193
15194        if std::env::var("MEMRA_MOE_STATS").is_ok() {
15195            let mut sizes: Vec<usize> = by_expert
15196                .iter()
15197                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
15198                .collect();
15199            sizes.sort_unstable();
15200            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
15201            println!(
15202                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
15203                 m_e: min={} median={} mean={mean:.1} max={}",
15204                sizes.len(),
15205                n_expert,
15206                sizes.first().copied().unwrap_or(0),
15207                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
15208                sizes.last().copied().unwrap_or(0),
15209            );
15210        }
15211        Ok(moe_out)
15212    }
15213
15214    /// MEMRA_SHEXP_SPLIT worker: the shared expert's gate/up/down rows split across both
15215    /// devices (dev1 idles during E3), act halves exchanged both ways, downs row-split —
15216    /// per-element/per-row programs identical, so `sh` is BIT-IDENTICAL to the single-device
15217    /// arm. Process-static workspace pinned by the gate tensor pointer; rank1 holds one-time
15218    /// row-half replicas (~10MB/layer x 42). Returns None when ineligible.
15219    #[allow(clippy::too_many_arguments)]
15220    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
15221    fn shexp_split_matvec(
15222        e: &Engine,
15223        rank1: &Engine,
15224        wg: &CudaSlice<u8>,
15225        wu: &CudaSlice<u8>,
15226        wd: &CudaSlice<u8>,
15227        z: &CudaSlice<f32>,
15228        lim: Option<SwigluClamp>,
15229        cfg: &ModelConfig,
15230        il: u16,
15231        n_embd: usize,
15232        n_ff_sh: usize,
15233    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15234        use cudarc::driver::DevicePtr;
15235        if !n_ff_sh.is_multiple_of(2) || !n_embd.is_multiple_of(2) {
15236            return Ok(None);
15237        }
15238        let hf = n_ff_sh / 2;
15239        let nd = n_embd / 2;
15240        struct Rep {
15241            wg1: CudaSlice<u8>,
15242            wu1: CudaSlice<u8>,
15243            wd1: CudaSlice<u8>,
15244        }
15245        struct SplitWs {
15246            pin_dev: usize,
15247            // e side
15248            gate0: CudaSlice<f32>,
15249            up0: CudaSlice<f32>,
15250            act: CudaSlice<f32>,
15251            sh_buf: CudaSlice<f32>,
15252            ev_z: cudarc::driver::CudaEvent,
15253            ev_act0: cudarc::driver::CudaEvent,
15254            // rank1 side
15255            z1: CudaSlice<f32>,
15256            g1: CudaSlice<f32>,
15257            u1: CudaSlice<f32>,
15258            a1h: CudaSlice<f32>,
15259            act1: CudaSlice<f32>,
15260            y1: CudaSlice<f32>,
15261            ev_act1: cudarc::driver::CudaEvent,
15262            ev_y1: cudarc::driver::CudaEvent,
15263            raw_act_e: u64,
15264            raw_sh_e: u64,
15265            raw_z1: u64,
15266            raw_a1h: u64,
15267            raw_act1: u64,
15268            raw_y1: u64,
15269        }
15270        static WS: std::sync::Mutex<Option<SplitWs>> = std::sync::Mutex::new(None);
15271        static REPS: std::sync::Mutex<Option<std::collections::HashMap<u64, Rep>>> =
15272            std::sync::Mutex::new(None);
15273        let mut guard = WS.lock().map_err(|_| "shexp split lock is poisoned")?;
15274        let mut reps_guard = REPS.lock().map_err(|_| "shexp reps lock is poisoned")?;
15275        let reps = reps_guard.get_or_insert_with(std::collections::HashMap::new);
15276        let pins = e.ctx().ordinal();
15277        if guard.as_ref().is_none_or(|w| w.pin_dev != pins) {
15278            let (gate0, up0, act, sh_buf, ev_z, ev_act0) = {
15279                let _m = e.gpu.enter_main()?;
15280                (
15281                    e.htod(&vec![0.0f32; hf])?,
15282                    e.htod(&vec![0.0f32; hf])?,
15283                    e.htod(&vec![0.0f32; n_ff_sh])?,
15284                    e.htod(&vec![0.0f32; n_embd])?,
15285                    e.ctx().new_event(None)?,
15286                    e.ctx().new_event(None)?,
15287                )
15288            };
15289            let (z1, g1, u1, a1h, act1, y1, ev_act1, ev_y1) = {
15290                let _r = rank1.gpu.enter_main()?;
15291                (
15292                    rank1.htod(&vec![0.0f32; n_embd])?,
15293                    rank1.htod(&vec![0.0f32; hf])?,
15294                    rank1.htod(&vec![0.0f32; hf])?,
15295                    rank1.htod(&vec![0.0f32; hf])?,
15296                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
15297                    rank1.htod(&vec![0.0f32; nd])?,
15298                    rank1.ctx().new_event(None)?,
15299                    rank1.ctx().new_event(None)?,
15300                )
15301            };
15302            let (raw_act_e, raw_sh_e) = {
15303                let _m = e.gpu.enter_main()?;
15304                let stream = e.stream();
15305                let (a, _g0) = act.device_ptr(&stream);
15306                let (b, _g1) = sh_buf.device_ptr(&stream);
15307                (a, b)
15308            };
15309            let (raw_z1, raw_a1h, raw_act1, raw_y1) = {
15310                let _r = rank1.gpu.enter_main()?;
15311                let rs = rank1.stream();
15312                let (a, _g0) = z1.device_ptr(&rs);
15313                let (b, _g1) = a1h.device_ptr(&rs);
15314                let (c, _g2) = act1.device_ptr(&rs);
15315                let (d, _g3) = y1.device_ptr(&rs);
15316                (a, b, c, d)
15317            };
15318            *guard = Some(SplitWs {
15319                pin_dev: pins,
15320                gate0,
15321                up0,
15322                act,
15323                sh_buf,
15324                ev_z,
15325                ev_act0,
15326                z1,
15327                g1,
15328                u1,
15329                a1h,
15330                act1,
15331                y1,
15332                ev_act1,
15333                ev_y1,
15334                raw_act_e,
15335                raw_sh_e,
15336                raw_z1,
15337                raw_a1h,
15338                raw_act1,
15339                raw_y1,
15340            });
15341        }
15342        let ws = guard.as_mut().expect("armed above");
15343        let wg_pin = {
15344            let _m = e.gpu.enter_main()?;
15345            let stream = e.stream();
15346            let (p, _g) = wg.device_ptr(&stream);
15347            p
15348        };
15349        if !reps.contains_key(&wg_pin) {
15350            // One-time per layer: upload rank1's row halves (gate/up rows [hf..], down rows [nd..]).
15351            let up = |src: &CudaSlice<u8>,
15352                      off_bytes: usize,
15353                      len: usize|
15354             -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
15355                use cudarc::driver::sys;
15356                let sptr = {
15357                    let _m = e.gpu.enter_main()?;
15358                    let stream = e.stream();
15359                    let (p, _g) = src.device_ptr(&stream);
15360                    p + off_bytes as u64
15361                };
15362                let dst = {
15363                    let _r = rank1.gpu.enter_main()?;
15364                    rank1.alloc_u8_uninit(len)?
15365                };
15366                let dptr = {
15367                    let _r = rank1.gpu.enter_main()?;
15368                    let rs = rank1.stream();
15369                    let (p, _g) = dst.device_ptr(&rs);
15370                    p
15371                };
15372                let _r = rank1.gpu.enter_main()?;
15373                let r = unsafe {
15374                    sys::cuMemcpyAsync(
15375                        dptr as sys::CUdeviceptr,
15376                        sptr as sys::CUdeviceptr,
15377                        len,
15378                        rank1.stream().cu_stream() as sys::CUstream,
15379                    )
15380                };
15381                if r != sys::CUresult::CUDA_SUCCESS {
15382                    return Err(format!("shexp split replica upload: {r:?}").into());
15383                }
15384                rank1.stream().synchronize()?;
15385                Ok(dst)
15386            };
15387            let wg1 = up(wg, hf * n_embd * 2, hf * n_embd * 2)?;
15388            let wu1 = up(wu, hf * n_embd * 2, hf * n_embd * 2)?;
15389            let wd1 = up(wd, nd * n_ff_sh * 2, nd * n_ff_sh * 2)?;
15390            reps.insert(wg_pin, Rep { wg1, wu1, wd1 });
15391        }
15392        let _ = il;
15393        // Per token, evented split flow.
15394        let raw_z = {
15395            let _m = e.gpu.enter_main()?;
15396            let stream = e.stream();
15397            let (p, _g) = z.device_ptr(&stream);
15398            ws.ev_z.record(&stream)?;
15399            p
15400        };
15401        // rank1: pull z, its dual half, its act half; push act half to e; pull e's act half.
15402        {
15403            let rep = reps.get(&wg_pin).expect("uploaded above");
15404            let _r = rank1.gpu.enter_main()?;
15405            rank1.stream().wait(&ws.ev_z)?;
15406            crate::tp::raw_copy_bytes(ws.raw_z1, raw_z, n_embd * 4, rank1)?;
15407            let SplitWs {
15408                z1, g1, u1, a1h, ..
15409            } = &mut *ws;
15410            rank1.matvec_bf16_dual_into(&rep.wg1, &rep.wu1, z1, g1, u1, n_embd, hf)?;
15411            Self::ffn_act_lim(rank1, cfg, g1, u1, 1.0, 1.0, lim, a1h, hf)?;
15412            // local place into act1[hf..] + P2P push into e's act[hf..]
15413            crate::tp::raw_copy_bytes(ws.raw_act1 + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15414            crate::tp::raw_copy_bytes(ws.raw_act_e + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15415            ws.ev_act1.record(&rank1.stream())?;
15416        }
15417        // e: its dual half + act half; publish it; then wait rank1's half and run down lo.
15418        {
15419            let _m = e.gpu.enter_main()?;
15420            let SplitWs {
15421                gate0, up0, act, ..
15422            } = &mut *ws;
15423            let wg_lo = wg.slice(0..hf * n_embd * 2);
15424            let wu_lo = wu.slice(0..hf * n_embd * 2);
15425            e.matvec_bf16_dual_view_into(&wg_lo, &wu_lo, z, gate0, up0, n_embd, hf)?;
15426            Self::ffn_act_lim(e, cfg, gate0, up0, 1.0, 1.0, lim, act, hf)?;
15427            ws.ev_act0.record(&e.stream())?;
15428        }
15429        // rank1: pull e's act half into act1[0..hf], run down hi, push y half.
15430        {
15431            let rep = reps.get(&wg_pin).expect("uploaded above");
15432            let _r = rank1.gpu.enter_main()?;
15433            rank1.stream().wait(&ws.ev_act0)?;
15434            crate::tp::raw_copy_bytes(ws.raw_act1, ws.raw_act_e, hf * 4, rank1)?;
15435            let SplitWs { act1, y1, .. } = &mut *ws;
15436            rank1.matvec_bf16_into(&rep.wd1, act1, y1, n_ff_sh, nd)?;
15437            crate::tp::raw_copy_bytes(ws.raw_sh_e + (nd * 4) as u64, ws.raw_y1, nd * 4, rank1)?;
15438            ws.ev_y1.record(&rank1.stream())?;
15439        }
15440        // e: down lo into sh_buf[0..nd]; join rank1's half; hand back an owned sh.
15441        {
15442            let _m = e.gpu.enter_main()?;
15443            e.stream().wait(&ws.ev_act1)?;
15444            let SplitWs { act, sh_buf, .. } = &mut *ws;
15445            let wd_lo = wd.slice(0..nd * n_ff_sh * 2);
15446            e.matvec_bf16_view_into(&wd_lo, act, sh_buf, n_ff_sh, nd)?;
15447            e.stream().wait(&ws.ev_y1)?;
15448            let mut sh = e.uninit(n_embd)?;
15449            {
15450                let mut dst = sh.slice_mut(0..n_embd);
15451                e.stream()
15452                    .memcpy_dtod(&ws.sh_buf.slice(0..n_embd), &mut dst)?;
15453            }
15454            Ok(Some(sh))
15455        }
15456    }
15457
15458    /// SHEXP OVERLAP issue (MEMRA_SHEXP_OVERLAP=1): the shared expert reads only `z`, so
15459    /// its kernels (dual matvec+SwiGLU, down) are issued on e's stream from the routes
15460    /// PREJOIN hook — they execute while the peer rank drains its sweep, filling dev0's
15461    /// join wait. The down lands in ITS OWN row (the pre-#2e split program — receipted
15462    /// bit-identical to the fused down+addscale) and `shexp_overlap_apply` adds it after
15463    /// the join with the exact add_scaled_rows expression: values unchanged.
15464    fn shexp_overlap_issue(
15465        e: &Engine,
15466        m: &MoeWeights,
15467        z: &CudaSlice<f32>,
15468        cfg: &ModelConfig,
15469        il: u16,
15470        n_embd: usize,
15471    ) -> Result<bool, Box<dyn std::error::Error>> {
15472        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15473            return Ok(false);
15474        }
15475        let (
15476            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15477            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15478            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15479        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15480        else {
15481            return Ok(false);
15482        };
15483        let n_ff_sh = m
15484            .gate_shexp
15485            .as_ref()
15486            .expect("matched Some above")
15487            .out_features();
15488        // The dual-silu epilogue is step35's POST form only; a PRE-clamped layer declines here
15489        // and takes the unfused seam.
15490        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15491            return Ok(false);
15492        };
15493        let mut guard = SHEXP_OV_WS
15494            .lock()
15495            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15496        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15497        if guard
15498            .as_ref()
15499            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15500        {
15501            *guard = Some((
15502                pins.0,
15503                pins.1,
15504                pins.2,
15505                e.uninit(n_ff_sh)?,
15506                e.uninit(n_embd)?,
15507            ));
15508        }
15509        let (_, _, _, act, sh) = guard.as_mut().expect("armed above");
15510        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim)?;
15511        e.matvec_bf16_into(wd, act, sh, n_ff_sh, n_embd)?;
15512        drop(guard);
15513        Ok(true)
15514    }
15515
15516    /// SHEXP ON DEV1 issue (MEMRA_SHEXP_DEV1=1): the shared expert runs on rank1 — the
15517    /// idle device — with replica weights (one-time P2P upload), the SAME kernels and the
15518    /// SAME split program as the dev0 overlap (dual matvec+SwiGLU, f32acc down), so the
15519    /// values are bit-identical. z rides one 16KB P2P pull behind an e-stream event; the
15520    /// down row lands root-resident (single P2P store pass); apply waits ev_done on e.
15521    #[allow(clippy::too_many_arguments)]
15522    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
15523    fn shexp_dev1_issue(
15524        e: &Engine,
15525        rank1: &Engine,
15526        m: &MoeWeights,
15527        z: &CudaSlice<f32>,
15528        cfg: &ModelConfig,
15529        il: u16,
15530        n_embd: usize,
15531    ) -> Result<bool, Box<dyn std::error::Error>> {
15532        use cudarc::driver::DevicePtr;
15533        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15534            return Ok(false);
15535        }
15536        let (
15537            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15538            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15539            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15540        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15541        else {
15542            return Ok(false);
15543        };
15544        let n_ff_sh = m
15545            .gate_shexp
15546            .as_ref()
15547            .expect("matched Some above")
15548            .out_features();
15549        // POST-form epilogue only (see `fused_post_limit`): a PRE-clamped layer declines
15550        // (Ok(false) = nothing issued, caller falls back per column).
15551        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15552            return Ok(false);
15553        };
15554        // Shared scratch, geometry-keyed.
15555        let mut ws_guard = SHEXP_D1_WS
15556            .lock()
15557            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15558        if ws_guard
15559            .as_ref()
15560            .is_none_or(|(k, ..)| *k != (n_embd, n_ff_sh))
15561        {
15562            let (act1, z1, ev_done) = {
15563                let _r1 = rank1.gpu.enter_main()?;
15564                (
15565                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
15566                    rank1.htod(&vec![0.0f32; n_embd])?,
15567                    rank1.ctx().new_event(None)?,
15568                )
15569            };
15570            let (sh_root, ev_z) = {
15571                let _main = e.gpu.enter_main()?;
15572                (e.htod(&vec![0.0f32; n_embd])?, e.ctx().new_event(None)?)
15573            };
15574            *ws_guard = Some(((n_embd, n_ff_sh), act1, z1, sh_root, ev_z, ev_done));
15575        }
15576        // Per-LAYER weight replicas (gate/up/down differ per layer): one-time P2P upload.
15577        let mut reps_guard = SHEXP_D1_REPS
15578            .lock()
15579            .map_err(|_| "shexp dev1 replica lock is poisoned")?;
15580        let reps = reps_guard.get_or_insert_with(Default::default);
15581        if !reps.contains_key(&il) {
15582            let (wg1, wu1, wd1) = {
15583                let _r1 = rank1.gpu.enter_main()?;
15584                (
15585                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15586                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15587                    rank1.alloc_u8_uninit(n_embd * n_ff_sh * 2)?,
15588                )
15589            };
15590            for (src, dst) in [(wg, &wg1), (wu, &wu1), (wd, &wd1)] {
15591                let s_ptr = {
15592                    let _main = e.gpu.enter_main()?;
15593                    let stream = e.stream();
15594                    let (p, _g) = src.device_ptr(&stream);
15595                    p
15596                };
15597                let d_ptr = {
15598                    let _r1 = rank1.gpu.enter_main()?;
15599                    let stream = rank1.stream();
15600                    let (p, _g) = dst.device_ptr(&stream);
15601                    p
15602                };
15603                let _r1 = rank1.gpu.enter_main()?;
15604                crate::tp::raw_copy_bytes(d_ptr, s_ptr, src.len(), rank1)?;
15605            }
15606            {
15607                let _r1 = rank1.gpu.enter_main()?;
15608                rank1.stream().synchronize()?;
15609            }
15610            reps.insert(il, (wg1, wu1, wd1));
15611        }
15612        let (wg1, wu1, wd1) = reps.get(&il).expect("armed above");
15613        let (_, act1, z1, sh_root, ev_z, ev_done) = ws_guard.as_mut().expect("armed above");
15614        // z ready on e's stream -> rank1 pulls it, runs the split shexp, pushes the down
15615        // row root-side (single store pass), rings ev_done.
15616        let (raw_z, raw_sh) = {
15617            let _main = e.gpu.enter_main()?;
15618            let stream = e.stream();
15619            let (a, _g0) = z.device_ptr(&stream);
15620            let (b, _g1) = sh_root.device_ptr(&stream);
15621            ev_z.record(&stream)?;
15622            (a, b)
15623        };
15624        {
15625            let _r1 = rank1.gpu.enter_main()?;
15626            rank1.stream().wait(ev_z)?;
15627            let raw_z1 = {
15628                let stream = rank1.stream();
15629                let (p, _g) = z1.device_ptr(&stream);
15630                p
15631            };
15632            crate::tp::raw_copy_bytes(raw_z1, raw_z, n_embd * 4, rank1)?;
15633            rank1.matvec_bf16_dual_silu_into(wg1, wu1, z1, act1, n_embd, n_ff_sh, lim)?;
15634            // down writes the ROOT-resident row over P2P via the raw-output twin of
15635            // matvec_bf16_into: reuse the view launcher with a slice view is not possible
15636            // cross-device, so launch on the raw pointer.
15637            rank1.matvec_bf16_raw_out(wd1, act1, raw_sh, n_ff_sh, n_embd)?;
15638            ev_done.record(&rank1.stream())?;
15639        }
15640        Ok(true)
15641    }
15642
15643    /// Apply the dev1 shared expert: wait ev_done on e, then the exact add_scaled_rows.
15644    fn shexp_dev1_apply(
15645        e: &Engine,
15646        output: &mut CudaSlice<f32>,
15647        n_embd: usize,
15648    ) -> Result<(), Box<dyn std::error::Error>> {
15649        let guard = SHEXP_D1_WS
15650            .lock()
15651            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15652        let (pin, _, _, sh_root, _, ev_done) =
15653            guard.as_ref().ok_or("shexp dev1 apply without issue")?;
15654        if pin.0 != n_embd {
15655            return Err("shexp dev1 width drifted".into());
15656        }
15657        let _main = e.gpu.enter_main()?;
15658        e.stream().wait(ev_done)?;
15659        static ONES_D1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15660            std::sync::Mutex::new(None);
15661        let mut og = ONES_D1.lock().map_err(|_| "ones lock is poisoned")?;
15662        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15663            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15664        }
15665        let ones = &og.as_ref().expect("armed above").1;
15666        e.add_scaled_rows(sh_root, ones, output, n_embd, 1)?;
15667        Ok(())
15668    }
15669
15670    /// MOE TAIL FUSION M1 helper: pre-arm the overlap ws + persistent ones row and
15671    /// return their RAW pointers (None when the overlap is ineligible — the caller then
15672    /// takes the split path). Mirrors shexp_overlap_issue's eligibility exactly.
15673    fn shexp_overlap_tail_ptrs(
15674        e: &Engine,
15675        m: &MoeWeights,
15676        cfg: &ModelConfig,
15677        n_embd: usize,
15678    ) -> Result<Option<(u64, u64)>, Box<dyn std::error::Error>> {
15679        use cudarc::driver::DevicePtr;
15680        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15681            return Ok(None);
15682        }
15683        let (
15684            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15685            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15686            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15687        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15688        else {
15689            return Ok(None);
15690        };
15691        let n_ff_sh = m
15692            .gate_shexp
15693            .as_ref()
15694            .expect("matched Some above")
15695            .out_features();
15696        let mut guard = SHEXP_OV_WS
15697            .lock()
15698            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15699        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15700        if guard
15701            .as_ref()
15702            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15703        {
15704            *guard = Some((
15705                pins.0,
15706                pins.1,
15707                pins.2,
15708                e.uninit(n_ff_sh)?,
15709                e.uninit(n_embd)?,
15710            ));
15711        }
15712        let sh_raw = {
15713            let (_, _, _, _, sh) = guard.as_ref().expect("armed above");
15714            let stream = e.stream();
15715            let (p, _g) = sh.device_ptr(&stream);
15716            p
15717        };
15718        static ONES_T3: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15719            std::sync::Mutex::new(None);
15720        let mut og = ONES_T3.lock().map_err(|_| "ones lock is poisoned")?;
15721        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15722            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15723        }
15724        let ones_raw = {
15725            let stream = e.stream();
15726            let (p, _g) = og.as_ref().expect("armed above").1.device_ptr(&stream);
15727            p
15728        };
15729        Ok(Some((sh_raw, ones_raw)))
15730    }
15731
15732    /// Apply the overlapped shared expert: output[r] += sh[r] * 1.0 — the exact
15733    /// add_scaled_rows program the split path used (persistent ones row, no htod).
15734    fn shexp_overlap_apply(
15735        e: &Engine,
15736        output: &mut CudaSlice<f32>,
15737        n_embd: usize,
15738    ) -> Result<(), Box<dyn std::error::Error>> {
15739        let guard = SHEXP_OV_WS
15740            .lock()
15741            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15742        let (_, ne, _, _, sh) = guard.as_ref().ok_or("shexp overlap apply without issue")?;
15743        if *ne != n_embd {
15744            return Err("shexp overlap width drifted".into());
15745        }
15746        static ONES_OV: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15747            std::sync::Mutex::new(None);
15748        let mut og = ONES_OV.lock().map_err(|_| "ones lock is poisoned")?;
15749        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15750            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15751        }
15752        let ones = &og.as_ref().expect("armed above").1;
15753        e.add_scaled_rows(sh, ones, output, n_embd, 1)?;
15754        Ok(())
15755    }
15756
15757    fn moe_ffn_grouped_add_shared(
15758        e: &Engine,
15759        m: &MoeWeights,
15760        z: &CudaSlice<f32>,
15761        t: usize,
15762        cfg: &ModelConfig,
15763        il: u16,
15764        moe_out: &mut CudaSlice<f32>,
15765    ) -> Result<(), Box<dyn std::error::Error>> {
15766        // MEMRA_STEP_TP_TIMING=1: shared-expert wall (syncs e's stream at exit to bill the
15767        // queued matmuls here rather than at the next host readback).
15768        static SHEXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15769        static SHEXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15770        let shexp_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
15771        let shexp_started = shexp_timing.then(std::time::Instant::now);
15772        let result = Self::moe_ffn_grouped_add_shared_inner(e, m, z, t, cfg, il, moe_out);
15773        if let Some(started) = shexp_started {
15774            use std::sync::atomic::Ordering;
15775            e.stream().synchronize()?;
15776            let ns = SHEXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
15777                + started.elapsed().as_nanos() as u64;
15778            let calls = SHEXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
15779            if calls.is_multiple_of(430) {
15780                eprintln!(
15781                    "[moe-shexp-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
15782                    ns as f64 / 1.0e6,
15783                    ns as f64 / calls as f64 / 1.0e3,
15784                );
15785            }
15786        }
15787        result
15788    }
15789
15790    #[allow(clippy::too_many_arguments)]
15791    fn moe_ffn_grouped_add_shared_inner(
15792        e: &Engine,
15793        m: &MoeWeights,
15794        z: &CudaSlice<f32>,
15795        t: usize,
15796        cfg: &ModelConfig,
15797        il: u16,
15798        moe_out: &mut CudaSlice<f32>,
15799    ) -> Result<(), Box<dyn std::error::Error>> {
15800        let n_embd = cfg.n_embd as usize;
15801        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
15802            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15803        {
15804            let n_ff_sh = gate_shexp.out_features();
15805            let lim = cfg.clamp_shexp_at(il as u32);
15806            // T=1 DECODE FUSION (2026-08-20): the ffn_swiglu_decode fast-path program, ported
15807            // — one shared quantize feeds gate+up (matmul at m=1 quantizes internally per
15808            // call with identical bytes, so sharing it is bit-identical), the dual NVFP4/Q8
15809            // launch covers both when available, and silu_mul_scaled_q8_1 emits down's
15810            // operand pre-quantized (kernel_check-proven identities). This path measured
15811            // 167us/layer as separate matmuls + 5 allocs at decode.
15812            let fused = t == 1
15813                && lim.is_none()
15814                && cfg.m3.is_none()
15815                && e.uses_q8_1_fast(gate_shexp)
15816                && e.uses_q8_1_fast(up_shexp);
15817            let canonical_w4a16_rows =
15818                t <= 32 && m.step_ep.as_ref().is_some_and(|ep| ep.nvfp4_device_routes);
15819            // MEMRA_BF16_MMV class: both projections in ONE launch (bit-identical per row to
15820            // the two matvec_bf16 launches matmul would issue). W4A16 distributed execution
15821            // uses the same row program for decode and verify; a t=1-only fusion accumulated
15822            // sub-ULP residual drift from the first MoE layer onward.
15823            let bf16_dual = if (t == 1 || canonical_w4a16_rows)
15824                && crate::Engine::bf16_mmv_on()
15825                && n_embd.is_multiple_of(8)
15826            {
15827                match (gate_shexp, up_shexp) {
15828                    (
15829                        crate::model::GpuTensor::FloatBf16 { data: wg, .. },
15830                        crate::model::GpuTensor::FloatBf16 { data: wu, .. },
15831                    ) => Some((wg, wu)),
15832                    _ => None,
15833                }
15834            } else {
15835                None
15836            };
15837            let sh = if let Some((wg, wu)) = bf16_dual {
15838                // Persistent shared-expert workspace: sizes are constant across every MoE
15839                // layer, so one process-level set pinned by (device, n_embd, n_ff_sh) removes
15840                // the four per-layer allocations. Buffers are fully overwritten each call.
15841                type SharedExpertWorkspace = (
15842                    usize,
15843                    usize,
15844                    usize,
15845                    CudaSlice<f32>,
15846                    CudaSlice<f32>,
15847                    CudaSlice<f32>,
15848                    CudaSlice<f32>,
15849                );
15850                static SHEXP_WS: std::sync::Mutex<
15851                    Option<std::collections::HashMap<usize, SharedExpertWorkspace>>,
15852                > = std::sync::Mutex::new(None);
15853                let down_bf16 = match down_shexp {
15854                    crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15855                    _ => None,
15856                };
15857                let mut guard = SHEXP_WS
15858                    .lock()
15859                    .map_err(|_| "shexp workspace lock is poisoned")?;
15860                let capacity = if canonical_w4a16_rows { 32 } else { 1 };
15861                let device = e.ctx().ordinal();
15862                let workspaces = guard.get_or_insert_with(Default::default);
15863                if workspaces
15864                    .get(&device)
15865                    .is_none_or(|(ne, nf, cap, ..)| (*ne, *nf, *cap) != (n_embd, n_ff_sh, capacity))
15866                {
15867                    workspaces.insert(
15868                        device,
15869                        (
15870                            n_embd,
15871                            n_ff_sh,
15872                            capacity,
15873                            e.uninit(capacity * n_ff_sh)?,
15874                            e.uninit(capacity * n_ff_sh)?,
15875                            e.uninit(capacity * n_ff_sh)?,
15876                            e.uninit(capacity * n_embd)?,
15877                        ),
15878                    );
15879                }
15880                // MEMRA_SHEXP_SPLIT=1: bit-identical row split across both devices; falls
15881                // through to the single-device arm when ineligible.
15882                {
15883                    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15884                    let split_on = *ON
15885                        .get_or_init(|| std::env::var("MEMRA_SHEXP_SPLIT").as_deref() == Ok("1"));
15886                    if split_on
15887                        && t == 1
15888                        && let (Some(wd), Some(rank1)) = (
15889                            match down_shexp {
15890                                crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15891                                _ => None,
15892                            },
15893                            m.step_tp.as_ref().and_then(|st| st.runtime.rank_engine(1)),
15894                        )
15895                        && let Some(sh) = Self::shexp_split_matvec(
15896                            e, rank1, wg, wu, wd, z, lim, cfg, il, n_embd, n_ff_sh,
15897                        )?
15898                    {
15899                        drop(guard);
15900                        let gate = match &m.gate_inp_shexp {
15901                            Some(gate_inp_shexp) => {
15902                                e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15903                            }
15904                            None => e.htod(&vec![1.0f32; t])?,
15905                        };
15906                        e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
15907                        return Ok(());
15908                    }
15909                }
15910                let (_, _, _, gate, up, act, sh_buf) = workspaces
15911                    .get_mut(&device)
15912                    .expect("shexp workspace initialized above");
15913                // FUSION #2b needs the POST form its epilogue hardcodes; m3's swigluoai and
15914                // glm5_next's PRE clamp both take the unfused dual-matmul + ffn_act_lim arm.
15915                if let (true, Ok(lim_post)) = (cfg.m3.is_none(), Self::fused_post_limit(lim)) {
15916                    // dual matvec + SwiGLU act in one launch — exact dual per-row program +
15917                    // exact silu/clamped expression, bit-identical.
15918                    if canonical_w4a16_rows {
15919                        e.matvec_bf16_dual_silu_rows_into(
15920                            wg, wu, z, act, n_embd, n_ff_sh, lim_post, t,
15921                        )?;
15922                    } else {
15923                        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim_post)?;
15924                    }
15925                    let _ = (&gate, &up);
15926                } else {
15927                    e.matvec_bf16_dual_into(wg, wu, z, gate, up, n_embd, n_ff_sh)?;
15928                    Self::ffn_act_lim(e, cfg, gate, up, 1.0, 1.0, lim, act, n_ff_sh)?;
15929                }
15930                if let Some(down) = down_bf16 {
15931                    if canonical_w4a16_rows {
15932                        e.matvec_bf16_rows_into(down, act, sh_buf, n_ff_sh, n_embd, t)?;
15933                        let mut sh = e.uninit(t * n_embd)?;
15934                        {
15935                            let mut dst = sh.slice_mut(0..t * n_embd);
15936                            e.stream()
15937                                .memcpy_dtod(&sh_buf.slice(0..t * n_embd), &mut dst)?;
15938                        }
15939                        sh
15940                    } else {
15941                        // FUSION #2e (gate-less shexp only, MEMRA_FUSE_DOWN_ADDSCALE=0 reverts):
15942                        // down matvec + scaled accumulate straight into moe_out in ONE launch —
15943                        // exact f32acc per-row program + the exact add_scaled_rows expression
15944                        // (dst[r] += y_r * 1.0). Replaces down + ownership alloc + 16KB copy +
15945                        // add_scaled (3 launches + alloc -> 1 launch); bit-identical because the
15946                        // accumulate consumes the same f32 the split path stored and reloaded.
15947                        static FUSE_DA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15948                        let fuse_da = *FUSE_DA.get_or_init(|| {
15949                            std::env::var("MEMRA_FUSE_DOWN_ADDSCALE").as_deref() != Ok("0")
15950                        });
15951                        if fuse_da && m.gate_inp_shexp.is_none() {
15952                            static ONES1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15953                                std::sync::Mutex::new(None);
15954                            let mut og = ONES1.lock().map_err(|_| "shexp ones lock is poisoned")?;
15955                            if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15956                                *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15957                            }
15958                            let ones = &og.as_ref().expect("armed above").1;
15959                            e.matvec_bf16_down_addscale_into(
15960                                down, act, ones, moe_out, n_ff_sh, n_embd,
15961                            )?;
15962                            return Ok(());
15963                        }
15964                        e.matvec_bf16_into(down, act, sh_buf, n_ff_sh, n_embd)?;
15965                        let sh = e.uninit(n_embd)?;
15966                        // One alloc keeps the ownership contract; the copy is 16KB on-stream.
15967                        let mut sh = sh;
15968                        {
15969                            let mut dst = sh.slice_mut(0..n_embd);
15970                            e.stream().memcpy_dtod(&sh_buf.slice(0..n_embd), &mut dst)?;
15971                        }
15972                        sh
15973                    }
15974                } else {
15975                    e.matmul(down_shexp, act, t)?
15976                }
15977            } else if fused {
15978                let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
15979                let pair = match e.matmul_pre_dual_noscale(gate_shexp, up_shexp, &zq, &zd, 1)? {
15980                    Some((gate, up)) => Some((gate, up)),
15981                    None => {
15982                        match (
15983                            e.matmul_pre_noscale(gate_shexp, &zq, &zd, 1)?,
15984                            e.matmul_pre_noscale(up_shexp, &zq, &zd, 1)?,
15985                        ) {
15986                            (Some(gate), Some(up)) => Some((gate, up)),
15987                            _ => None,
15988                        }
15989                    }
15990                };
15991                match pair {
15992                    Some(((gate, gs), (up, us))) => {
15993                        if e.uses_q8_1_fast(down_shexp) {
15994                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff_sh)?;
15995                            e.matmul_pre(down_shexp, &aq, &ad, &gate, 1)?
15996                        } else {
15997                            let mut act = e.uninit(n_ff_sh)?;
15998                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff_sh)?;
15999                            e.matmul(down_shexp, &act, 1)?
16000                        }
16001                    }
16002                    None => {
16003                        let gate = e.matmul_pre(gate_shexp, &zq, &zd, z, 1)?;
16004                        let up = e.matmul_pre(up_shexp, &zq, &zd, z, 1)?;
16005                        let mut act = e.uninit(n_ff_sh)?;
16006                        Self::ffn_act(e, cfg, &gate, &up, &mut act, n_ff_sh)?;
16007                        e.matmul(down_shexp, &act, 1)?
16008                    }
16009                }
16010            } else {
16011                let sg_gate = e.matmul(gate_shexp, z, t)?;
16012                let sg_up = e.matmul(up_shexp, z, t)?;
16013                let mut sa = e.uninit(t * n_ff_sh)?;
16014                Self::ffn_act_lim(
16015                    e,
16016                    cfg,
16017                    &sg_gate,
16018                    &sg_up,
16019                    1.0,
16020                    1.0,
16021                    lim,
16022                    &mut sa,
16023                    t * n_ff_sh,
16024                )?;
16025                e.matmul(down_shexp, &sa, t)?
16026            };
16027            let gate = match &m.gate_inp_shexp {
16028                Some(gate_inp_shexp) => {
16029                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
16030                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
16031                    } else {
16032                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
16033                        let mut gate = e.uninit(t)?;
16034                        e.sigmoid(&raw, &mut gate, t)?;
16035                        gate
16036                    }
16037                }
16038                // t=1 hot path: the per-layer htod of a ones row is a PAGEABLE H2D that
16039                // synchronizes the stream — measured as the biggest per-layer host gap
16040                // (44.6us x 42, eager gap table 2026-08-21). One persistent ones row per
16041                // device serves every layer; larger t (prefill) keeps the plain htod.
16042                None if t == 1 => {
16043                    static ONES: std::sync::Mutex<
16044                        Option<std::collections::HashMap<usize, CudaSlice<f32>>>,
16045                    > = std::sync::Mutex::new(None);
16046                    let mut guard = ONES.lock().map_err(|_| "shexp ones lock is poisoned")?;
16047                    let device = e.ctx().ordinal();
16048                    let rows = guard.get_or_insert_with(Default::default);
16049                    // One entry lookup, not three (contains_key + insert + get). The vacant arm
16050                    // stays fallible, which is why this is `match` and not `or_insert_with`.
16051                    use std::collections::hash_map::Entry;
16052                    let ones = match rows.entry(device) {
16053                        Entry::Occupied(occupied) => occupied.into_mut(),
16054                        Entry::Vacant(vacant) => vacant.insert(e.htod(&[1.0f32])?),
16055                    };
16056                    e.add_scaled_rows(&sh, ones, moe_out, n_embd, t)?;
16057                    return Ok(());
16058                }
16059                None => e.htod(&vec![1.0f32; t])?,
16060            };
16061            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
16062        }
16063        Ok(())
16064    }
16065
16066    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
16067    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
16068    pub(crate) fn moe_ffn_grouped(
16069        e: &Engine,
16070        m: &MoeWeights,
16071        z: &CudaSlice<f32>,
16072        t: usize,
16073        cfg: &ModelConfig,
16074        il: u16,
16075        max_block: usize,
16076    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16077        let moe = cfg.moe.as_ref().unwrap();
16078        let n_embd = cfg.n_embd as usize;
16079        let n_expert = moe.expert_count as usize;
16080        let n_used = moe.expert_used_count as usize;
16081        let n_ff_exp = moe.expert_ff_length as usize;
16082        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
16083        let lim_exp = cfg.clamp_exp_at(il as u32);
16084
16085        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
16086        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
16087        // enters the softmax-only pairs/dev router.
16088        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
16089        if let Some(sig) = cfg.sigmoid_router() {
16090            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
16091        }
16092        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
16093            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
16094        } else {
16095            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
16096        };
16097        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
16098        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
16099        Self::trace_moe_input(e, il, t, n_embd, z)?;
16100
16101        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
16102        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
16103        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
16104        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
16105        let no_exp_macros = m.gate_exps.macros.is_none()
16106            && m.up_exps.macros.is_none()
16107            && m.down_exps.macros.is_none();
16108        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
16109            m.has_uniform_expert_layout()
16110                && no_exp_macros
16111                && moe_q8_enabled_for_model(cfg, m)
16112                && moe_slab_enabled()
16113                && dev.dev == e.ctx().ordinal()
16114        });
16115        if let Some(dev) = resident_q8 {
16116            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
16117                e,
16118                m,
16119                z,
16120                t,
16121                cfg,
16122                il,
16123                &sel_all,
16124                &w_all,
16125                &dev.ptr_row,
16126                dev.gu_il,
16127            )?;
16128            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
16129            return Ok(moe_out);
16130        }
16131
16132        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
16133        // For each expert e, we need: which tokens use it, their positions in z, their top-k
16134        // slot index (for bit-identical accumulation), and their weights.
16135        struct ExpertGroup {
16136            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
16137            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
16138            weights: Vec<f32>,      // renormalized weight for that token-expert pair
16139        }
16140        let mut groups: Vec<ExpertGroup> = (0..n_expert)
16141            .map(|_| ExpertGroup {
16142                tok_indices: Vec::new(),
16143                slot_indices: Vec::new(),
16144                weights: Vec::new(),
16145            })
16146            .collect();
16147
16148        for tok in 0..t {
16149            for j in 0..n_used {
16150                let ex = sel_all[tok * n_used + j] as usize;
16151                let w = w_all[tok * n_used + j];
16152                groups[ex].tok_indices.push(tok as i32);
16153                groups[ex].slot_indices.push(j as i32);
16154                groups[ex].weights.push(w);
16155            }
16156        }
16157
16158        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
16159        // Each token's 8 expert contributions land in their respective slots.
16160        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
16161        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
16162
16163        // Expert weight dimensions (used in both cache and staging paths).
16164        let g_len = m.gate_exps.max_expert_bytes();
16165        let u_len = m.up_exps.max_expert_bytes();
16166        let d_len = m.down_exps.max_expert_bytes();
16167        let moe_q8 = moe_q8_enabled_for_model(cfg, m);
16168        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
16169        // Interleaved GU slabs require the pointer-table fast path above.
16170        let slab_local = m
16171            .dev_exps
16172            .as_ref()
16173            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
16174        let use_cache =
16175            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
16176        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
16177        // also does: a local resident slab or a live SLRU dispatch.
16178        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
16179
16180        // GPU scratch for staging (only allocated without a local slab or cache).
16181        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
16182            (
16183                Some(e.alloc_u8(g_len)?),
16184                Some(e.alloc_u8(u_len)?),
16185                Some(e.alloc_u8(d_len)?),
16186            )
16187        } else {
16188            (None, None, None)
16189        };
16190
16191        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
16192        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
16193        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
16194        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
16195        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
16196        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
16197        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
16198        // at long prompts where every expert stages regardless. Order is FREE to change without
16199        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
16200        // regardless of expert processing order (the whole point of the slots).
16201        let mut order: Vec<usize> = (0..n_expert)
16202            .filter(|&ex| !groups[ex].tok_indices.is_empty())
16203            .collect();
16204        order.sort_by(|&a, &b| {
16205            groups[b]
16206                .tok_indices
16207                .len()
16208                .cmp(&groups[a].tok_indices.len())
16209                .then(a.cmp(&b))
16210        });
16211        let mut m_dist: Vec<usize> = Vec::new(); // for stats
16212        let page_window = moe_page_prefetch_window();
16213        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
16214        if worker_disk_prefetch
16215            && let Some(first) = grouped_worker_prefetch_position(order.len(), None)
16216        {
16217            Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
16218        }
16219        for (order_pos, &ex) in order.iter().enumerate() {
16220            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
16221                Self::moe_prefetch_host_expert(order[next], m);
16222            }
16223            if worker_disk_prefetch
16224                && let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos))
16225            {
16226                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16227                let keep = [
16228                    BlockId::new(il, PROJ_GATE, ex as u16),
16229                    BlockId::new(il, PROJ_UP, ex as u16),
16230                    BlockId::new(il, PROJ_DOWN, ex as u16),
16231                ];
16232                Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
16233            }
16234            let grp = &groups[ex];
16235            let m_e = grp.tok_indices.len();
16236            m_dist.push(m_e);
16237            let gl = m.gate_exps.expert_layout(ex);
16238            let ul = m.up_exps.expert_layout(ex);
16239            let dl = m.down_exps.expert_layout(ex);
16240
16241            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
16242            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
16243            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
16244            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
16245            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
16246            let dmac = m.down_exps.macro_scale(ex);
16247            let weight_d = if dmac == 1.0 {
16248                e.htod(&grp.weights)?
16249            } else {
16250                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
16251                e.htod(&scaled)?
16252            };
16253
16254            // GATHER: collect m_e activation rows from z into a contiguous buffer.
16255            let mut gathered = e.zeros(m_e * n_embd)?;
16256            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
16257            let gv = gathered.slice(0..m_e * n_embd);
16258
16259            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
16260            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
16261            let y = if let Some(dev) = slab_local {
16262                let gate_start = ex * m.gate_exps.expert_stride;
16263                let up_start = ex * m.up_exps.expert_stride;
16264                let down_start = ex * m.down_exps.expert_stride;
16265                if grouped_q8 {
16266                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16267                    let gate = e.qmatvec_expert_q8(
16268                        &dev.gate,
16269                        gate_start..gate_start + gl.len,
16270                        &zq,
16271                        &zd,
16272                        m_e,
16273                        m.gate_exps.in_f,
16274                        m.gate_exps.out_f,
16275                        gl.qtype,
16276                        gl.row_bytes,
16277                    )?;
16278                    let up = e.qmatvec_expert_q8(
16279                        &dev.up,
16280                        up_start..up_start + ul.len,
16281                        &zq,
16282                        &zd,
16283                        m_e,
16284                        m.up_exps.in_f,
16285                        m.up_exps.out_f,
16286                        ul.qtype,
16287                        ul.row_bytes,
16288                    )?;
16289                    let mut act = e.uninit(m_e * n_ff_exp)?;
16290                    Self::ffn_act_lim(
16291                        e,
16292                        cfg,
16293                        &gate,
16294                        &up,
16295                        m.gate_exps.macro_scale(ex),
16296                        m.up_exps.macro_scale(ex),
16297                        lim_exp,
16298                        &mut act,
16299                        m_e * n_ff_exp,
16300                    )?;
16301                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16302                    e.qmatvec_expert_q8(
16303                        &dev.down,
16304                        down_start..down_start + dl.len,
16305                        &aq2,
16306                        &ad2,
16307                        m_e,
16308                        m.down_exps.in_f,
16309                        m.down_exps.out_f,
16310                        dl.qtype,
16311                        dl.row_bytes,
16312                    )?
16313                } else {
16314                    let gate = m.qmatvec_view(
16315                        e,
16316                        &dev.gate,
16317                        gate_start..gate_start + gl.len,
16318                        &gv,
16319                        m_e,
16320                        m.gate_exps.in_f,
16321                        m.gate_exps.out_f,
16322                        gl.qtype,
16323                        gl.row_bytes,
16324                    )?;
16325                    let up = m.qmatvec_view(
16326                        e,
16327                        &dev.up,
16328                        up_start..up_start + ul.len,
16329                        &gv,
16330                        m_e,
16331                        m.up_exps.in_f,
16332                        m.up_exps.out_f,
16333                        ul.qtype,
16334                        ul.row_bytes,
16335                    )?;
16336                    let mut act = e.uninit(m_e * n_ff_exp)?;
16337                    Self::ffn_act_lim(
16338                        e,
16339                        cfg,
16340                        &gate,
16341                        &up,
16342                        m.gate_exps.macro_scale(ex),
16343                        m.up_exps.macro_scale(ex),
16344                        lim_exp,
16345                        &mut act,
16346                        m_e * n_ff_exp,
16347                    )?;
16348                    let actv = act.slice(0..m_e * n_ff_exp);
16349                    m.qmatvec_view(
16350                        e,
16351                        &dev.down,
16352                        down_start..down_start + dl.len,
16353                        &actv,
16354                        m_e,
16355                        m.down_exps.in_f,
16356                        m.down_exps.out_f,
16357                        dl.qtype,
16358                        dl.row_bytes,
16359                    )?
16360                }
16361            } else if use_cache {
16362                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16363                if grouped_q8 {
16364                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16365                    let gate = e.with_moe_cache(max_block, |cache, eng| {
16366                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
16367                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16368                        eng.qmatvec_expert_q8(
16369                            cache.buf(slot),
16370                            0..gl.len,
16371                            &zq,
16372                            &zd,
16373                            m_e,
16374                            m.gate_exps.in_f,
16375                            m.gate_exps.out_f,
16376                            gl.qtype,
16377                            gl.row_bytes,
16378                        )
16379                    })?;
16380                    let up = e.with_moe_cache(max_block, |cache, eng| {
16381                        let id = BlockId::new(il, PROJ_UP, ex as u16);
16382                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16383                        eng.qmatvec_expert_q8(
16384                            cache.buf(slot),
16385                            0..ul.len,
16386                            &zq,
16387                            &zd,
16388                            m_e,
16389                            m.up_exps.in_f,
16390                            m.up_exps.out_f,
16391                            ul.qtype,
16392                            ul.row_bytes,
16393                        )
16394                    })?;
16395                    let mut act = e.uninit(m_e * n_ff_exp)?;
16396                    Self::ffn_act_lim(
16397                        e,
16398                        cfg,
16399                        &gate,
16400                        &up,
16401                        m.gate_exps.macro_scale(ex),
16402                        m.up_exps.macro_scale(ex),
16403                        lim_exp,
16404                        &mut act,
16405                        m_e * n_ff_exp,
16406                    )?;
16407                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16408                    e.with_moe_cache(max_block, |cache, eng| {
16409                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16410                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16411                        eng.qmatvec_expert_q8(
16412                            cache.buf(slot),
16413                            0..dl.len,
16414                            &aq2,
16415                            &ad2,
16416                            m_e,
16417                            m.down_exps.in_f,
16418                            m.down_exps.out_f,
16419                            dl.qtype,
16420                            dl.row_bytes,
16421                        )
16422                    })?
16423                } else {
16424                    let gate = e.with_moe_cache(max_block, |cache, eng| {
16425                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
16426                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16427                        m.qmatvec_view(
16428                            eng,
16429                            cache.buf(slot),
16430                            0..gl.len,
16431                            &gv,
16432                            m_e,
16433                            m.gate_exps.in_f,
16434                            m.gate_exps.out_f,
16435                            gl.qtype,
16436                            gl.row_bytes,
16437                        )
16438                    })?;
16439                    let up = e.with_moe_cache(max_block, |cache, eng| {
16440                        let id = BlockId::new(il, PROJ_UP, ex as u16);
16441                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16442                        m.qmatvec_view(
16443                            eng,
16444                            cache.buf(slot),
16445                            0..ul.len,
16446                            &gv,
16447                            m_e,
16448                            m.up_exps.in_f,
16449                            m.up_exps.out_f,
16450                            ul.qtype,
16451                            ul.row_bytes,
16452                        )
16453                    })?;
16454                    let mut act = e.uninit(m_e * n_ff_exp)?;
16455                    Self::ffn_act_lim(
16456                        e,
16457                        cfg,
16458                        &gate,
16459                        &up,
16460                        m.gate_exps.macro_scale(ex),
16461                        m.up_exps.macro_scale(ex),
16462                        lim_exp,
16463                        &mut act,
16464                        m_e * n_ff_exp,
16465                    )?;
16466                    let actv = act.slice(0..m_e * n_ff_exp);
16467                    e.with_moe_cache(max_block, |cache, eng| {
16468                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16469                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16470                        m.qmatvec_view(
16471                            eng,
16472                            cache.buf(slot),
16473                            0..dl.len,
16474                            &actv,
16475                            m_e,
16476                            m.down_exps.in_f,
16477                            m.down_exps.out_f,
16478                            dl.qtype,
16479                            dl.row_bytes,
16480                        )
16481                    })?
16482                }
16483            } else {
16484                let sg = scratch_g.as_mut().unwrap();
16485                let su = scratch_u.as_mut().unwrap();
16486                let sd = scratch_d.as_mut().unwrap();
16487                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
16488                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
16489                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
16490                if grouped_q8 {
16491                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16492                    let gate = e.qmatvec_expert_q8(
16493                        sg,
16494                        0..gl.len,
16495                        &zq,
16496                        &zd,
16497                        m_e,
16498                        m.gate_exps.in_f,
16499                        m.gate_exps.out_f,
16500                        gl.qtype,
16501                        gl.row_bytes,
16502                    )?;
16503                    let up = e.qmatvec_expert_q8(
16504                        su,
16505                        0..ul.len,
16506                        &zq,
16507                        &zd,
16508                        m_e,
16509                        m.up_exps.in_f,
16510                        m.up_exps.out_f,
16511                        ul.qtype,
16512                        ul.row_bytes,
16513                    )?;
16514                    let mut act = e.uninit(m_e * n_ff_exp)?;
16515                    Self::ffn_act_lim(
16516                        e,
16517                        cfg,
16518                        &gate,
16519                        &up,
16520                        m.gate_exps.macro_scale(ex),
16521                        m.up_exps.macro_scale(ex),
16522                        lim_exp,
16523                        &mut act,
16524                        m_e * n_ff_exp,
16525                    )?;
16526                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16527                    e.qmatvec_expert_q8(
16528                        sd,
16529                        0..dl.len,
16530                        &aq2,
16531                        &ad2,
16532                        m_e,
16533                        m.down_exps.in_f,
16534                        m.down_exps.out_f,
16535                        dl.qtype,
16536                        dl.row_bytes,
16537                    )?
16538                } else {
16539                    let gate = m.qmatvec_view(
16540                        e,
16541                        sg,
16542                        0..gl.len,
16543                        &gv,
16544                        m_e,
16545                        m.gate_exps.in_f,
16546                        m.gate_exps.out_f,
16547                        gl.qtype,
16548                        gl.row_bytes,
16549                    )?;
16550                    let up = m.qmatvec_view(
16551                        e,
16552                        su,
16553                        0..ul.len,
16554                        &gv,
16555                        m_e,
16556                        m.up_exps.in_f,
16557                        m.up_exps.out_f,
16558                        ul.qtype,
16559                        ul.row_bytes,
16560                    )?;
16561                    let mut act = e.uninit(m_e * n_ff_exp)?;
16562                    Self::ffn_act_lim(
16563                        e,
16564                        cfg,
16565                        &gate,
16566                        &up,
16567                        m.gate_exps.macro_scale(ex),
16568                        m.up_exps.macro_scale(ex),
16569                        lim_exp,
16570                        &mut act,
16571                        m_e * n_ff_exp,
16572                    )?;
16573                    let actv = act.slice(0..m_e * n_ff_exp);
16574                    m.qmatvec_view(
16575                        e,
16576                        sd,
16577                        0..dl.len,
16578                        &actv,
16579                        m_e,
16580                        m.down_exps.in_f,
16581                        m.down_exps.out_f,
16582                        dl.qtype,
16583                        dl.row_bytes,
16584                    )?
16585                }
16586            };
16587
16588            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
16589            e.scatter_slot(
16590                &y,
16591                &tok_idx_d,
16592                &slot_idx_d,
16593                &weight_d,
16594                &mut slot_buf,
16595                &mut wbuf,
16596                n_embd,
16597                n_used,
16598                m_e,
16599            )?;
16600        }
16601
16602        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
16603        let mut moe_out = e.zeros(t * n_embd)?;
16604        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
16605
16606        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
16607        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
16608            m_dist.sort_unstable();
16609            let active = m_dist.len();
16610            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
16611            let median = m_dist[active / 2];
16612            let max_m = *m_dist.last().unwrap();
16613            let min_m = m_dist[0];
16614            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
16615            println!(
16616                "moe-grouped il={il} t={t} active={active}/{n_expert} \
16617                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
16618                      above_gemm_threshold(>=16)={above16}/{active}"
16619            );
16620        }
16621
16622        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
16623        Ok(moe_out)
16624    }
16625
16626    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
16627    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
16628    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
16629    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
16630    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
16631    /// expert-sum order identical to the sequential path.
16632    pub(crate) fn moe_ffn_lockstep(
16633        &self,
16634        e: &Engine,
16635        m: &MoeWeights,
16636        zbatch: &CudaSlice<f32>,
16637        mrows: usize,
16638        il: u16,
16639        max_block: usize,
16640    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16641        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16642        let cfg = &self.cfg;
16643        let moe = cfg.moe.as_ref().unwrap();
16644        let n_embd = cfg.n_embd as usize;
16645        let n_expert = moe.expert_count as usize;
16646        let n_used = moe.expert_used_count as usize;
16647        let n_ff_exp = moe.expert_ff_length as usize;
16648        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
16649        let lim_exp = cfg.clamp_exp_at(il as u32);
16650        let lim_shexp = cfg.clamp_shexp_at(il as u32);
16651
16652        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
16653        if let Some(sig) = cfg.sigmoid_router() {
16654            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
16655        }
16656        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
16657            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
16658        } else {
16659            Self::moe_route_cfg(
16660                e,
16661                &logits,
16662                mrows,
16663                n_expert,
16664                n_used,
16665                m.active_experts.as_deref(),
16666            )?
16667        };
16668        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
16669
16670        // Residency split at whole-expert granularity against the (frozen) cache.
16671        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
16672            Ok((0..n_expert)
16673                .map(|ex| {
16674                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
16675                        .into_iter()
16676                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
16677                })
16678                .collect())
16679        })?;
16680
16681        struct Group {
16682            rows: Vec<i32>,
16683            slots: Vec<i32>,
16684            weights: Vec<f32>,
16685        }
16686        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
16687        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
16688        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
16689            Default::default();
16690        for row in 0..mrows {
16691            for j in 0..n_used {
16692                let ex = sel_all[row * n_used + j] as usize;
16693                let w = w_all[row * n_used + j];
16694                if resident_expert[ex] {
16695                    let group = groups.entry(ex).or_insert_with(|| Group {
16696                        rows: Vec::new(),
16697                        slots: Vec::new(),
16698                        weights: Vec::new(),
16699                    });
16700                    group.rows.push(row as i32);
16701                    group.slots.push(j as i32);
16702                    group.weights.push(w);
16703                } else {
16704                    crate::cpu_experts::record_incomplete_gpu_residency(0);
16705                    cpu_rows[row].push((ex, w));
16706                    cpu_by_expert.entry(ex).or_default().push((row, w));
16707                }
16708            }
16709        }
16710
16711        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
16712        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
16713        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
16714        // order per row differs from the sequential single-call chunk — part of the
16715        // documented lockstep numeric class.
16716        let host_rows = e.dtoh(zbatch)?;
16717        let rows_ok = crate::cpu_experts::rows_supported();
16718        enum CpuPart {
16719            Single { row: usize },
16720            Rows { rows: Vec<usize> },
16721        }
16722        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
16723        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
16724        if rows_ok {
16725            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
16726                .into_iter()
16727                .filter(|(_, rows)| rows.len() >= 2)
16728                .collect();
16729            shared.sort_by_key(|(ex, _)| *ex);
16730            for (ex, mut row_weights) in shared {
16731                row_weights.sort_by_key(|(row, _)| *row);
16732                let inputs: Vec<(&[f32], f32)> = row_weights
16733                    .iter()
16734                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
16735                    .collect();
16736                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
16737                    .map_err(std::io::Error::other)?;
16738                for &(row, _) in &row_weights {
16739                    rows_served.insert((row, ex));
16740                }
16741                tickets.push((
16742                    CpuPart::Rows {
16743                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
16744                    },
16745                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
16746                ));
16747            }
16748        }
16749        for (row, selected) in cpu_rows.iter().enumerate() {
16750            let leftover: Vec<(usize, f32)> = selected
16751                .iter()
16752                .copied()
16753                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
16754                .collect();
16755            if leftover.is_empty() {
16756                continue;
16757            }
16758            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
16759            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
16760                .map_err(std::io::Error::other)?;
16761            tickets.push((
16762                CpuPart::Single { row },
16763                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
16764            ));
16765        }
16766
16767        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
16768        let mut wbuf = e.zeros(mrows * n_used)?;
16769        let mut order: Vec<usize> = groups.keys().copied().collect();
16770        order.sort_by(|&a, &b| {
16771            groups[&b]
16772                .rows
16773                .len()
16774                .cmp(&groups[&a].rows.len())
16775                .then(a.cmp(&b))
16776        });
16777        for &ex in &order {
16778            let group = &groups[&ex];
16779            let m_e = group.rows.len();
16780            let gl = m.gate_exps.expert_layout(ex);
16781            let ul = m.up_exps.expert_layout(ex);
16782            let dl = m.down_exps.expert_layout(ex);
16783            let row_idx_d = e.htod_i32(&group.rows)?;
16784            let slot_idx_d = e.htod_i32(&group.slots)?;
16785            let dmac = m.down_exps.macro_scale(ex);
16786            let weight_d = if dmac == 1.0 {
16787                e.htod(&group.weights)?
16788            } else {
16789                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
16790                e.htod(&scaled)?
16791            };
16792            let mut gathered = e.zeros(m_e * n_embd)?;
16793            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
16794            let gv = gathered.slice(0..m_e * n_embd);
16795            let gate = e.with_moe_cache(max_block, |c, eng| {
16796                let slot = c
16797                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
16798                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16799                m.qmatvec_view(
16800                    eng,
16801                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16802                    0..gl.len,
16803                    &gv,
16804                    m_e,
16805                    m.gate_exps.in_f,
16806                    m.gate_exps.out_f,
16807                    gl.qtype,
16808                    gl.row_bytes,
16809                )
16810            })?;
16811            let up = e.with_moe_cache(max_block, |c, eng| {
16812                let slot = c
16813                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
16814                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16815                m.qmatvec_view(
16816                    eng,
16817                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16818                    0..ul.len,
16819                    &gv,
16820                    m_e,
16821                    m.up_exps.in_f,
16822                    m.up_exps.out_f,
16823                    ul.qtype,
16824                    ul.row_bytes,
16825                )
16826            })?;
16827            let mut act = e.zeros(m_e * n_ff_exp)?;
16828            Self::ffn_act_lim(
16829                e,
16830                cfg,
16831                &gate,
16832                &up,
16833                m.gate_exps.macro_scale(ex),
16834                m.up_exps.macro_scale(ex),
16835                lim_exp,
16836                &mut act,
16837                m_e * n_ff_exp,
16838            )?;
16839            let actv = act.slice(0..m_e * n_ff_exp);
16840            let y = e.with_moe_cache(max_block, |c, eng| {
16841                let slot = c
16842                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
16843                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16844                m.qmatvec_view(
16845                    eng,
16846                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16847                    0..dl.len,
16848                    &actv,
16849                    m_e,
16850                    m.down_exps.in_f,
16851                    m.down_exps.out_f,
16852                    dl.qtype,
16853                    dl.row_bytes,
16854                )
16855            })?;
16856            e.scatter_slot(
16857                &y,
16858                &row_idx_d,
16859                &slot_idx_d,
16860                &weight_d,
16861                &mut slot_buf,
16862                &mut wbuf,
16863                n_embd,
16864                n_used,
16865                m_e,
16866            )?;
16867        }
16868        let mut moe_out = e.zeros(mrows * n_embd)?;
16869        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
16870
16871        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
16872        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
16873        for (part, ticket) in tickets {
16874            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
16875            let mut add_row = |row: usize, chunk: &[f32]| {
16876                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
16877                for (accumulator, value) in sum.iter_mut().zip(chunk) {
16878                    *accumulator += value;
16879                }
16880            };
16881            match part {
16882                CpuPart::Single { row } => add_row(row, &cpu_output),
16883                CpuPart::Rows { rows } => {
16884                    for (slot, row) in rows.into_iter().enumerate() {
16885                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
16886                    }
16887                }
16888            }
16889        }
16890        for (row, sum) in row_sums.into_iter().enumerate() {
16891            let Some(sum) = sum else { continue };
16892            let cpu_output = e.htod(&sum)?;
16893            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
16894            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
16895        }
16896
16897        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
16898            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
16899        {
16900            let n_ff_sh = gate_shexp.out_features();
16901            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
16902            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
16903            let mut sa = e.zeros(mrows * n_ff_sh)?;
16904            Self::ffn_act_lim(
16905                e,
16906                cfg,
16907                &sg_gate,
16908                &sg_up,
16909                1.0,
16910                1.0,
16911                lim_shexp,
16912                &mut sa,
16913                mrows * n_ff_sh,
16914            )?;
16915            let sh = e.matmul(down_shexp, &sa, mrows)?;
16916            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
16917            // decode matches the single-sequence decode chain bit-for-bit.
16918            let g = match &m.gate_inp_shexp {
16919                Some(gate_inp_shexp) => {
16920                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
16921                }
16922                None => e.htod(&vec![1.0f32; mrows])?,
16923            };
16924            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
16925        }
16926
16927        Ok(moe_out)
16928    }
16929}
16930
16931// ============================ gemma4 (R8 verified wiring) ==================================
16932// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
16933// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
16934// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
16935// gemma variants after the correctness gate).
16936impl HybridModel {
16937    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
16938    /// Rotary width the gemma-4 config DECLARES for layer `il`, per SWA/global class — the
16939    /// `n_rot` argument of the fused `rms_norm_qkv_rope*` kernels
16940    /// (`Engine::full_width_rope_only`, lane/graph-s-key-exactness-20260819).
16941    ///
16942    /// It exists so the fusion's full-width assumption is checked against the CONFIG instead of
16943    /// against itself: passing `hd` twice would guard nothing. On every gemma-4 artifact both
16944    /// routes set `rope_dims_{global,swa} == key_length_{global,swa}` (GGUF reads
16945    /// `rope.dimension_count(_swa)`, the safetensors route derives them from
16946    /// `global_head_dim`/`head_dim`), which is exactly WHY the fusion is legal here — gemma's
16947    /// "partial" rotary is per-dimension frequency scaling through the `rope_freqs` factors
16948    /// tensor, not an n_rot truncation. An arch that truncates instead (qwen3.5: 64 of 256;
16949    /// step35 full-attn: 64 of 128) must not reach these kernels, and now cannot silently.
16950    pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
16951        let g = self
16952            .cfg
16953            .gemma4
16954            .as_ref()
16955            .expect("gemma4_rope_dims on a non-gemma4 config");
16956        if g.swa_pattern[il] {
16957            g.rope_dims_swa as usize
16958        } else {
16959            g.rope_dims_global as usize
16960        }
16961    }
16962
16963    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
16964        let g = self.cfg.gemma4.as_ref().unwrap();
16965        let swa = g.swa_pattern[il];
16966        let hd = if swa {
16967            g.key_length_swa
16968        } else {
16969            g.key_length_global
16970        } as usize;
16971        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
16972        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
16973        // rows exact (softmax over one element) while every later position drifted).
16974        (
16975            hd,
16976            g.head_count_kv[il] as usize,
16977            self.cfg.n_head as usize,
16978            if swa {
16979                g.rope_base_swa
16980            } else {
16981                g.rope_base_global
16982            },
16983            1.0,
16984            swa,
16985        )
16986    }
16987
16988    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
16989    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
16990    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
16991    pub(crate) fn gemma4_suppress(
16992        &self,
16993        e: &Engine,
16994        ld: &mut CudaSlice<f32>,
16995        t: usize,
16996    ) -> Result<(), Box<dyn std::error::Error>> {
16997        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
16998            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
16999            // stage as primary, and this tail runs only after the last stage). The assert turns
17000            // that argued invariant into a checked one: any topology violating primary==head
17001            // trips here in debug instead of silently peer-reading a device-0 buffer.
17002            #[cfg(debug_assertions)]
17003            crate::debug_assert_tensor_stream_device(
17004                ids,
17005                &e.stream(),
17006                "gemma4_suppress.suppress_d",
17007            );
17008            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
17009        }
17010        Ok(())
17011    }
17012
17013    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
17014    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
17015    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
17016    /// only (v0): attends within `tokens` via the f32 sdpa.
17017    #[allow(clippy::too_many_arguments)]
17018    /// ONE-PROGRAM seam (lane/gemma-restore-exactness-20260819, DEFAULT OFF): route dense
17019    /// gemma SWA prefill attention through the WINDOWED arm at every prompt length instead of
17020    /// switching program at `t > sliding_window`. The door is the measured cause of the
17021    /// prefix-restore exactness violation: below it the SWA layers attend unwindowed on f32
17022    /// operands, above it windowed on bf16 operands, so rows[0..932) of a 932-row prime and of
17023    /// a 1048-row prime differ in 59 of 60 layers. `=1` collapses it to one program so the
17024    /// published prefix KV stops depending on the total prompt length. Off by default because
17025    /// collapsing changes the bytes short prompts get today; the flip is an owner call.
17026    fn gemma_fa_one_program() -> bool {
17027        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17028        *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
17029    }
17030
17031    #[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
17032    fn gemma4_attn_prime(
17033        &self,
17034        e: &Engine,
17035        fa: &crate::hybrid::FullAttnLayer,
17036        il: usize,
17037        h: &CudaSlice<f32>,
17038        pos_d: &CudaSlice<i32>,
17039        t: usize,
17040        cache: Option<&mut Cache>,
17041        island: Option<&CudaSlice<i32>>,
17042    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17043        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
17044        let eps = self.cfg.rms_eps;
17045        let aux = self.gemma4_aux.as_ref().unwrap();
17046        let ones = aux.ones(e);
17047        #[cfg(debug_assertions)]
17048        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
17049
17050        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
17051        // (h stays borrowed across the triple, so the cache key can't go stale).
17052        e.mmq_act_begin();
17053        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
17054        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
17055            let v = e.dtoh(&q0)?;
17056            let nan = v.iter().filter(|x| x.is_nan()).count();
17057            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
17058            eprintln!(
17059                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
17060                v.len()
17061            );
17062        }
17063        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
17064        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
17065        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
17066        let v0 = if swa {
17067            e.matmul(&fa.wv, h, t)?
17068        } else {
17069            e.clone_dtod(&k0)?
17070        };
17071        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
17072            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
17073                let v = e.dtoh(buf)?;
17074                let nan = v.iter().filter(|x| x.is_nan()).count();
17075                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
17076                eprintln!(
17077                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
17078                    v.len()
17079                );
17080            }
17081        }
17082
17083        let mut q = e.uninit(t * nh * hd)?;
17084        let mut k = e.uninit(t * nkv * hd)?;
17085        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
17086        let mut v = e.uninit(t * nkv * hd)?;
17087        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
17088        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
17089        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
17090        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17091        // Island primes take the mask-capable naive kernel below; keep the operands f32
17092        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
17093        let emit = island.is_none()
17094            && t >= 16
17095            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
17096            && *EMIT.get_or_init(|| {
17097                std::env::var("MEMRA_FA_EMIT")
17098                    .map(|s| s != "0")
17099                    .unwrap_or(true)
17100            });
17101        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
17102        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
17103        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
17104        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
17105        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
17106        let v_f16 = emit
17107            && crate::fa_f16pv_on()
17108            && match hd {
17109                512 => true,
17110                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
17111                _ => false,
17112            };
17113        if emit {
17114            e.rms_norm_qkv_w4b(
17115                &q0,
17116                &k0,
17117                &v0,
17118                fa.q_norm.float_data(),
17119                fa.k_norm.float_data(),
17120                ones,
17121                &mut q,
17122                &mut k,
17123                &mut v,
17124                &mut vb,
17125                hd,
17126                nh * t,
17127                nkv * t,
17128                eps,
17129                v_f16,
17130            )?;
17131        } else {
17132            e.rms_norm_qkv(
17133                &q0,
17134                &k0,
17135                &v0,
17136                fa.q_norm.float_data(),
17137                fa.k_norm.float_data(),
17138                ones,
17139                &mut q,
17140                &mut k,
17141                &mut v,
17142                hd,
17143                nh * t,
17144                nkv * t,
17145                eps,
17146            )?;
17147        }
17148
17149        let ff = if swa {
17150            None
17151        } else {
17152            Some(
17153                aux.rope_freqs(e)
17154                    .expect("gemma4 global rope needs rope_freqs.weight"),
17155            )
17156        };
17157        #[cfg(debug_assertions)]
17158        if let Some(ff) = ff {
17159            crate::debug_assert_tensor_stream_device(
17160                ff,
17161                &e.stream(),
17162                "gemma4_attn_prime.rope_freqs",
17163            );
17164        }
17165        if emit {
17166            e.rope_neox2_bf16e(
17167                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
17168            )?;
17169        } else {
17170            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
17171        }
17172
17173        if let Some(cache) = cache {
17174            let kvl = cache.kv[il].as_mut().unwrap();
17175            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
17176            e.append_kv_quantized_rows(
17177                &k,
17178                &v,
17179                &mut kvl.k,
17180                &mut kvl.v,
17181                kvl.len,
17182                t,
17183                kvl.kv_dim_k,
17184                kvl.kv_dim_v,
17185                kvl.k_tok_bytes,
17186                kvl.v_tok_bytes,
17187                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
17188            )?;
17189            kvl.len += t;
17190        }
17191        let mut attn = e.zeros(t * nh * hd)?;
17192        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
17193        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
17194        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
17195        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
17196        if let Some(span) = island {
17197            // Masked-prefill arm: every layer routes through the island-aware naive
17198            // kernel (correctness-first, same posture as the vision tower v1). The
17199            // window argument keeps the R6 shortcut: 0 while the prompt fits the
17200            // window, the real window beyond it.
17201            let w = if swa && t > win { win } else { 0 };
17202            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
17203        } else if swa && (t > win || Self::gemma_fa_one_program()) {
17204            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17205                if emit {
17206                    e.fa_prefill_w_pre(
17207                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
17208                    )?;
17209                } else {
17210                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17211                }
17212            } else {
17213                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17214            }
17215        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17216            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17217        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
17218            if emit {
17219                e.fa_prefill_hd512_pre(
17220                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
17221                )?;
17222            } else {
17223                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17224            }
17225        } else {
17226            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17227        }
17228        e.matmul(&fa.wo, &attn, t)
17229    }
17230
17231    /// Back-compat wrapper (pure prefill, no cache).
17232    fn gemma4_attn(
17233        &self,
17234        e: &Engine,
17235        fa: &crate::hybrid::FullAttnLayer,
17236        il: usize,
17237        h: &CudaSlice<f32>,
17238        pos_d: &CudaSlice<i32>,
17239        t: usize,
17240    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17241        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
17242    }
17243
17244    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
17245    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
17246    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
17247    /// the q8z epilogue is quantize_q8_1 verbatim).
17248    fn gemma4_moe_q8(
17249        &self,
17250        e: &Engine,
17251        m: &crate::hybrid::MoeWeights,
17252        bits: &crate::hybrid::Gemma4MoeBits,
17253        mq: &(CudaSlice<i8>, CudaSlice<f32>),
17254        router_in: &CudaSlice<f32>,
17255        t: usize,
17256    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17257        let cfg = &self.cfg;
17258        let moe = cfg.moe.as_ref().unwrap();
17259        let n_embd = cfg.n_embd as usize;
17260        let n_expert = moe.expert_count as usize;
17261        let n_used = moe.expert_used_count as usize;
17262        let n_ff_exp = moe.expert_ff_length as usize;
17263        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
17264        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
17265        // the pair's 12us is kernel time, not launch gaps.
17266        let logits = if crate::router_kernel_on() {
17267            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17268        } else {
17269            e.matmul(&m.gate_inp, router_in, t)?
17270        };
17271        let dev = m.dev_exps.as_ref().unwrap();
17272        let (sel_d, w_d) =
17273            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17274        let (zq, zd) = mq;
17275        if t == 1 {
17276            let selv = sel_d.slice(0..n_used);
17277            let wv = w_d.slice(0..n_used);
17278            let act = e.moe_gate_up_gelu8_dev_q8(
17279                &dev.ptr_row,
17280                &selv,
17281                zq,
17282                zd,
17283                n_embd,
17284                n_ff_exp,
17285                n_used,
17286                n_expert,
17287                m.gate_exps.qtype,
17288                m.up_exps.qtype,
17289                m.gate_exps.row_bytes,
17290                m.up_exps.row_bytes,
17291            )?;
17292            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17293            let mut moe_out = e.uninit(n_embd)?;
17294            e.moe_down8_fma_dev_q8(
17295                &dev.ptr_row,
17296                &selv,
17297                &wv,
17298                &aq2,
17299                &ad2,
17300                &mut moe_out.slice_mut(0..n_embd),
17301                n_ff_exp,
17302                n_embd,
17303                n_used,
17304                n_expert,
17305                m.down_exps.qtype,
17306                m.down_exps.row_bytes,
17307            )?;
17308            return Ok(moe_out);
17309        }
17310        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17311        let act = if csr {
17312            e.moe_gate_up_gelu8_dev_q8_csr(
17313                &dev.ptr_row,
17314                &sel_d,
17315                zq,
17316                zd,
17317                t * n_used,
17318                n_embd,
17319                n_ff_exp,
17320                n_used,
17321                n_expert,
17322                m.gate_exps.qtype,
17323                m.up_exps.qtype,
17324                m.gate_exps.row_bytes,
17325                m.up_exps.row_bytes,
17326            )?
17327        } else {
17328            e.moe_gate_up_gelu8_dev_q8_rows(
17329                &dev.ptr_row,
17330                &sel_d,
17331                zq,
17332                zd,
17333                t,
17334                n_embd,
17335                n_ff_exp,
17336                n_used,
17337                n_expert,
17338                m.gate_exps.qtype,
17339                m.up_exps.qtype,
17340                m.gate_exps.row_bytes,
17341                m.up_exps.row_bytes,
17342            )?
17343        };
17344        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17345        let mut moe_out = e.uninit(t * n_embd)?;
17346        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
17347        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
17348        e.moe_down8_fma_dev_q8_rows_g(
17349            &dev.ptr_row,
17350            &sel_d,
17351            &w_d,
17352            &aq2,
17353            &ad2,
17354            &mut moe_out,
17355            t,
17356            n_ff_exp,
17357            n_embd,
17358            n_used,
17359            n_expert,
17360            m.down_exps.qtype,
17361            m.down_exps.row_bytes,
17362        )?;
17363        Ok(moe_out)
17364    }
17365
17366    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
17367    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
17368    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
17369    fn gemma4_moe(
17370        &self,
17371        e: &Engine,
17372        m: &crate::hybrid::MoeWeights,
17373        bits: &crate::hybrid::Gemma4MoeBits,
17374        moe_in: &CudaSlice<f32>,
17375        router_in: &CudaSlice<f32>,
17376        t: usize,
17377    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17378        let cfg = &self.cfg;
17379        let moe = cfg.moe.as_ref().unwrap();
17380        let n_embd = cfg.n_embd as usize;
17381        let n_expert = moe.expert_count as usize;
17382        let n_used = moe.expert_used_count as usize;
17383        let n_ff_exp = moe.expert_ff_length as usize;
17384
17385        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
17386        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
17387        // batched matmul only at real prefill.
17388        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
17389            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17390        } else {
17391            e.matmul(&m.gate_inp, router_in, t)?
17392        };
17393
17394        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
17395        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
17396        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
17397        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
17398        if t < PRIME_MIN_T
17399            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17400            && expert_dp4a_supported(m.gate_exps.qtype)
17401            && expert_dp4a_supported(m.up_exps.qtype)
17402            && expert_dp4a_supported(m.down_exps.qtype)
17403            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
17404        {
17405            let dev = m.dev_exps.as_ref().unwrap();
17406            let (sel_d, w_d) =
17407                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17408            if t == 1 {
17409                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
17410                let selv = sel_d.slice(0..n_used);
17411                let wv = w_d.slice(0..n_used);
17412                let act = e.moe_gate_up_gelu8_dev_q8(
17413                    &dev.ptr_row,
17414                    &selv,
17415                    &zq,
17416                    &zd,
17417                    n_embd,
17418                    n_ff_exp,
17419                    n_used,
17420                    n_expert,
17421                    m.gate_exps.qtype,
17422                    m.up_exps.qtype,
17423                    m.gate_exps.row_bytes,
17424                    m.up_exps.row_bytes,
17425                )?;
17426                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17427                let mut moe_out = e.uninit(n_embd)?;
17428                e.moe_down8_fma_dev_q8(
17429                    &dev.ptr_row,
17430                    &selv,
17431                    &wv,
17432                    &aq2,
17433                    &ad2,
17434                    &mut moe_out.slice_mut(0..n_embd),
17435                    n_ff_exp,
17436                    n_embd,
17437                    n_used,
17438                    n_expert,
17439                    m.down_exps.qtype,
17440                    m.down_exps.row_bytes,
17441                )?;
17442                return Ok(moe_out);
17443            }
17444            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
17445            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
17446            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
17447            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
17448            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17449            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17450            let act = if csr {
17451                e.moe_gate_up_gelu8_dev_q8_csr(
17452                    &dev.ptr_row,
17453                    &sel_d,
17454                    &zq,
17455                    &zd,
17456                    t * n_used,
17457                    n_embd,
17458                    n_ff_exp,
17459                    n_used,
17460                    n_expert,
17461                    m.gate_exps.qtype,
17462                    m.up_exps.qtype,
17463                    m.gate_exps.row_bytes,
17464                    m.up_exps.row_bytes,
17465                )?
17466            } else {
17467                e.moe_gate_up_gelu8_dev_q8_rows(
17468                    &dev.ptr_row,
17469                    &sel_d,
17470                    &zq,
17471                    &zd,
17472                    t,
17473                    n_embd,
17474                    n_ff_exp,
17475                    n_used,
17476                    n_expert,
17477                    m.gate_exps.qtype,
17478                    m.up_exps.qtype,
17479                    m.gate_exps.row_bytes,
17480                    m.up_exps.row_bytes,
17481                )?
17482            };
17483            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17484            let mut moe_out = e.uninit(t * n_embd)?;
17485            e.moe_down8_fma_dev_q8_rows_g(
17486                &dev.ptr_row,
17487                &sel_d,
17488                &w_d,
17489                &aq2,
17490                &ad2,
17491                &mut moe_out,
17492                t,
17493                n_ff_exp,
17494                n_embd,
17495                n_used,
17496                n_expert,
17497                m.down_exps.qtype,
17498                m.down_exps.row_bytes,
17499            )?;
17500            return Ok(moe_out);
17501        }
17502
17503        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
17504        for (i, &sx) in sel_all.iter().enumerate() {
17505            w_all[i] *= bits.per_expert_scale[sx as usize];
17506        }
17507
17508        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
17509        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
17510        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
17511        if t >= PRIME_MIN_T
17512            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17513            && expert_dp4a_supported(m.gate_exps.qtype)
17514            && expert_dp4a_supported(m.up_exps.qtype)
17515            && expert_dp4a_supported(m.down_exps.qtype)
17516            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
17517        {
17518            let dev = m.dev_exps.as_ref().unwrap();
17519            let n_pairs = t * n_used;
17520            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
17521            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
17522            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
17523            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
17524            let pt = e.htod_i32(&pair_tok)?;
17525            let pw = e.htod(&w_all)?;
17526            let toff = e.htod_i32(&tok_off)?;
17527            let tids = e.htod_i32(&tok_ids)?;
17528            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
17529            for p in 0..n_pairs {
17530                by_ex[pair_ex[p] as usize].push(p as i32);
17531            }
17532            let mut ex_ids: Vec<i32> = Vec::new();
17533            let mut ex_off: Vec<i32> = vec![0];
17534            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
17535            for (ex, list) in by_ex.iter().enumerate() {
17536                if list.is_empty() {
17537                    continue;
17538                }
17539                ex_ids.push(ex as i32);
17540                ex_pairs.extend_from_slice(list);
17541                ex_off.push(ex_pairs.len() as i32);
17542            }
17543            let n_active = ex_ids.len();
17544            let exi = e.htod_i32(&ex_ids)?;
17545            let exo = e.htod_i32(&ex_off)?;
17546            let exp_d = e.htod_i32(&ex_pairs)?;
17547            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
17548            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
17549            // end-to-end (gelu is elementwise), one row permute before the scatter. The
17550            // ragged down k (704) needs no padding here — cublas takes any k.
17551            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
17552            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
17553            // Hopper default — see moe_f16g_gemma_on.
17554            if crate::moe_f16g_gemma_on()
17555                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
17556                && f16g_proj_ok(m.up_exps.qtype, n_embd)
17557                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
17558            {
17559                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
17560                let csr_tok_d = e.htod_i32(&csr_tok)?;
17561                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
17562                let g_csr = e.moe_f16_grouped(
17563                    &dev.ptr_row,
17564                    0,
17565                    n_expert,
17566                    &exi,
17567                    &ex_off,
17568                    &exo,
17569                    &z_f16,
17570                    &z_s,
17571                    n_embd,
17572                    n_ff_exp,
17573                    n_active,
17574                    n_pairs,
17575                    m.gate_exps.qtype,
17576                    m.gate_exps.row_bytes,
17577                )?;
17578                let u_csr = e.moe_f16_grouped(
17579                    &dev.ptr_row,
17580                    1,
17581                    n_expert,
17582                    &exi,
17583                    &ex_off,
17584                    &exo,
17585                    &z_f16,
17586                    &z_s,
17587                    n_embd,
17588                    n_ff_exp,
17589                    n_active,
17590                    n_pairs,
17591                    m.up_exps.qtype,
17592                    m.up_exps.row_bytes,
17593                )?;
17594                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
17595                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
17596                let d_csr = e.moe_f16_grouped(
17597                    &dev.ptr_row,
17598                    2,
17599                    n_expert,
17600                    &exi,
17601                    &ex_off,
17602                    &exo,
17603                    &a_f16,
17604                    &a_s,
17605                    n_ff_exp,
17606                    n_embd,
17607                    n_active,
17608                    n_pairs,
17609                    m.down_exps.qtype,
17610                    m.down_exps.row_bytes,
17611                )?;
17612                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
17613                let mut moe_out = e.uninit(t * n_embd)?;
17614                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17615                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
17616                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
17617                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
17618                    eprintln!(
17619                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
17620                        scan(&yd),
17621                        scan(&mo)
17622                    );
17623                }
17624                return Ok(moe_out);
17625            }
17626            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
17627            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
17628            let mma = n_embd.is_multiple_of(256)
17629                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
17630            let (gate, up) = if mma {
17631                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
17632                (
17633                    e.mmq_iq_experts(
17634                        &dev.ptr_row,
17635                        0,
17636                        n_expert,
17637                        &exi,
17638                        &exo,
17639                        &exp_d,
17640                        &pt,
17641                        &z_scr,
17642                        n_embd,
17643                        n_ff_exp,
17644                        n_active,
17645                        n_pairs,
17646                        t,
17647                        m.gate_exps.qtype,
17648                        m.gate_exps.row_bytes,
17649                    )?,
17650                    e.mmq_iq_experts(
17651                        &dev.ptr_row,
17652                        1,
17653                        n_expert,
17654                        &exi,
17655                        &exo,
17656                        &exp_d,
17657                        &pt,
17658                        &z_scr,
17659                        n_embd,
17660                        n_ff_exp,
17661                        n_active,
17662                        n_pairs,
17663                        t,
17664                        m.up_exps.qtype,
17665                        m.up_exps.row_bytes,
17666                    )?,
17667                )
17668            } else {
17669                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17670                (
17671                    e.moe_pairs_matvec_q8_dec(
17672                        &dev.ptr_row,
17673                        0,
17674                        &exi,
17675                        &exo,
17676                        &exp_d,
17677                        &pt,
17678                        &zq,
17679                        &zd,
17680                        n_embd,
17681                        n_ff_exp,
17682                        n_expert,
17683                        n_active,
17684                        n_pairs,
17685                        m.gate_exps.qtype,
17686                        m.gate_exps.row_bytes,
17687                    )?,
17688                    e.moe_pairs_matvec_q8_dec(
17689                        &dev.ptr_row,
17690                        1,
17691                        &exi,
17692                        &exo,
17693                        &exp_d,
17694                        &pt,
17695                        &zq,
17696                        &zd,
17697                        n_embd,
17698                        n_ff_exp,
17699                        n_expert,
17700                        n_active,
17701                        n_pairs,
17702                        m.up_exps.qtype,
17703                        m.up_exps.row_bytes,
17704                    )?,
17705                )
17706            };
17707            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
17708            let pself = e.htod_i32(&pair_self)?;
17709            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
17710            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
17711            // to the 256-val superblock (768) while the act quantizer's zero padding
17712            // makes every padded-k product exactly zero (weight overread bytes multiply
17713            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
17714            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
17715            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
17716            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
17717            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
17718            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
17719            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
17720            let y_down = if mma {
17721                let in_pad = n_ff_exp.div_ceil(256) * 256;
17722                let a_scr = if crate::moe_fuse_actq_on() {
17723                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
17724                } else {
17725                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17726                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
17727                };
17728                e.mmq_iq_experts(
17729                    &dev.ptr_row,
17730                    2,
17731                    n_expert,
17732                    &exi,
17733                    &exo,
17734                    &exp_d,
17735                    &pself,
17736                    &a_scr,
17737                    in_pad,
17738                    n_embd,
17739                    n_active,
17740                    n_pairs,
17741                    n_pairs,
17742                    m.down_exps.qtype,
17743                    m.down_exps.row_bytes,
17744                )?
17745            } else {
17746                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17747                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
17748                e.moe_pairs_matvec_q8_dec(
17749                    &dev.ptr_row,
17750                    2,
17751                    &exi,
17752                    &exo,
17753                    &exp_d,
17754                    &pself,
17755                    &aq2,
17756                    &ad2,
17757                    n_ff_exp,
17758                    n_embd,
17759                    n_expert,
17760                    n_active,
17761                    n_pairs,
17762                    m.down_exps.qtype,
17763                    m.down_exps.row_bytes,
17764                )?
17765            };
17766            let mut moe_out = e.uninit(t * n_embd)?;
17767            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17768            return Ok(moe_out);
17769        }
17770
17771        let g_len = m.gate_exps.expert_stride;
17772        let u_len = m.up_exps.expert_stride;
17773        let d_len = m.down_exps.expert_stride;
17774        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
17775        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
17776        // the spill fallback.
17777        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
17778        let (mut sg, mut su, mut sd) = if dev.is_some() {
17779            (None, None, None)
17780        } else {
17781            (
17782                Some(e.alloc_u8_uninit(g_len)?),
17783                Some(e.alloc_u8_uninit(u_len)?),
17784                Some(e.alloc_u8_uninit(d_len)?),
17785            )
17786        };
17787        let mut moe_out = e.zeros(t * n_embd)?;
17788        for tok in 0..t {
17789            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
17790            let w = &w_all[tok * n_used..(tok + 1) * n_used];
17791            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
17792            for (j, &ex) in sel.iter().enumerate() {
17793                let ex = ex as usize;
17794                let gate = match dev {
17795                    Some(d) => m.qmatvec_view(
17796                        e,
17797                        &d.gate,
17798                        ex * g_len..(ex + 1) * g_len,
17799                        &zt,
17800                        1,
17801                        m.gate_exps.in_f,
17802                        m.gate_exps.out_f,
17803                        m.gate_exps.qtype,
17804                        m.gate_exps.row_bytes,
17805                    )?,
17806                    None => {
17807                        let sg = sg.as_mut().unwrap();
17808                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
17809                        m.qmatvec_view(
17810                            e,
17811                            sg,
17812                            0..g_len,
17813                            &zt,
17814                            1,
17815                            m.gate_exps.in_f,
17816                            m.gate_exps.out_f,
17817                            m.gate_exps.qtype,
17818                            m.gate_exps.row_bytes,
17819                        )?
17820                    }
17821                };
17822                let up = match dev {
17823                    Some(d) => m.qmatvec_view(
17824                        e,
17825                        &d.up,
17826                        ex * u_len..(ex + 1) * u_len,
17827                        &zt,
17828                        1,
17829                        m.up_exps.in_f,
17830                        m.up_exps.out_f,
17831                        m.up_exps.qtype,
17832                        m.up_exps.row_bytes,
17833                    )?,
17834                    None => {
17835                        let su = su.as_mut().unwrap();
17836                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
17837                        m.qmatvec_view(
17838                            e,
17839                            su,
17840                            0..u_len,
17841                            &zt,
17842                            1,
17843                            m.up_exps.in_f,
17844                            m.up_exps.out_f,
17845                            m.up_exps.qtype,
17846                            m.up_exps.row_bytes,
17847                        )?
17848                    }
17849                };
17850                let mut act = e.uninit(n_ff_exp)?;
17851                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
17852                let actv = act.slice(0..n_ff_exp);
17853                let y = match dev {
17854                    Some(d) => m.qmatvec_view(
17855                        e,
17856                        &d.down,
17857                        ex * d_len..(ex + 1) * d_len,
17858                        &actv,
17859                        1,
17860                        m.down_exps.in_f,
17861                        m.down_exps.out_f,
17862                        m.down_exps.qtype,
17863                        m.down_exps.row_bytes,
17864                    )?,
17865                    None => {
17866                        let sd = sd.as_mut().unwrap();
17867                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
17868                        m.qmatvec_view(
17869                            e,
17870                            sd,
17871                            0..d_len,
17872                            &actv,
17873                            1,
17874                            m.down_exps.in_f,
17875                            m.down_exps.out_f,
17876                            m.down_exps.qtype,
17877                            m.down_exps.row_bytes,
17878                        )?
17879                    }
17880                };
17881                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
17882                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
17883            }
17884        }
17885        Ok(moe_out)
17886    }
17887
17888    /// One gemma4 trunk layer (R8): x -> x_next.
17889    fn gemma4_layer(
17890        &self,
17891        e: &Engine,
17892        il: usize,
17893        layer: &crate::hybrid::HybridLayer,
17894        x: &CudaSlice<f32>,
17895        pos_d: &CudaSlice<i32>,
17896        t: usize,
17897    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17898        let n_embd = self.cfg.n_embd as usize;
17899        let eps = self.cfg.rms_eps;
17900
17901        let mut h = e.zeros(t * n_embd)?;
17902        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
17903        let Mixer::Full(fa) = &layer.mixer else {
17904            panic!("gemma4 layer {il} not full-attn")
17905        };
17906        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
17907        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
17908        let mut cur = e.zeros(t * n_embd)?;
17909        e.rms_norm(
17910            &o,
17911            layer.post_attn_norm.float_data(),
17912            &mut cur,
17913            n_embd,
17914            t,
17915            eps,
17916        )?;
17917        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
17918    }
17919
17920    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
17921    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
17922    /// layer scale — shared verbatim by the prefill, decode and verify paths.
17923    fn gemma4_layer_tail_add(
17924        &self,
17925        e: &Engine,
17926        layer: &crate::hybrid::HybridLayer,
17927        cur: &CudaSlice<f32>,
17928        x: &CudaSlice<f32>,
17929        t: usize,
17930    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17931        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
17932    }
17933
17934    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
17935    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
17936    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17937    fn gemma4_layer_tail_add_n(
17938        &self,
17939        e: &Engine,
17940        layer: &crate::hybrid::HybridLayer,
17941        cur: &CudaSlice<f32>,
17942        x: &CudaSlice<f32>,
17943        t: usize,
17944        next_norm: Option<&CudaSlice<f32>>,
17945    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
17946        let n_embd = self.cfg.n_embd as usize;
17947        let bits = layer.gemma4.as_ref().unwrap();
17948        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
17949        let mut xn = e.uninit(t * n_embd)?;
17950        match next_norm {
17951            Some(w) => {
17952                let mut hn = e.uninit(t * n_embd)?;
17953                e.add_scale_rms_norm(
17954                    &sn,
17955                    &attn_out,
17956                    bits.layer_scale,
17957                    w,
17958                    &mut xn,
17959                    &mut hn,
17960                    n_embd,
17961                    t,
17962                    self.cfg.rms_eps,
17963                )?;
17964                Ok((xn, Some(hn)))
17965            }
17966            None => {
17967                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
17968                Ok((xn, None))
17969            }
17970        }
17971    }
17972
17973    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
17974    /// norm — returns (sn, attn_out) for the closing add+scale variants.
17975    fn gemma4_layer_tail_core(
17976        &self,
17977        e: &Engine,
17978        layer: &crate::hybrid::HybridLayer,
17979        cur: &CudaSlice<f32>,
17980        x: &CudaSlice<f32>,
17981        t: usize,
17982    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17983        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
17984    }
17985
17986    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
17987    /// means `cur` is the RAW attention output and the dense entry runs
17988    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
17989    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
17990    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
17991    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
17992    #[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
17993    fn gemma4_layer_tail_core_pn(
17994        &self,
17995        e: &Engine,
17996        layer: &crate::hybrid::HybridLayer,
17997        cur: &CudaSlice<f32>,
17998        x: &CudaSlice<f32>,
17999        t: usize,
18000        pre_norm: Option<&CudaSlice<f32>>,
18001        defer_post_norm: bool,
18002    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18003        let n_embd = self.cfg.n_embd as usize;
18004        let eps = self.cfg.rms_eps;
18005        let bits = layer.gemma4.as_ref().unwrap();
18006
18007        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
18008        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
18009        let Some(mbits) = bits.moe_bits.as_ref() else {
18010            let crate::hybrid::Ffn::Dense {
18011                ffn_gate,
18012                ffn_up,
18013                ffn_down,
18014            } = &layer.ffn
18015            else {
18016                panic!("gemma4 dense layer without Dense ffn")
18017            };
18018            let mut attn_out = e.uninit(t * n_embd)?;
18019            let mut zsh = e.uninit(t * n_embd)?;
18020            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
18021            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
18022            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
18023            match pre_norm {
18024                Some(wa) if t == 1 => {
18025                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
18026                        cur,
18027                        wa,
18028                        x,
18029                        bits.ffn_norm.float_data(),
18030                        &mut attn_out,
18031                        &mut zsh,
18032                        n_embd,
18033                        t,
18034                        eps,
18035                    )?);
18036                }
18037                Some(wa) => e.rms_pre_add_rms_norm(
18038                    cur,
18039                    wa,
18040                    x,
18041                    bits.ffn_norm.float_data(),
18042                    &mut attn_out,
18043                    &mut zsh,
18044                    n_embd,
18045                    t,
18046                    eps,
18047                )?,
18048                None => e.add_rms_norm(
18049                    cur,
18050                    x,
18051                    bits.ffn_norm.float_data(),
18052                    &mut attn_out,
18053                    &mut zsh,
18054                    n_embd,
18055                    t,
18056                    eps,
18057                )?,
18058            }
18059            let n_ff = ffn_gate.out_features();
18060            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
18061            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
18062            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
18063            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
18064            // rescue segment C — the megakernel front is closed for the dense tail.
18065            let (gate, up) = if t == 1 {
18066                let (zq, zd) = match zpair {
18067                    Some(p) => p,
18068                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
18069                };
18070                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
18071                    Some(p) => p,
18072                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
18073                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
18074                        Some(p) => p,
18075                        None => (
18076                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
18077                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
18078                        ),
18079                    },
18080                }
18081            } else {
18082                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
18083                // launch for the verify's gate+up — the up segment's blocks fill SMs as
18084                // the gate segment drains (the launch-tail mechanism behind the b-tier
18085                // plateau; first positive after six falsified in-kernel variants).
18086                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18087                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
18088                let fused = if f2b {
18089                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
18090                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
18091                } else {
18092                    None
18093                };
18094                match fused {
18095                    Some(p) => p,
18096                    None => {
18097                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
18098                        e.mmq_act_begin();
18099                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
18100                    }
18101                }
18102            };
18103            let mut act = e.uninit(t * n_ff)?;
18104            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
18105            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
18106            let f0 = if e.uses_q8_1_fast(ffn_down) {
18107                let upv = e.view(&up, t * n_ff);
18108                let up_all = upv.slice(0..t * n_ff);
18109                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
18110                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
18111            } else {
18112                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
18113                e.matmul(ffn_down, &act, t)?
18114            };
18115            if defer_post_norm {
18116                return Ok((f0, attn_out));
18117            }
18118            let mut sn = e.uninit(t * n_embd)?;
18119            e.rms_norm(
18120                &f0,
18121                bits.post_ffw_norm.float_data(),
18122                &mut sn,
18123                n_embd,
18124                t,
18125                eps,
18126            )?;
18127            return Ok((sn, attn_out));
18128        };
18129
18130        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
18131        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
18132        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
18133        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
18134        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
18135        let mut attn_out = e.uninit(t * n_embd)?;
18136        let mut router_in = e.uninit(t * n_embd)?;
18137        let fast_moe = match &layer.ffn {
18138            crate::hybrid::Ffn::Moe(m) => {
18139                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
18140                    && expert_dp4a_supported(m.gate_exps.qtype)
18141                    && expert_dp4a_supported(m.up_exps.qtype)
18142                    && expert_dp4a_supported(m.down_exps.qtype)
18143                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
18144            }
18145            _ => false,
18146        };
18147        let q8z = t < PRIME_MIN_T && fast_moe;
18148        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
18149            let (z0, m2) = e.add_rms_norm3_q8z(
18150                cur,
18151                x,
18152                bits.ffn_norm.float_data(),
18153                &mbits.router_scale_pre,
18154                mbits.pre_ffw_norm_2.float_data(),
18155                &mut attn_out,
18156                &mut router_in,
18157                n_embd,
18158                t,
18159                eps,
18160            )?;
18161            (None, Some(z0), Some(m2))
18162        } else {
18163            let mut zsh = e.uninit(t * n_embd)?;
18164            let mut moe_in = e.uninit(t * n_embd)?;
18165            e.add_rms_norm3(
18166                cur,
18167                x,
18168                bits.ffn_norm.float_data(),
18169                &mbits.router_scale_pre,
18170                mbits.pre_ffw_norm_2.float_data(),
18171                &mut attn_out,
18172                &mut zsh,
18173                &mut router_in,
18174                &mut moe_in,
18175                n_embd,
18176                t,
18177                eps,
18178            )?;
18179            (Some((zsh, moe_in)), None, None)
18180        };
18181        let attn_out2 = attn_out;
18182        #[allow(unused_variables)]
18183        let attn_out = &attn_out2;
18184        let n_ff = mbits.shared_gate.out_features();
18185        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
18186            if t == 1 {
18187                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
18188                    Some(p) => p,
18189                    None => match e.matmul_nvfp4_fused2(
18190                        &mbits.shared_gate,
18191                        &mbits.shared_up,
18192                        zq,
18193                        zd,
18194                        1,
18195                    )? {
18196                        Some(p) => p,
18197                        None => {
18198                            let h0 = e.zeros(0)?;
18199                            (
18200                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
18201                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
18202                            )
18203                        }
18204                    },
18205                }
18206            } else {
18207                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
18208                let h0 = e.zeros(0)?;
18209                (
18210                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
18211                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
18212                )
18213            }
18214        } else {
18215            let (zsh, _) = zsh_f32.as_ref().unwrap();
18216            (
18217                e.matmul(&mbits.shared_gate, zsh, t)?,
18218                e.matmul(&mbits.shared_up, zsh, t)?,
18219            )
18220        };
18221        let mut act = e.uninit(t * n_ff)?;
18222        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
18223        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
18224        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
18225            panic!("gemma4 layer not MoE")
18226        };
18227        let moe0 = match (&moe_q8, &zsh_f32) {
18228            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
18229            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
18230            _ => unreachable!(),
18231        };
18232        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
18233        let mut mlp = e.uninit(t * n_embd)?;
18234        let mut moe = e.uninit(t * n_embd)?;
18235        e.rms_norm2x(
18236            &mlp0,
18237            &moe0,
18238            mbits.post_ffw_norm_1.float_data(),
18239            mbits.post_ffw_norm_2.float_data(),
18240            &mut mlp,
18241            &mut moe,
18242            n_embd,
18243            t,
18244            eps,
18245        )?;
18246
18247        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
18248        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
18249        let mut sum = e.uninit(t * n_embd)?;
18250        let mut sn = e.uninit(t * n_embd)?;
18251        e.add_rms_norm(
18252            &mlp,
18253            &moe,
18254            bits.post_ffw_norm.float_data(),
18255            &mut sum,
18256            &mut sn,
18257            n_embd,
18258            t,
18259            eps,
18260        )?;
18261        Ok((sn, attn_out2))
18262    }
18263
18264    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
18265    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
18266    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
18267    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
18268    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
18269    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
18270    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
18271    /// decode == verify == graph parity holds by construction at either seam value.
18272    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
18273    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18274    pub(crate) fn gemma4_layer_tail_add_nq_pn(
18275        &self,
18276        e: &Engine,
18277        layer: &crate::hybrid::HybridLayer,
18278        o: &CudaSlice<f32>,
18279        x: &CudaSlice<f32>,
18280        t: usize,
18281        next_norm: Option<&CudaSlice<f32>>,
18282    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18283    {
18284        let n_embd = self.cfg.n_embd as usize;
18285        let eps = self.cfg.rms_eps;
18286        let bits = layer.gemma4.as_ref().unwrap();
18287        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
18288            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
18289                e,
18290                layer,
18291                o,
18292                x,
18293                t,
18294                Some(layer.post_attn_norm.float_data()),
18295                true,
18296            )?;
18297            let mut xn = e.uninit(t * n_embd)?;
18298            return match next_norm {
18299                Some(w) => {
18300                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
18301                        &f0,
18302                        bits.post_ffw_norm.float_data(),
18303                        &attn_out,
18304                        bits.layer_scale,
18305                        w,
18306                        &mut xn,
18307                        n_embd,
18308                        t,
18309                        eps,
18310                    )?;
18311                    Ok((xn, Some(pair)))
18312                }
18313                None => {
18314                    let mut sn = e.uninit(t * n_embd)?;
18315                    e.rms_norm(
18316                        &f0,
18317                        bits.post_ffw_norm.float_data(),
18318                        &mut sn,
18319                        n_embd,
18320                        t,
18321                        eps,
18322                    )?;
18323                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18324                    Ok((xn, None))
18325                }
18326            };
18327        }
18328        let mut cur = e.uninit(t * n_embd)?;
18329        e.rms_norm(
18330            o,
18331            layer.post_attn_norm.float_data(),
18332            &mut cur,
18333            n_embd,
18334            t,
18335            eps,
18336        )?;
18337        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
18338    }
18339
18340    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18341    pub(crate) fn gemma4_layer_tail_add_nq(
18342        &self,
18343        e: &Engine,
18344        layer: &crate::hybrid::HybridLayer,
18345        cur: &CudaSlice<f32>,
18346        x: &CudaSlice<f32>,
18347        t: usize,
18348        next_norm: Option<&CudaSlice<f32>>,
18349    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18350    {
18351        let n_embd = self.cfg.n_embd as usize;
18352        let bits = layer.gemma4.as_ref().unwrap();
18353        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
18354        let mut xn = e.uninit(t * n_embd)?;
18355        match next_norm {
18356            Some(w) => {
18357                let pair = e.add_scale_rms_norm_q8_1(
18358                    &sn,
18359                    &attn_out,
18360                    bits.layer_scale,
18361                    w,
18362                    &mut xn,
18363                    n_embd,
18364                    t,
18365                    self.cfg.rms_eps,
18366                )?;
18367                Ok((xn, Some(pair)))
18368            }
18369            None => {
18370                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18371                Ok((xn, None))
18372            }
18373        }
18374    }
18375
18376    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
18377    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
18378    fn gemma4_forward(
18379        &self,
18380        e: &Engine,
18381        tokens: &[u32],
18382        last_only: bool,
18383    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
18384        // E4B routes to its own forward regardless of the caller's entry point (forward /
18385        // forward_last / prime paths all funnel here for gemma4).
18386        if self.is_gemma4_e4b() {
18387            return self.gemma4_e4b_forward(e, tokens, last_only);
18388        }
18389        let n_embd = self.cfg.n_embd as usize;
18390        let t = tokens.len();
18391        let pos: Vec<i32> = (0..t as i32).collect();
18392        let pos_d = e.htod_i32(&pos)?;
18393
18394        let mut x = self.embed(e, tokens)?;
18395        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18396        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
18397        // the bring-up bisect vs llama-eval-callback node stats.
18398        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
18399        let stat =
18400            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
18401                let h = e.dtoh(x)?;
18402                let bad = h.iter().filter(|v| !v.is_finite()).count();
18403                let mx = h
18404                    .iter()
18405                    .filter(|v| v.is_finite())
18406                    .fold(0.0f32, |m, v| m.max(v.abs()));
18407                eprintln!(
18408                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
18409                    &h[..3]
18410                );
18411                Ok(())
18412            };
18413        if probe {
18414            stat(e, &x, "embed")?;
18415        }
18416        for (il, layer) in self.layers.iter().enumerate() {
18417            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
18418            if probe {
18419                stat(e, &x, &format!("L{il}"))?;
18420            }
18421        }
18422        let mut hn = e.zeros(t * n_embd)?;
18423        e.rms_norm(
18424            &x,
18425            self.output_norm.float_data(),
18426            &mut hn,
18427            n_embd,
18428            t,
18429            self.cfg.rms_eps,
18430        )?;
18431        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18432        let n_vocab = self.output.out_features();
18433        let logits = if last_only {
18434            let hv = e.view(&hn, t * n_embd);
18435            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
18436            let mut hlast = e.zeros(n_embd)?;
18437            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
18438            let mut ld = e.matmul(&self.output, &hlast, 1)?;
18439            e.softcap(&mut ld, cap, n_vocab)?;
18440            self.gemma4_suppress(e, &mut ld, 1)?;
18441            e.dtoh(&ld)?
18442        } else {
18443            let mut ld = e.matmul(&self.output, &hn, t)?;
18444            e.softcap(&mut ld, cap, t * n_vocab)?;
18445            self.gemma4_suppress(e, &mut ld, t)?;
18446            e.dtoh(&ld)?
18447        };
18448        Ok(logits)
18449    }
18450
18451    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
18452    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
18453    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
18454    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
18455    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18456    pub(crate) fn gemma4_prime(
18457        &self,
18458        e: &Engine,
18459        tokens: &[u32],
18460        cache: &mut Cache,
18461        overlay: Option<&crate::vision::EmbedOverlay>,
18462    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18463        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
18464        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
18465        // whole worker process on this line. The worker now primes gemma4 monolithically and
18466        // routes continuation suffixes tokenwise; this is the per-request backstop.
18467        if cache.pos != 0 {
18468            return Err(
18469                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
18470                        — prime the full prompt in one call or decode tokenwise"
18471                    .into(),
18472            );
18473        }
18474        let n_embd = self.cfg.n_embd as usize;
18475        let eps = self.cfg.rms_eps;
18476        let t = tokens.len();
18477        let pos: Vec<i32> = (0..t as i32).collect();
18478        let pos_d = e.htod_i32(&pos)?;
18479        let mut x = self.embed(e, tokens)?;
18480        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18481        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
18482        // sqrt(n_embd) text scale — the reference scales token batches only
18483        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
18484        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
18485        // bidirectional within itself, causal+SWA everywhere else, matching the
18486        // reference's llama_set_causal_attn(false) image batch exactly.
18487        let island: Option<CudaSlice<i32>> = match overlay {
18488            Some(ov) => {
18489                // The residency law reaches this arm too (lane/glm53-vision-ppn): the splice
18490                // arithmetic below differs from the shared helper on purpose (post-scale
18491                // placement + island ids), but the pointer it reads obeys the same rule.
18492                ov.require_resident(e)?;
18493                let mut span_id = vec![-1i32; t];
18494                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
18495                    if pos + n_rows > t {
18496                        return Err(format!(
18497                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
18498                            pos + n_rows
18499                        )
18500                        .into());
18501                    }
18502                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
18503                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
18504                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
18505                        *s = i as i32;
18506                    }
18507                }
18508                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
18509                // keep the plain causal mask. Exists only so the decisive probe can show
18510                // the island mask itself changes the answer; never on in serving.
18511                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
18512                    None
18513                } else {
18514                    Some(e.htod_i32(&span_id)?)
18515                }
18516            }
18517            None => None,
18518        };
18519        for (il, layer) in self.layers.iter().enumerate() {
18520            let mut h = e.zeros(t * n_embd)?;
18521            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
18522            let Mixer::Full(fa) = &layer.mixer else {
18523                panic!("gemma4 layer not full-attn")
18524            };
18525            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
18526            if trace {
18527                let v = e.dtoh(&h)?;
18528                let nan = v.iter().filter(|x| x.is_nan()).count();
18529                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
18530            }
18531            let o =
18532                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
18533            if trace {
18534                let v = e.dtoh(&o)?;
18535                let nan = v.iter().filter(|x| x.is_nan()).count();
18536                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
18537            }
18538            let mut cur = e.zeros(t * n_embd)?;
18539            e.rms_norm(
18540                &o,
18541                layer.post_attn_norm.float_data(),
18542                &mut cur,
18543                n_embd,
18544                t,
18545                eps,
18546            )?;
18547            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
18548            self.dflash_tap(e, cache, il, &x, t)?;
18549            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
18550            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
18551                let h = e.dtoh(&x)?;
18552                let nan = h.iter().filter(|v| v.is_nan()).count();
18553                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
18554                eprintln!(
18555                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
18556                    h.len()
18557                );
18558                if nan > 0 {
18559                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
18560                }
18561            }
18562        }
18563        cache.pos += t;
18564        let hiddens = e.clone_dtod(&x)?;
18565        let xv = e.view(&x, t * n_embd);
18566        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
18567        let mut h_seed = e.zeros(n_embd)?;
18568        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
18569        let mut hn = e.uninit(n_embd)?;
18570        e.rms_norm(
18571            &h_seed,
18572            self.output_norm.float_data(),
18573            &mut hn,
18574            n_embd,
18575            1,
18576            eps,
18577        )?;
18578        let mut ld = e.matmul(&self.output, &hn, 1)?;
18579        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18580        e.softcap(&mut ld, cap, self.output.out_features())?;
18581        self.gemma4_suppress(e, &mut ld, 1)?;
18582        let logits = e.dtoh(&ld)?;
18583        Ok((logits, h_seed, hiddens))
18584    }
18585
18586    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
18587    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
18588    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
18589    /// fused norm emits q8 directly — the f32 h never materializes).
18590    #[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
18591    fn gemma4_decode_attn(
18592        &self,
18593        e: &Engine,
18594        fa: &crate::hybrid::FullAttnLayer,
18595        il: usize,
18596        hq: &CudaSlice<i8>,
18597        hdq: &CudaSlice<f32>,
18598        pos_d: &CudaSlice<i32>,
18599        cache: &mut Cache,
18600    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18601        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18602        let eps = self.cfg.rms_eps;
18603        let aux = self.gemma4_aux.as_ref().unwrap();
18604        let ones = aux.ones(e);
18605        #[cfg(debug_assertions)]
18606        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
18607        let (hq, hdq) = (hq, hdq);
18608        let h0 = e.zeros(0)?;
18609        let h = &h0;
18610        let (q0, k0, v0) = if swa {
18611            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
18612                Some(t3) => t3,
18613                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
18614                // match — fuse the uniform (q,k) pair and take v as its own single.
18615                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18616                    Some((q0, k0)) => {
18617                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, h, 1)?;
18618                        (q0, k0, v0)
18619                    }
18620                    None => (
18621                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18622                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18623                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
18624                    ),
18625                },
18626            }
18627        } else {
18628            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
18629                Some(p) => p,
18630                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18631                    Some(p) => p,
18632                    None => (
18633                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18634                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18635                    ),
18636                },
18637            };
18638            let v0 = e.clone_dtod(&k0)?;
18639            (q0, k0, v0)
18640        };
18641        let mut q = e.uninit(nh * hd)?;
18642        let mut k = e.uninit(nkv * hd)?;
18643        let mut v = e.uninit(nkv * hd)?;
18644        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
18645        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
18646        let ff = if swa {
18647            None
18648        } else {
18649            Some(
18650                aux.rope_freqs(e)
18651                    .expect("gemma4 global rope needs rope_freqs.weight"),
18652            )
18653        };
18654        #[cfg(debug_assertions)]
18655        if let Some(ff) = ff {
18656            crate::debug_assert_tensor_stream_device(
18657                ff,
18658                &e.stream(),
18659                "gemma4_decode_attn.rope_freqs",
18660            );
18661        }
18662        let kvl = cache.kv[il].as_mut().unwrap();
18663        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
18664        if crate::Engine::qkv_append_on() {
18665            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
18666            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
18667            // twin of the dc fold — bit-identical bodies, one launch per layer.
18668            e.rms_norm_qkv_rope_append(
18669                &q0,
18670                &k0,
18671                &v0,
18672                fa.q_norm.float_data(),
18673                fa.k_norm.float_data(),
18674                ones,
18675                &mut q,
18676                &mut k,
18677                &mut v,
18678                hd,
18679                self.gemma4_rope_dims(il),
18680                nh,
18681                nkv,
18682                pos_d,
18683                nh,
18684                nkv,
18685                base,
18686                1.0,
18687                ff,
18688                eps,
18689                &mut kvl.k,
18690                &mut kvl.v,
18691                kvl.len,
18692                kvl.k_tok_bytes,
18693                kvl.v_tok_bytes,
18694                kv_fp8,
18695            )?;
18696        } else {
18697            e.rms_norm_qkv_rope(
18698                &q0,
18699                &k0,
18700                &v0,
18701                fa.q_norm.float_data(),
18702                fa.k_norm.float_data(),
18703                ones,
18704                &mut q,
18705                &mut k,
18706                &mut v,
18707                hd,
18708                self.gemma4_rope_dims(il),
18709                nh,
18710                nkv,
18711                pos_d,
18712                nh,
18713                nkv,
18714                base,
18715                1.0,
18716                ff,
18717                eps,
18718            )?;
18719            e.append_kv_quantized(
18720                &k,
18721                &v,
18722                &mut kvl.k,
18723                &mut kvl.v,
18724                kvl.len,
18725                kvl.kv_dim_k,
18726                kvl.kv_dim_v,
18727                kvl.k_tok_bytes,
18728                kvl.v_tok_bytes,
18729                kv_fp8,
18730            )?;
18731        }
18732        kvl.len += 1;
18733        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
18734        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
18735        // positional). Globals attend the full history.
18736        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
18737        let mut attn = e.uninit(nh * hd)?;
18738        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
18739        if !swa
18740            && hd == 512
18741            && kvl.len >= crate::fa512_min_tkv()
18742            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18743        {
18744            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18745            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18746            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
18747            let base = kvl.len as i32;
18748            e.i32_set_k(&mut kvl.len_d, base)?;
18749            e.fa_decode_rows(
18750                &q,
18751                &kp,
18752                &vp,
18753                &mut attn,
18754                hd,
18755                nh,
18756                nkv,
18757                kvl.len - 1,
18758                1,
18759                scale,
18760                kvl.k_tok_bytes,
18761                kvl.v_tok_bytes,
18762                Some((&kvl.len_d, -1)),
18763                false,
18764                false,
18765                None,
18766            )?;
18767            return e.matmul(&fa.wo, &attn, 1);
18768        }
18769        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
18770        if swa
18771            && kvl.len > win
18772            && hd == 256
18773            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18774        {
18775            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18776            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18777            let base = kvl.len as i32;
18778            e.i32_set_k(&mut kvl.len_d, base)?;
18779            e.fa_decode_rows_w(
18780                &q,
18781                &kp,
18782                &vp,
18783                &mut attn,
18784                hd,
18785                nh,
18786                nkv,
18787                &kvl.len_d,
18788                -1,
18789                1,
18790                scale,
18791                win,
18792                kvl.k_tok_bytes,
18793                kvl.v_tok_bytes,
18794                None,
18795            )?;
18796            return e.matmul(&fa.wo, &attn, 1);
18797        }
18798        let (off_tok, t_kv) = if swa && kvl.len > win {
18799            (kvl.len - win, win)
18800        } else {
18801            (0, kvl.len)
18802        };
18803        let k_view = e.view_u8_range(
18804            &kvl.k,
18805            off_tok * kvl.k_tok_bytes,
18806            (off_tok + t_kv) * kvl.k_tok_bytes,
18807        );
18808        let v_view = e.view_u8_range(
18809            &kvl.v,
18810            off_tok * kvl.v_tok_bytes,
18811            (off_tok + t_kv) * kvl.v_tok_bytes,
18812        );
18813        e.fa_decode_kvmod(
18814            &q,
18815            &k_view,
18816            &v_view,
18817            &mut attn,
18818            hd,
18819            nh,
18820            nkv,
18821            t_kv,
18822            scale,
18823            kvl.k_tok_bytes,
18824            kvl.v_tok_bytes,
18825            swa && crate::Engine::wkv_on(),
18826        )?;
18827        e.matmul(&fa.wo, &attn, 1)
18828    }
18829
18830    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
18831    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
18832    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
18833    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
18834    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
18835    /// in-graph; the driver gates).
18836    #[allow(clippy::too_many_arguments)]
18837    pub fn gemma4_decode_step_dc(
18838        &self,
18839        e: &Engine,
18840        token_d: &CudaSlice<u32>,
18841        pos_d: &mut CudaSlice<i32>,
18842        embd_gpu: &CudaSlice<u8>,
18843        embd_qt: i32,
18844        embd_rb: usize,
18845        cache: &mut Cache,
18846        n_vocab: usize,
18847        cap_bucket_max: Option<(usize, usize)>,
18848    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
18849        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
18850        self.gemma4_decode_step_dc_into(
18851            e,
18852            token_d,
18853            pos_d,
18854            embd_gpu,
18855            embd_qt,
18856            embd_rb,
18857            cache,
18858            n_vocab,
18859            cap_bucket_max,
18860            &mut tok_out,
18861        )?;
18862        Ok(tok_out)
18863    }
18864
18865    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
18866    /// every replay; pass `token_d` itself for the self-feeding graph loop).
18867    #[allow(clippy::too_many_arguments)]
18868    pub fn gemma4_decode_step_dc_into(
18869        &self,
18870        e: &Engine,
18871        token_d: &CudaSlice<u32>,
18872        pos_d: &mut CudaSlice<i32>,
18873        embd_gpu: &CudaSlice<u8>,
18874        embd_qt: i32,
18875        embd_rb: usize,
18876        cache: &mut Cache,
18877        n_vocab: usize,
18878        cap_bucket_max: Option<(usize, usize)>,
18879        tok_out: &mut CudaSlice<u32>,
18880    ) -> Result<(), Box<dyn std::error::Error>> {
18881        let n_embd = self.cfg.n_embd as usize;
18882        let eps = self.cfg.rms_eps;
18883        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
18884        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
18885        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
18886        let n_layers = self.layers.len();
18887        for (il, layer) in self.layers.iter().enumerate() {
18888            let (hq, hdq) = match h_carry.take() {
18889                Some(p) => p,
18890                None => {
18891                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
18892                }
18893            };
18894            let Mixer::Full(fa) = &layer.mixer else {
18895                panic!("gemma4 layer {il} not full-attn")
18896            };
18897            let o =
18898                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
18899            let next_norm = if il + 1 < n_layers {
18900                Some(self.layers[il + 1].attn_norm.float_data())
18901            } else {
18902                None
18903            };
18904            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
18905            x = xn;
18906            h_carry = hn;
18907        }
18908        let mut hn = e.uninit(n_embd)?;
18909        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
18910        let mut logits = e.matmul(&self.output, &hn, 1)?;
18911        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
18912        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
18913        e.inc_seqlen(pos_d)?;
18914        if cap_bucket_max.is_none() {
18915            cache.pos += 1;
18916        }
18917        Ok(())
18918    }
18919
18920    // Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
18921    // every buffer the step produces per token lives here, allocated ONCE pre-capture, so
18922    // the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
18923    // osrt 2026-07-23). Sized for the model's max per-layer shapes.
18924
18925    /// Build the slot set (call OUTSIDE any capture).
18926    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
18927        let n_embd = self.cfg.n_embd as usize;
18928        let n_vocab = self.output.out_features();
18929        let n_layers = self.layers.len();
18930        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
18931        for il in 0..n_layers {
18932            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
18933            qmax = qmax.max(nh * hd);
18934            kvmax = kvmax.max(nkv * hd);
18935            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
18936                ffmax = ffmax.max(ffn_gate.out_features());
18937            }
18938        }
18939        Ok(G4DcSlots {
18940            x: e.uninit(n_embd)?,
18941            xn: e.uninit(n_embd)?,
18942            cur: e.uninit(n_embd)?,
18943            hq: e.alloc_i8_uninit(n_embd)?,
18944            hd_: e.uninit(n_embd / 32)?,
18945            q0: e.uninit(qmax)?,
18946            k0: e.uninit(kvmax)?,
18947            v0: e.uninit(kvmax)?,
18948            q: e.uninit(qmax)?,
18949            k: e.uninit(kvmax)?,
18950            v: e.uninit(kvmax)?,
18951            attn: e.uninit(qmax)?,
18952            o: e.uninit(n_embd)?,
18953            attn_out: e.uninit(n_embd)?,
18954            zsh: e.uninit(n_embd)?,
18955            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
18956            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
18957            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
18958            zd: e.uninit(n_embd.max(qmax) / 32)?,
18959            gate: e.uninit(ffmax)?,
18960            up: e.uninit(ffmax)?,
18961            act: e.uninit(ffmax)?,
18962            actq: e.alloc_i8_uninit(ffmax)?,
18963            actd: e.uninit(ffmax / 32)?,
18964            f0: e.uninit(n_embd)?,
18965            sn: e.uninit(n_embd)?,
18966            hn: e.uninit(n_embd)?,
18967            logits: e.uninit(n_vocab)?,
18968        })
18969    }
18970
18971    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
18972    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
18973    fn g4_matvec_m1_into(
18974        &self,
18975        e: &Engine,
18976        w: &crate::model::GpuTensor,
18977        aq: &CudaSlice<i8>,
18978        ad: &CudaSlice<f32>,
18979        y: &mut CudaSlice<f32>,
18980    ) -> Result<(), Box<dyn std::error::Error>> {
18981        use crate::model::GpuTensor;
18982        let (bytes, qtype, row_bytes, scale, rp) = match w {
18983            GpuTensor::Quant {
18984                bytes,
18985                qtype,
18986                row_bytes,
18987                scale,
18988                rp,
18989                ..
18990            } => (bytes, *qtype, *row_bytes, *scale, *rp),
18991            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
18992        };
18993        let (mbytes, mrp) = match w {
18994            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
18995            _ => (bytes, rp),
18996        };
18997        e.qmatvec_mmvq_into(
18998            mbytes,
18999            aq,
19000            ad,
19001            1,
19002            w.in_features(),
19003            w.out_features(),
19004            qtype,
19005            row_bytes,
19006            scale,
19007            mrp,
19008            y,
19009        )
19010    }
19011
19012    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
19013    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
19014    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
19015    #[allow(clippy::too_many_arguments)]
19016    pub fn gemma4_decode_step_dc_slotted(
19017        &self,
19018        e: &Engine,
19019        token_d: &CudaSlice<u32>,
19020        pos_d: &mut CudaSlice<i32>,
19021        embd_gpu: &CudaSlice<u8>,
19022        embd_qt: i32,
19023        embd_rb: usize,
19024        cache: &mut Cache,
19025        n_vocab: usize,
19026        cap_bucket_max: Option<(usize, usize)>,
19027        sl: &mut G4DcSlots,
19028        tok_out: &mut CudaSlice<u32>,
19029        ring: Option<(&mut CudaSlice<u32>, usize)>,
19030    ) -> Result<(), Box<dyn std::error::Error>> {
19031        let n_embd = self.cfg.n_embd as usize;
19032        let eps = self.cfg.rms_eps;
19033        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
19034        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
19035        let n_layers = self.layers.len();
19036        let mut has_carry = false;
19037        for il in 0..n_layers {
19038            if !has_carry {
19039                e.rms_norm_q8_1_into(
19040                    &sl.x,
19041                    self.layers[il].attn_norm.float_data(),
19042                    n_embd,
19043                    1,
19044                    eps,
19045                    &mut sl.hq,
19046                    &mut sl.hd_,
19047                )?;
19048            }
19049            has_carry = true;
19050            let layer = &self.layers[il];
19051            let Mixer::Full(fa) = &layer.mixer else {
19052                panic!("gemma4 layer {il} not full-attn")
19053            };
19054            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
19055            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
19056            // the standalone norm only survives on the unfused seam arm.
19057            if !Engine::g4_pnfold_on() {
19058                e.rms_norm(
19059                    &sl.o,
19060                    layer.post_attn_norm.float_data(),
19061                    &mut sl.cur,
19062                    n_embd,
19063                    1,
19064                    eps,
19065                )?;
19066            }
19067            let next_norm = if il + 1 < n_layers {
19068                Some(self.layers[il + 1].attn_norm.float_data())
19069            } else {
19070                None
19071            };
19072            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
19073            std::mem::swap(&mut sl.x, &mut sl.xn);
19074        }
19075        e.rms_norm(
19076            &sl.x,
19077            self.output_norm.float_data(),
19078            &mut sl.hn,
19079            n_embd,
19080            1,
19081            eps,
19082        )?;
19083        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
19084        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
19085        {
19086            let (zq, zd) = (&sl.zq, &sl.zd);
19087            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
19088            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
19089            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
19090        }
19091        self.gemma4_suppress(e, &mut sl.logits, 1)?;
19092        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
19093        if let Some((ring, base)) = ring {
19094            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
19095            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
19096            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
19097            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
19098        }
19099        e.inc_seqlen(pos_d)?;
19100        if cap_bucket_max.is_none() {
19101            cache.pos += 1;
19102        }
19103        Ok(())
19104    }
19105
19106    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
19107    #[allow(clippy::too_many_arguments)]
19108    fn gemma4_decode_attn_dc_slotted(
19109        &self,
19110        e: &Engine,
19111        fa: &crate::hybrid::FullAttnLayer,
19112        il: usize,
19113        pos_d: &CudaSlice<i32>,
19114        cache: &mut Cache,
19115        cap_bucket_max: Option<(usize, usize)>,
19116        sl: &mut G4DcSlots,
19117    ) -> Result<(), Box<dyn std::error::Error>> {
19118        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
19119        let eps = self.cfg.rms_eps;
19120        let aux = self.gemma4_aux.as_ref().unwrap();
19121        let ones = aux.ones(e);
19122        #[cfg(debug_assertions)]
19123        crate::debug_assert_tensor_stream_device(
19124            ones,
19125            &e.stream(),
19126            "gemma4_decode_attn_dc_slotted.ones",
19127        );
19128        {
19129            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
19130            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
19131            if swa {
19132                if !e.matmul_q4_fused3_into(
19133                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
19134                )? {
19135                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
19136                    // (q,k) pair, v through the generic m1 slot matvec — the same two
19137                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
19138                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
19139                    {
19140                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
19141                    } else {
19142                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
19143                    }
19144                }
19145            } else {
19146                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
19147                    && !e
19148                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
19149                {
19150                    return Err("slotted step: fused2 unavailable".into());
19151                }
19152                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
19153                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
19154            }
19155        }
19156        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
19157        // kernel-for-kernel (graph stream-identity gate).
19158        let ff = if swa {
19159            None
19160        } else {
19161            Some(
19162                aux.rope_freqs(e)
19163                    .expect("gemma4 global rope needs rope_freqs.weight"),
19164            )
19165        };
19166        #[cfg(debug_assertions)]
19167        if let Some(ff) = ff {
19168            crate::debug_assert_tensor_stream_device(
19169                ff,
19170                &e.stream(),
19171                "gemma4_decode_attn_dc_slotted.rope_freqs",
19172            );
19173        }
19174        let kvl = cache.kv[il].as_mut().unwrap();
19175        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19176        if crate::Engine::qkv_append_on() {
19177            // append fold (2026-07-23): mirrors dc_into.
19178            e.rms_norm_qkv_rope_append_dc(
19179                &sl.q0,
19180                &sl.k0,
19181                &sl.v0,
19182                fa.q_norm.float_data(),
19183                fa.k_norm.float_data(),
19184                ones,
19185                &mut sl.q,
19186                &mut sl.k,
19187                &mut sl.v,
19188                hd,
19189                self.gemma4_rope_dims(il),
19190                nh,
19191                nkv,
19192                pos_d,
19193                nh,
19194                nkv,
19195                base,
19196                1.0,
19197                ff,
19198                eps,
19199                &mut kvl.k,
19200                &mut kvl.v,
19201                &kvl.len_d,
19202                kvl.k_tok_bytes,
19203                kvl.v_tok_bytes,
19204                kv_fp8,
19205            )?;
19206        } else {
19207            e.rms_norm_qkv_rope(
19208                &sl.q0,
19209                &sl.k0,
19210                &sl.v0,
19211                fa.q_norm.float_data(),
19212                fa.k_norm.float_data(),
19213                ones,
19214                &mut sl.q,
19215                &mut sl.k,
19216                &mut sl.v,
19217                hd,
19218                self.gemma4_rope_dims(il),
19219                nh,
19220                nkv,
19221                pos_d,
19222                nh,
19223                nkv,
19224                base,
19225                1.0,
19226                ff,
19227                eps,
19228            )?;
19229            e.append_kv_quantized_dc(
19230                &sl.k,
19231                &sl.v,
19232                &mut kvl.k,
19233                &mut kvl.v,
19234                &kvl.len_d,
19235                kvl.kv_dim_k,
19236                kvl.kv_dim_v,
19237                kvl.k_tok_bytes,
19238                kvl.v_tok_bytes,
19239                kv_fp8,
19240            )?;
19241        }
19242        e.inc_seqlen(&mut kvl.len_d)?;
19243        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
19244        let k_view = e.view_u8(&kvl.k, kvl.k.len());
19245        let v_view = e.view_u8(&kvl.v, kvl.v.len());
19246        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19247        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19248        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
19249        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
19250        // the dc_into arm branch-for-branch (stream gate).
19251        let mut fa_q8 = false;
19252        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19253            e.fa_decode_rows(
19254                &sl.q,
19255                &k_view,
19256                &v_view,
19257                &mut sl.attn,
19258                hd,
19259                nh,
19260                nkv,
19261                b_glob - 1,
19262                1,
19263                scale,
19264                kvl.k_tok_bytes,
19265                kvl.v_tok_bytes,
19266                Some((&kvl.len_d, -1)),
19267                false,
19268                false,
19269                Some((&mut sl.zq, &mut sl.zd)),
19270            )?;
19271            fa_q8 = true;
19272        } else if swa && b_swa > win && hd == 256 && rows_on {
19273            e.fa_decode_rows_w(
19274                &sl.q,
19275                &k_view,
19276                &v_view,
19277                &mut sl.attn,
19278                hd,
19279                nh,
19280                nkv,
19281                &kvl.len_d,
19282                -1,
19283                1,
19284                scale,
19285                win,
19286                kvl.k_tok_bytes,
19287                kvl.v_tok_bytes,
19288                Some((&mut sl.zq, &mut sl.zd)),
19289            )?;
19290            fa_q8 = true;
19291        } else {
19292            let b = if swa { b_swa } else { b_glob };
19293            e.fa_decode_dc(
19294                &sl.q,
19295                &k_view,
19296                &v_view,
19297                &mut sl.attn,
19298                hd,
19299                nh,
19300                nkv,
19301                &kvl.len_d,
19302                b,
19303                scale,
19304                kvl.k_tok_bytes,
19305                kvl.v_tok_bytes,
19306                swa && crate::Engine::wkv_on(),
19307            )?;
19308        }
19309        if !fa_q8 {
19310            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
19311            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
19312        }
19313        {
19314            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19315            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19316            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
19317        }
19318        Ok(())
19319    }
19320
19321    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
19322    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
19323    fn gemma4_layer_tail_slotted(
19324        &self,
19325        e: &Engine,
19326        layer: &crate::hybrid::HybridLayer,
19327        next_norm: Option<&CudaSlice<f32>>,
19328        sl: &mut G4DcSlots,
19329    ) -> Result<(), Box<dyn std::error::Error>> {
19330        let n_embd = self.cfg.n_embd as usize;
19331        let eps = self.cfg.rms_eps;
19332        let bits = layer.gemma4.as_ref().unwrap();
19333        let crate::hybrid::Ffn::Dense {
19334            ffn_gate,
19335            ffn_up,
19336            ffn_down,
19337        } = &layer.ffn
19338        else {
19339            return Err("slotted tail: dense ffn only".into());
19340        };
19341        let pnfold = Engine::g4_pnfold_on();
19342        if pnfold {
19343            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
19344            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
19345            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
19346            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
19347            e.rms_pre_add_rms_norm_q8z_into(
19348                or,
19349                layer.post_attn_norm.float_data(),
19350                xr,
19351                bits.ffn_norm.float_data(),
19352                &mut sl.attn_out,
19353                &mut sl.zsh,
19354                n_embd,
19355                1,
19356                eps,
19357                &mut sl.zq,
19358                &mut sl.zd,
19359            )?;
19360        } else {
19361            e.add_rms_norm(
19362                &sl.cur,
19363                &sl.x,
19364                bits.ffn_norm.float_data(),
19365                &mut sl.attn_out,
19366                &mut sl.zsh,
19367                n_embd,
19368                1,
19369                eps,
19370            )?;
19371        }
19372        let n_ff = ffn_gate.out_features();
19373        if !pnfold {
19374            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
19375            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
19376        }
19377        {
19378            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19379            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19380            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
19381                && !e.matmul_nvfp4_fused2_into(
19382                    ffn_gate,
19383                    ffn_up,
19384                    zq,
19385                    zd,
19386                    &mut sl.gate,
19387                    &mut sl.up,
19388                )?
19389            {
19390                return Err("slotted tail: ffn fused2 unavailable".into());
19391            }
19392        }
19393        debug_assert!(e.uses_q8_1_fast(ffn_down));
19394        {
19395            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
19396            let upv = e.view(upr, n_ff);
19397            let up_all = upv.slice(0..n_ff);
19398            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
19399            e.gelu_tanh_mul_q8_1_into(
19400                gr,
19401                &up_all,
19402                &mut sl.act,
19403                n_ff,
19404                1,
19405                &mut sl.actq,
19406                &mut sl.actd,
19407            )?;
19408        }
19409        {
19410            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
19411            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
19412            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
19413        }
19414        if pnfold {
19415            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
19416            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
19417            if let Some(w) = next_norm {
19418                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
19419                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
19420                e.rms_pre_add_scale_rms_norm_q8_1_into(
19421                    f0r,
19422                    bits.post_ffw_norm.float_data(),
19423                    aor,
19424                    bits.layer_scale,
19425                    w,
19426                    &mut sl.xn,
19427                    n_embd,
19428                    1,
19429                    eps,
19430                    &mut sl.hq,
19431                    &mut sl.hd_,
19432                )?;
19433                return Ok(());
19434            }
19435        }
19436        e.rms_norm(
19437            &sl.f0,
19438            bits.post_ffw_norm.float_data(),
19439            &mut sl.sn,
19440            n_embd,
19441            1,
19442            eps,
19443        )?;
19444        match next_norm {
19445            Some(w) => {
19446                e.add_scale_rms_norm_q8_1_into(
19447                    &sl.sn,
19448                    &sl.attn_out,
19449                    bits.layer_scale,
19450                    w,
19451                    &mut sl.xn,
19452                    n_embd,
19453                    1,
19454                    eps,
19455                    &mut sl.hq,
19456                    &mut sl.hd_,
19457                )?;
19458            }
19459            None => {
19460                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
19461            }
19462        }
19463        Ok(())
19464    }
19465
19466    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
19467    #[allow(clippy::too_many_arguments)]
19468    fn gemma4_decode_attn_dc(
19469        &self,
19470        e: &Engine,
19471        fa: &crate::hybrid::FullAttnLayer,
19472        il: usize,
19473        hq: &CudaSlice<i8>,
19474        hdq: &CudaSlice<f32>,
19475        pos_d: &CudaSlice<i32>,
19476        cache: &mut Cache,
19477        cap_bucket_max: Option<(usize, usize)>,
19478    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19479        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
19480        let eps = self.cfg.rms_eps;
19481        let aux = self.gemma4_aux.as_ref().unwrap();
19482        let ones = aux.ones(e);
19483        #[cfg(debug_assertions)]
19484        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
19485        let (q0, k0, v0) = if swa {
19486            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
19487                Some(t3) => t3,
19488                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
19489                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19490                    Some((q0, k0)) => {
19491                        let h0 = e.zeros(0)?;
19492                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
19493                        (q0, k0, v0)
19494                    }
19495                    None => {
19496                        let h0 = e.zeros(0)?;
19497                        (
19498                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19499                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19500                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
19501                        )
19502                    }
19503                },
19504            }
19505        } else {
19506            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
19507                Some(p) => p,
19508                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19509                    Some(p) => p,
19510                    None => {
19511                        let h0 = e.zeros(0)?;
19512                        (
19513                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19514                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19515                        )
19516                    }
19517                },
19518            };
19519            let v0 = e.clone_dtod(&k0)?;
19520            (q0, k0, v0)
19521        };
19522        let mut q = e.uninit(nh * hd)?;
19523        let mut k = e.uninit(nkv * hd)?;
19524        let mut v = e.uninit(nkv * hd)?;
19525        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
19526        let ff = if swa {
19527            None
19528        } else {
19529            Some(
19530                aux.rope_freqs(e)
19531                    .expect("gemma4 global rope needs rope_freqs.weight"),
19532            )
19533        };
19534        #[cfg(debug_assertions)]
19535        if let Some(ff) = ff {
19536            crate::debug_assert_tensor_stream_device(
19537                ff,
19538                &e.stream(),
19539                "gemma4_decode_attn_dc.rope_freqs",
19540            );
19541        }
19542        let kvl = cache.kv[il].as_mut().unwrap();
19543        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19544        if crate::Engine::qkv_append_on() {
19545            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
19546            e.rms_norm_qkv_rope_append_dc(
19547                &q0,
19548                &k0,
19549                &v0,
19550                fa.q_norm.float_data(),
19551                fa.k_norm.float_data(),
19552                ones,
19553                &mut q,
19554                &mut k,
19555                &mut v,
19556                hd,
19557                self.gemma4_rope_dims(il),
19558                nh,
19559                nkv,
19560                pos_d,
19561                nh,
19562                nkv,
19563                base,
19564                1.0,
19565                ff,
19566                eps,
19567                &mut kvl.k,
19568                &mut kvl.v,
19569                &kvl.len_d,
19570                kvl.k_tok_bytes,
19571                kvl.v_tok_bytes,
19572                kv_fp8,
19573            )?;
19574        } else {
19575            e.rms_norm_qkv_rope(
19576                &q0,
19577                &k0,
19578                &v0,
19579                fa.q_norm.float_data(),
19580                fa.k_norm.float_data(),
19581                ones,
19582                &mut q,
19583                &mut k,
19584                &mut v,
19585                hd,
19586                self.gemma4_rope_dims(il),
19587                nh,
19588                nkv,
19589                pos_d,
19590                nh,
19591                nkv,
19592                base,
19593                1.0,
19594                ff,
19595                eps,
19596            )?;
19597            e.append_kv_quantized_dc(
19598                &k,
19599                &v,
19600                &mut kvl.k,
19601                &mut kvl.v,
19602                &kvl.len_d,
19603                kvl.kv_dim_k,
19604                kvl.kv_dim_v,
19605                kvl.k_tok_bytes,
19606                kvl.v_tok_bytes,
19607                kv_fp8,
19608            )?;
19609        }
19610        e.inc_seqlen(&mut kvl.len_d)?;
19611        let mut attn = e.uninit(nh * hd)?;
19612        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
19613        // rides g4_matvec_m1_into instead of matmul's internal quantize.
19614        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19615        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
19616        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
19617        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
19618        // (gemma4_e4b_attn, +0.65% valid window).
19619        match cap_bucket_max {
19620            None => {
19621                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
19622                // decode (SWA layers attend the last `sliding_window` keys); the device
19623                // counters carry only the append slot + the graph seam.
19624                kvl.len += 1;
19625                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19626                if !swa
19627                    && hd == 512
19628                    && kvl.len >= crate::fa512_min_tkv()
19629                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19630                {
19631                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
19632                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
19633                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19634                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19635                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19636                    e.fa_decode_rows(
19637                        &q,
19638                        &kp,
19639                        &vp,
19640                        &mut attn,
19641                        hd,
19642                        nh,
19643                        nkv,
19644                        kvl.len - 1,
19645                        1,
19646                        scale,
19647                        kvl.k_tok_bytes,
19648                        kvl.v_tok_bytes,
19649                        Some((&kvl.len_d, -1)),
19650                        false,
19651                        false,
19652                        Some((&mut aq8, &mut ad8)),
19653                    )?;
19654                    fa_q8 = Some((aq8, ad8));
19655                } else if swa
19656                    && kvl.len > win
19657                    && hd == 256
19658                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19659                {
19660                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
19661                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19662                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19663                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19664                    e.fa_decode_rows_w(
19665                        &q,
19666                        &kp,
19667                        &vp,
19668                        &mut attn,
19669                        hd,
19670                        nh,
19671                        nkv,
19672                        &kvl.len_d,
19673                        -1,
19674                        1,
19675                        scale,
19676                        win,
19677                        kvl.k_tok_bytes,
19678                        kvl.v_tok_bytes,
19679                        Some((&mut aq8, &mut ad8)),
19680                    )?;
19681                    fa_q8 = Some((aq8, ad8));
19682                } else {
19683                    let (off_tok, t_kv) = if swa && kvl.len > win {
19684                        (kvl.len - win, win)
19685                    } else {
19686                        (0, kvl.len)
19687                    };
19688                    let k_view = e.view_u8_range(
19689                        &kvl.k,
19690                        off_tok * kvl.k_tok_bytes,
19691                        (off_tok + t_kv) * kvl.k_tok_bytes,
19692                    );
19693                    let v_view = e.view_u8_range(
19694                        &kvl.v,
19695                        off_tok * kvl.v_tok_bytes,
19696                        (off_tok + t_kv) * kvl.v_tok_bytes,
19697                    );
19698                    e.fa_decode_kvmod(
19699                        &q,
19700                        &k_view,
19701                        &v_view,
19702                        &mut attn,
19703                        hd,
19704                        nh,
19705                        nkv,
19706                        t_kv,
19707                        scale,
19708                        kvl.k_tok_bytes,
19709                        kvl.v_tok_bytes,
19710                        swa && crate::Engine::wkv_on(),
19711                    )?;
19712                }
19713            }
19714            Some((b_swa, b_glob)) => {
19715                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
19716                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
19717                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
19718                // the RUNG max for the rows family (kernels derive per-replay splits from
19719                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
19720                let k_view = e.view_u8(&kvl.k, kvl.k.len());
19721                let v_view = e.view_u8(&kvl.v, kvl.v.len());
19722                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19723                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19724                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19725                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19726                    e.fa_decode_rows(
19727                        &q,
19728                        &k_view,
19729                        &v_view,
19730                        &mut attn,
19731                        hd,
19732                        nh,
19733                        nkv,
19734                        b_glob - 1,
19735                        1,
19736                        scale,
19737                        kvl.k_tok_bytes,
19738                        kvl.v_tok_bytes,
19739                        Some((&kvl.len_d, -1)),
19740                        false,
19741                        false,
19742                        Some((&mut aq8, &mut ad8)),
19743                    )?;
19744                    fa_q8 = Some((aq8, ad8));
19745                } else if swa && b_swa > win && hd == 256 && rows_on {
19746                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19747                    e.fa_decode_rows_w(
19748                        &q,
19749                        &k_view,
19750                        &v_view,
19751                        &mut attn,
19752                        hd,
19753                        nh,
19754                        nkv,
19755                        &kvl.len_d,
19756                        -1,
19757                        1,
19758                        scale,
19759                        win,
19760                        kvl.k_tok_bytes,
19761                        kvl.v_tok_bytes,
19762                        Some((&mut aq8, &mut ad8)),
19763                    )?;
19764                    fa_q8 = Some((aq8, ad8));
19765                } else {
19766                    let b = if swa { b_swa } else { b_glob };
19767                    e.fa_decode_dc(
19768                        &q,
19769                        &k_view,
19770                        &v_view,
19771                        &mut attn,
19772                        hd,
19773                        nh,
19774                        nkv,
19775                        &kvl.len_d,
19776                        b,
19777                        scale,
19778                        kvl.k_tok_bytes,
19779                        kvl.v_tok_bytes,
19780                        swa && crate::Engine::wkv_on(),
19781                    )?;
19782                }
19783            }
19784        }
19785        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
19786        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
19787        if let Some((aq8, ad8)) = fa_q8 {
19788            let mut y = e.uninit(fa.wo.out_features())?;
19789            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
19790            return Ok(y);
19791        }
19792        e.matmul(&fa.wo, &attn, 1)
19793    }
19794
19795    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
19796    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
19797    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
19798    /// views in-graph); caller gates and falls back to the dc-eager loop.
19799    #[allow(clippy::too_many_arguments)]
19800    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19801    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
19802    pub fn gemma4_generate_graph(
19803        &self,
19804        e: &Engine,
19805        prompt_pos: usize,
19806        first_token: u32,
19807        cache: &mut Cache,
19808        max_new: usize,
19809        eos: &[u32],
19810        mut on_token: impl FnMut(u32) -> bool,
19811    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
19812        if self.is_gemma4_e4b() {
19813            return Err(
19814                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
19815                    .into(),
19816            );
19817        }
19818        use crate::decode::StopReason;
19819        let n_vocab = self.output.out_features();
19820        let n_embd = self.cfg.n_embd as usize;
19821        let embd_gpu = self
19822            .embd_gpu
19823            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
19824        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
19825        for kvl in cache.kv.iter_mut().flatten() {
19826            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
19827        }
19828        let mut token_d = e.stream().clone_htod(&[first_token])?;
19829        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
19830        let g4 = self.cfg.gemma4.as_ref().unwrap();
19831        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
19832        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
19833        let nkv_s = g4
19834            .head_count_kv
19835            .iter()
19836            .zip(g4.swa_pattern.iter())
19837            .find(|p| *p.1)
19838            .map(|p| *p.0 as usize)
19839            .unwrap_or(8);
19840        let nkv_g = g4
19841            .head_count_kv
19842            .iter()
19843            .zip(g4.swa_pattern.iter())
19844            .find(|p| !*p.1)
19845            .map(|p| *p.0 as usize)
19846            .unwrap_or(2);
19847        #[allow(clippy::type_complexity)]
19848        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19849        let mut graphs: std::collections::HashMap<
19850            ((bool, usize), (bool, usize), bool, bool),
19851            (
19852                cudarc::driver::CudaGraph,
19853                Vec<Box<dyn std::any::Any + Send>>,
19854            ),
19855        > = Default::default();
19856        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
19857        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
19858        let mut slots = self.g4_dc_slots(e)?;
19859        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
19860        // baked at the door entry (the modulo keeps every capture valid indefinitely).
19861        const RING: usize = 64;
19862        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
19863        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
19864        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
19865        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
19866        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
19867        const DRAIN: usize = 1;
19868        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
19869        let ring_base = prompt_pos;
19870        let mut out = Vec::with_capacity(max_new);
19871        let mut reason = StopReason::MaxNew;
19872        let mut next = first_token;
19873        let mut captures = 0usize;
19874        for _ in 0..max_new {
19875            out.push(next);
19876            if eos.contains(&next) {
19877                reason = StopReason::Eos;
19878                break;
19879            }
19880            if !on_token(next) {
19881                reason = StopReason::Callback;
19882                break;
19883            }
19884            let t_kv = cache.pos + 1;
19885            // Bucket key per ARM (graph arc step 3):
19886            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
19887            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
19888            //    the component collapses to a single marker).
19889            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
19890            //    at/above it — the kernel derives splits from len_d per replay, so buckets
19891            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
19892            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19893            let f512 = crate::fa512_min_tkv();
19894            let key_s = if t_kv > win {
19895                (true, usize::MAX)
19896            } else {
19897                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
19898            };
19899            let (key_g, rung_end) = if t_kv >= f512 {
19900                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
19901                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
19902                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
19903                ((true, end), end)
19904            } else {
19905                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
19906            };
19907            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
19908            if !graphs.contains_key(&key) {
19909                let bucket_max = (t_kv, rung_end);
19910                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
19911                let snap = cache.snapshot(e)?;
19912                let pos_save = e.dtoh_i32_one(&pos_d)?;
19913                let len_save: Vec<Option<i32>> = cache
19914                    .kv
19915                    .iter()
19916                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
19917                    .collect();
19918                let tok_save = e.dtoh_u32_one(&token_d)?;
19919                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
19920                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
19921                // regression class, and this door's measured -8.8%. The keeper pins warmup
19922                // transients so the captured graph holds kernel nodes only.
19923                let graph = {
19924                    let tok_ref = &mut token_d;
19925                    let pos_ref = &mut pos_d;
19926                    let cache_ref = &mut *cache;
19927                    let slots_ref = &mut slots;
19928                    let ring_ref = &mut ring;
19929                    e.capture_graph_retained_flags(
19930                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
19931                        |e| {
19932                        // self-feeding: the argmax writes token_d itself.
19933                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
19934                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
19935                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
19936                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
19937                                                           cache_ref, n_vocab, Some(bucket_max),
19938                                                           sl, tok_ref, Some((rg, ring_base)))
19939                    })?
19940                };
19941                cache.rollback(e, &snap, 0)?;
19942                e.set_i32_one(&mut pos_d, pos_save)?;
19943                for (il, ls) in len_save.iter().enumerate() {
19944                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
19945                        e.set_i32_one(&mut kvl.len_d, *v)?;
19946                    }
19947                }
19948                e.set_u32_one(&mut token_d, tok_save)?;
19949                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
19950                    && let Ok(c) = crate::graph_update::node_census(&graph.0)
19951                {
19952                    eprintln!("[graph-census] {c:?}");
19953                }
19954                graphs.insert(key, graph);
19955                captures += 1;
19956            }
19957            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
19958            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
19959            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
19960            // the budget; capture warmups already emitted their tokens through the ring.
19961            let mut chunk = 1usize;
19962            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
19963                .ok()
19964                .and_then(|v| v.parse().ok())
19965                .unwrap_or(DRAIN);
19966            while chunk < drain_cap && out.len() + chunk < max_new {
19967                let t_next = cache.pos + 1 + chunk;
19968                let key_s2 = if t_next > win {
19969                    (true, usize::MAX)
19970                } else {
19971                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
19972                };
19973                let key_g2 = if t_next >= f512 {
19974                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
19975                } else {
19976                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
19977                };
19978                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
19979                    break;
19980                }
19981                chunk += 1;
19982            }
19983            let g = &graphs.get(&key).unwrap().0;
19984            for _ in 0..chunk {
19985                g.launch()?;
19986            }
19987            e.stream().synchronize()?;
19988            let ringh = e.dtoh_u32(&ring)?;
19989            for j in 0..chunk {
19990                let pos_j = cache.pos + j;
19991                let tok_j = ringh[(pos_j - ring_base) % RING];
19992                cache.pos += 0; // advanced below in one shot
19993                if j + 1 == chunk {
19994                    next = tok_j;
19995                } else {
19996                    out.push(tok_j);
19997                    if eos.contains(&tok_j) || !on_token(tok_j) {
19998                        reason = if eos.contains(&tok_j) {
19999                            StopReason::Eos
20000                        } else {
20001                            StopReason::Callback
20002                        };
20003                        // roll device/host state back to the stop point.
20004                        let keep = cache.pos + j + 1;
20005                        e.set_i32_one(&mut pos_d, keep as i32)?;
20006                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
20007                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
20008                            kvl.len = keep;
20009                        }
20010                        cache.pos = keep;
20011                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
20012                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
20013                        }
20014                        return Ok((out, reason));
20015                    }
20016                }
20017            }
20018            cache.pos += chunk;
20019            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
20020                kvl.len += chunk;
20021            }
20022        }
20023        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
20024            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
20025        }
20026        Ok((out, reason))
20027    }
20028
20029    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
20030    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
20031    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
20032    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
20033    /// logits (host) + advances cache.pos by t.
20034    pub(crate) fn gemma4_decode_step_t(
20035        &self,
20036        e: &Engine,
20037        tokens: &[u32],
20038        pos0: usize,
20039        cache: &mut Cache,
20040    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
20041        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
20042    }
20043
20044    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
20045    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
20046    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
20047    pub(crate) fn gemma4_decode_step_t_am(
20048        &self,
20049        e: &Engine,
20050        tokens: &[u32],
20051        pos0: usize,
20052        cache: &mut Cache,
20053    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20054        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
20055        let t = tokens.len();
20056        let n_vocab = self.output.out_features();
20057        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
20058        for i in 0..t {
20059            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
20060        }
20061        Ok((e.dtoh_u32(&toks)?, hn))
20062    }
20063
20064    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
20065    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
20066    pub(crate) fn gemma4_decode_step_t_am_dev(
20067        &self,
20068        e: &Engine,
20069        tok_d: &CudaSlice<u32>,
20070        t: usize,
20071        pos0: usize,
20072        cache: &mut Cache,
20073    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20074        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
20075        let n_vocab = self.output.out_features();
20076        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
20077        for i in 0..t {
20078            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
20079        }
20080        Ok((vam, hn))
20081    }
20082
20083    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
20084    /// llama's h_nextn convention).
20085    pub(crate) fn gemma4_decode_step_t_h(
20086        &self,
20087        e: &Engine,
20088        tokens: &[u32],
20089        pos0: usize,
20090        cache: &mut Cache,
20091    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20092        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
20093        let t = tokens.len();
20094        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20095        e.softcap(&mut ld, cap, t * self.output.out_features())?;
20096        Ok((e.dtoh(&ld)?, hn))
20097    }
20098
20099    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
20100    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
20101    pub(crate) fn verify_stream_scratch(
20102        &self,
20103        e: &Engine,
20104        cap: usize,
20105    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
20106        Ok(VerifyStreamScratch {
20107            pos_d: e.htod_i32(&vec![0i32; cap])?,
20108            row_ctrs: (0..cap)
20109                .map(|_| e.htod_i32(&[0]))
20110                .collect::<Result<_, _>>()?,
20111        })
20112    }
20113
20114    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
20115    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
20116    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
20117    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
20118    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
20119    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
20120    /// sync, exactly the turnaround the burst exists to remove.
20121    #[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
20122    pub(crate) fn gemma4_verify_t_am_stream(
20123        &self,
20124        e: &Engine,
20125        tok_d: &CudaSlice<u32>,
20126        t: usize,
20127        ctr: &CudaSlice<i32>,
20128        hint: usize,
20129        cache: &mut Cache,
20130        scr: &mut VerifyStreamScratch,
20131    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20132        let n_embd = self.cfg.n_embd as usize;
20133        let eps = self.cfg.rms_eps;
20134        assert!(t <= scr.row_ctrs.len() && t <= 64);
20135        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
20136        for i in 0..t {
20137            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
20138        }
20139        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
20140        let embd_gpu = self
20141            .embd_gpu
20142            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
20143        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
20144        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
20145        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20146        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20147        let n_layers = self.layers.len();
20148        for (il, layer) in self.layers.iter().enumerate() {
20149            let (hq, hdq) = match h_carry.take() {
20150                Some(p) => p,
20151                None => {
20152                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20153                }
20154            };
20155            let Mixer::Full(fa) = &layer.mixer else {
20156                panic!("gemma4 layer {il} not full-attn")
20157            };
20158            let o = self
20159                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
20160            let next_norm = if il + 1 < n_layers {
20161                Some(self.layers[il + 1].attn_norm.float_data())
20162            } else {
20163                None
20164            };
20165            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20166            x = xn;
20167            h_carry = hn;
20168            self.dflash_tap(e, cache, il, &x, t)?;
20169        }
20170        let mut hn = e.uninit(t * n_embd)?;
20171        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20172        let ld = e.matmul(&self.output, &hn, t)?;
20173        let n_vocab = self.output.out_features();
20174        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
20175        for i in 0..t {
20176            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
20177        }
20178        Ok((vam, hn))
20179    }
20180
20181    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
20182    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
20183    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
20184    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
20185    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
20186    /// kernel later if it shows in the profile).
20187    pub(crate) fn dflash_tap(
20188        &self,
20189        e: &Engine,
20190        cache: &mut Cache,
20191        il: usize,
20192        x: &CudaSlice<f32>,
20193        t: usize,
20194    ) -> Result<(), Box<dyn std::error::Error>> {
20195        let Some(taps) = cache.dflash_taps.as_mut() else {
20196            return Ok(());
20197        };
20198        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
20199            return Ok(());
20200        };
20201        let h = taps.hidden;
20202        let n_taps = taps.layer_ids.len();
20203        let base = taps.base;
20204        debug_assert!(
20205            base + t <= taps.t,
20206            "tap window {base}+{t} exceeds sink {}",
20207            taps.t
20208        );
20209        let xv = e.view(x, t * h);
20210        for r in 0..t {
20211            let row = xv.slice(r * h..(r + 1) * h);
20212            e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
20213        }
20214        Ok(())
20215    }
20216
20217    fn gemma4_verify_trunk(
20218        &self,
20219        e: &Engine,
20220        tokens: &[u32],
20221        pos0: usize,
20222        cache: &mut Cache,
20223        tok_dev: Option<&CudaSlice<u32>>,
20224    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20225        let n_embd = self.cfg.n_embd as usize;
20226        let eps = self.cfg.rms_eps;
20227        let t = tokens.len();
20228        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
20229        let pos_d = e.htod_i32(&pos)?;
20230        let mut x = match tok_dev {
20231            Some(td) => {
20232                let embd_gpu = self
20233                    .embd_gpu
20234                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
20235                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
20236                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
20237            }
20238            None => e.htod(&self.embd.gather(n_embd, tokens))?,
20239        };
20240        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20241        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20242        let n_layers = self.layers.len();
20243        for (il, layer) in self.layers.iter().enumerate() {
20244            let (hq, hdq) = match h_carry.take() {
20245                Some(p) => p,
20246                None => {
20247                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20248                }
20249            };
20250            let Mixer::Full(fa) = &layer.mixer else {
20251                panic!("gemma4 layer {il} not full-attn")
20252            };
20253            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
20254            let next_norm = if il + 1 < n_layers {
20255                Some(self.layers[il + 1].attn_norm.float_data())
20256            } else {
20257                None
20258            };
20259            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20260            x = xn;
20261            h_carry = hn;
20262            self.dflash_tap(e, cache, il, &x, t)?;
20263        }
20264        let mut hn = e.uninit(t * n_embd)?;
20265        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20266        let mut ld = e.matmul(&self.output, &hn, t)?;
20267        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
20268        cache.pos += t;
20269        Ok((ld, hn))
20270    }
20271
20272    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
20273    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
20274    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
20275    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
20276    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
20277    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
20278    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
20279    #[allow(clippy::too_many_arguments)]
20280    fn gemma4_verify_attn_stream(
20281        &self,
20282        e: &Engine,
20283        fa: &crate::hybrid::FullAttnLayer,
20284        il: usize,
20285        hq: &CudaSlice<i8>,
20286        hdq: &CudaSlice<f32>,
20287        pos_d: &CudaSlice<i32>,
20288        t: usize,
20289        cache: &mut Cache,
20290        hint: usize,
20291        row_ctrs: &[CudaSlice<i32>],
20292    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20293        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20294        let eps = self.cfg.rms_eps;
20295        let aux = self.gemma4_aux.as_ref().unwrap();
20296        let ones = aux.ones(e);
20297        #[cfg(debug_assertions)]
20298        crate::debug_assert_tensor_stream_device(
20299            ones,
20300            &e.stream(),
20301            "gemma4_verify_attn_stream.ones",
20302        );
20303        let h0 = e.zeros(0)?;
20304        let h = &h0;
20305        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
20306        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
20307        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20308        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20309        let fused_qkv = if f2b {
20310            if swa {
20311                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20312                    .map(|(a, b, c)| (a, b, Some(c)))
20313            } else {
20314                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20315                    .map(|(a, b)| (a, b, None))
20316            }
20317        } else {
20318            None
20319        };
20320        let (q0, k0, v0) = match fused_qkv {
20321            Some((a, b, cv)) => {
20322                let v = match cv {
20323                    Some(c) => c,
20324                    None => e.clone_dtod(&b)?,
20325                };
20326                (a, b, v)
20327            }
20328            None => {
20329                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20330                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20331                let v0 = if swa {
20332                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20333                } else {
20334                    e.clone_dtod(&k0)?
20335                };
20336                (q0, k0, v0)
20337            }
20338        };
20339        let mut q = e.uninit(t * nh * hd)?;
20340        let mut k = e.uninit(t * nkv * hd)?;
20341        let mut v = e.uninit(t * nkv * hd)?;
20342        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
20343        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
20344        let ff = if swa {
20345            None
20346        } else {
20347            Some(
20348                aux.rope_freqs(e)
20349                    .expect("gemma4 global rope needs rope_freqs.weight"),
20350            )
20351        };
20352        #[cfg(debug_assertions)]
20353        if let Some(ff) = ff {
20354            crate::debug_assert_tensor_stream_device(
20355                ff,
20356                &e.stream(),
20357                "gemma4_verify_attn_stream.rope_freqs",
20358            );
20359        }
20360        e.rms_norm_qkv_rope(
20361            &q0,
20362            &k0,
20363            &v0,
20364            fa.q_norm.float_data(),
20365            fa.k_norm.float_data(),
20366            ones,
20367            &mut q,
20368            &mut k,
20369            &mut v,
20370            hd,
20371            self.gemma4_rope_dims(il),
20372            nh * t,
20373            nkv * t,
20374            pos_d,
20375            nh,
20376            nkv,
20377            base,
20378            1.0,
20379            ff,
20380            eps,
20381        )?;
20382        let kvl = cache.kv[il].as_mut().unwrap();
20383        // append at the DEVICE slot; the counter advances by t on-device.
20384        e.append_kv_quantized_rows_dc(
20385            &k,
20386            &v,
20387            &mut kvl.k,
20388            &mut kvl.v,
20389            &kvl.len_d,
20390            t,
20391            kvl.kv_dim_k,
20392            kvl.kv_dim_v,
20393            kvl.k_tok_bytes,
20394            kvl.v_tok_bytes,
20395            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20396        )?;
20397        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
20398        // the sole len writer after this round's attention (base stays = old len, plus = 0).
20399        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20400        let mut attn = e.uninit(t * nh * hd)?;
20401        let k_view = e.view_u8(&kvl.k, kvl.k.len());
20402        let v_view = e.view_u8(&kvl.v, kvl.v.len());
20403        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
20404        // and a stable window regime — the same rung/regime keys as the draft graph).
20405        if swa && hint + 1 >= win {
20406            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
20407            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
20408            e.fa_decode_rows_w(
20409                &q,
20410                &k_view,
20411                &v_view,
20412                &mut attn,
20413                hd,
20414                nh,
20415                nkv,
20416                &kvl.len_d,
20417                0,
20418                t,
20419                scale,
20420                win,
20421                kvl.k_tok_bytes,
20422                kvl.v_tok_bytes,
20423                None,
20424            )?;
20425        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
20426            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
20427            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
20428            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
20429            // Burst entry gates the horizon onto one side of the crossover, so hint decides
20430            // for every row.
20431            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
20432            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
20433            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
20434            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
20435            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
20436            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
20437            // any bucket >= the live length is exact.
20438            let bucket = (hint + t + 2)
20439                .next_power_of_two()
20440                .min(crate::fa512_min_tkv().saturating_sub(1));
20441            let qv = e.view(&q, t * nh * hd);
20442            #[allow(clippy::needless_range_loop)]
20443            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
20444            for i in 0..t {
20445                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
20446                let mut q_one = e.uninit(nh * hd)?;
20447                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20448                let mut a_one = e.uninit(nh * hd)?;
20449                e.fa_decode_dc(
20450                    &q_one,
20451                    &k_view,
20452                    &v_view,
20453                    &mut a_one,
20454                    hd,
20455                    nh,
20456                    nkv,
20457                    &row_ctrs[i],
20458                    bucket,
20459                    scale,
20460                    kvl.k_tok_bytes,
20461                    kvl.v_tok_bytes,
20462                    false,
20463                )?;
20464                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20465            }
20466        } else if hd == 512 {
20467            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
20468            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
20469            e.fa_decode_rows(
20470                &q,
20471                &k_view,
20472                &v_view,
20473                &mut attn,
20474                hd,
20475                nh,
20476                nkv,
20477                hint,
20478                t,
20479                scale,
20480                kvl.k_tok_bytes,
20481                kvl.v_tok_bytes,
20482                Some((&kvl.len_d, 0)),
20483                false,
20484                false,
20485                None,
20486            )?;
20487        } else {
20488            // hd256 under-window: v4 device-len rows twin.
20489            e.fa_decode_rows_dc(
20490                &q,
20491                &k_view,
20492                &v_view,
20493                &mut attn,
20494                hd,
20495                nh,
20496                nkv,
20497                &kvl.len_d,
20498                hint + t,
20499                t,
20500                scale,
20501                kvl.k_tok_bytes,
20502                kvl.v_tok_bytes,
20503                0,
20504                swa && crate::Engine::wkv_on(),
20505            )?;
20506        }
20507        e.matmul(&fa.wo, &attn, t)
20508    }
20509
20510    #[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
20511    fn gemma4_verify_attn(
20512        &self,
20513        e: &Engine,
20514        fa: &crate::hybrid::FullAttnLayer,
20515        il: usize,
20516        hq: &CudaSlice<i8>,
20517        hdq: &CudaSlice<f32>,
20518        pos_d: &CudaSlice<i32>,
20519        t: usize,
20520        cache: &mut Cache,
20521    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20522        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20523        let eps = self.cfg.rms_eps;
20524        let aux = self.gemma4_aux.as_ref().unwrap();
20525        let ones = aux.ones(e);
20526        #[cfg(debug_assertions)]
20527        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
20528        let n_embd = self.cfg.n_embd as usize;
20529        let _ = n_embd;
20530
20531        let h0 = e.zeros(0)?;
20532        let h = &h0;
20533        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
20534        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
20535        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20536        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20537        let fused_qkv = if f2b {
20538            if swa {
20539                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20540                    .map(|(a, b, c)| (a, b, Some(c)))
20541            } else {
20542                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20543                    .map(|(a, b)| (a, b, None))
20544            }
20545        } else {
20546            None
20547        };
20548        let (q0, k0, v0) = match fused_qkv {
20549            Some((a, b, cv)) => {
20550                let v = match cv {
20551                    Some(c) => c,
20552                    None => e.clone_dtod(&b)?,
20553                };
20554                (a, b, v)
20555            }
20556            None => {
20557                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20558                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20559                let v0 = if swa {
20560                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20561                } else {
20562                    e.clone_dtod(&k0)?
20563                };
20564                (q0, k0, v0)
20565            }
20566        };
20567        let mut q = e.uninit(t * nh * hd)?;
20568        let mut k = e.uninit(t * nkv * hd)?;
20569        let mut v = e.uninit(t * nkv * hd)?;
20570        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
20571        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
20572        let ff = if swa {
20573            None
20574        } else {
20575            Some(
20576                aux.rope_freqs(e)
20577                    .expect("gemma4 global rope needs rope_freqs.weight"),
20578            )
20579        };
20580        #[cfg(debug_assertions)]
20581        if let Some(ff) = ff {
20582            crate::debug_assert_tensor_stream_device(
20583                ff,
20584                &e.stream(),
20585                "gemma4_verify_attn.rope_freqs",
20586            );
20587        }
20588        e.rms_norm_qkv_rope(
20589            &q0,
20590            &k0,
20591            &v0,
20592            fa.q_norm.float_data(),
20593            fa.k_norm.float_data(),
20594            ones,
20595            &mut q,
20596            &mut k,
20597            &mut v,
20598            hd,
20599            self.gemma4_rope_dims(il),
20600            nh * t,
20601            nkv * t,
20602            pos_d,
20603            nh,
20604            nkv,
20605            base,
20606            1.0,
20607            ff,
20608            eps,
20609        )?;
20610        let kvl = cache.kv[il].as_mut().unwrap();
20611        let base_len = kvl.len;
20612        e.append_kv_quantized_rows(
20613            &k,
20614            &v,
20615            &mut kvl.k,
20616            &mut kvl.v,
20617            base_len,
20618            t,
20619            kvl.kv_dim_k,
20620            kvl.kv_dim_v,
20621            kvl.k_tok_bytes,
20622            kvl.v_tok_bytes,
20623            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20624        )?;
20625        kvl.len += t;
20626        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20627        let mut attn = e.uninit(t * nh * hd)?;
20628        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
20629        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
20630        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
20631            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
20632            // decode rides the SAME symbol at t=1 (parity law).
20633            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
20634        if rows_ok && (!swa || base_len + t <= win) {
20635            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20636            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20637            if hd == 512 {
20638                // device-len twin: sync the counter to the verify base (async arg-store).
20639                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20640                e.fa_decode_rows(
20641                    &q,
20642                    &k_view,
20643                    &v_view,
20644                    &mut attn,
20645                    hd,
20646                    nh,
20647                    nkv,
20648                    base_len,
20649                    t,
20650                    scale,
20651                    kvl.k_tok_bytes,
20652                    kvl.v_tok_bytes,
20653                    Some((&kvl.len_d, 0)),
20654                    false,
20655                    swa && crate::Engine::wkv_on(),
20656                    None,
20657                )?;
20658            } else {
20659                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
20660                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
20661                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
20662                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20663                e.fa_decode_rows_dc(
20664                    &q,
20665                    &k_view,
20666                    &v_view,
20667                    &mut attn,
20668                    hd,
20669                    nh,
20670                    nkv,
20671                    &kvl.len_d,
20672                    base_len + t,
20673                    t,
20674                    scale,
20675                    kvl.k_tok_bytes,
20676                    kvl.v_tok_bytes,
20677                    0,
20678                    swa && crate::Engine::wkv_on(),
20679                )?;
20680            }
20681            return e.matmul(&fa.wo, &attn, t);
20682        }
20683        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
20684        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
20685        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
20686        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
20687        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
20688        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
20689        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
20690        if hd == 256
20691            && swa
20692            && base_len + 1 >= win
20693            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20694        {
20695            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20696            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20697            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20698            e.fa_decode_rows_w(
20699                &q,
20700                &k_view,
20701                &v_view,
20702                &mut attn,
20703                hd,
20704                nh,
20705                nkv,
20706                &kvl.len_d,
20707                0,
20708                t,
20709                scale,
20710                win,
20711                kvl.k_tok_bytes,
20712                kvl.v_tok_bytes,
20713                None,
20714            )?;
20715            return e.matmul(&fa.wo, &attn, t);
20716        }
20717        for i in 0..t {
20718            let avail = base_len + i + 1;
20719            let (off_tok, t_kv) = if swa && avail > win {
20720                (avail - win, win)
20721            } else {
20722                (0, avail)
20723            };
20724            let k_view = e.view_u8_range(
20725                &kvl.k,
20726                off_tok * kvl.k_tok_bytes,
20727                (off_tok + t_kv) * kvl.k_tok_bytes,
20728            );
20729            let v_view = e.view_u8_range(
20730                &kvl.v,
20731                off_tok * kvl.v_tok_bytes,
20732                (off_tok + t_kv) * kvl.v_tok_bytes,
20733            );
20734            let qi = e.view(&q, t * nh * hd);
20735            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
20736            let mut q_one = e.uninit(nh * hd)?;
20737            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20738            let mut a_one = e.uninit(nh * hd)?;
20739            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
20740            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
20741            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
20742            if swa
20743                && avail > win
20744                && hd == 256
20745                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20746            {
20747                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20748                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20749                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20750                e.fa_decode_rows_w(
20751                    &q_one,
20752                    &kp,
20753                    &vp,
20754                    &mut a_one,
20755                    hd,
20756                    nh,
20757                    nkv,
20758                    &kvl.len_d,
20759                    0,
20760                    1,
20761                    scale,
20762                    win,
20763                    kvl.k_tok_bytes,
20764                    kvl.v_tok_bytes,
20765                    None,
20766                )?;
20767            } else if !swa
20768                && hd == 512
20769                && avail >= crate::fa512_min_tkv()
20770                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20771            {
20772                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20773                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20774                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20775                e.fa_decode_rows(
20776                    &q_one,
20777                    &kp,
20778                    &vp,
20779                    &mut a_one,
20780                    hd,
20781                    nh,
20782                    nkv,
20783                    avail - 1,
20784                    1,
20785                    scale,
20786                    kvl.k_tok_bytes,
20787                    kvl.v_tok_bytes,
20788                    Some((&kvl.len_d, 0)),
20789                    false,
20790                    false,
20791                    None,
20792                )?;
20793            } else {
20794                e.fa_decode_kvmod(
20795                    &q_one,
20796                    &k_view,
20797                    &v_view,
20798                    &mut a_one,
20799                    hd,
20800                    nh,
20801                    nkv,
20802                    t_kv,
20803                    scale,
20804                    kvl.k_tok_bytes,
20805                    kvl.v_tok_bytes,
20806                    swa && crate::Engine::wkv_on(),
20807                )?;
20808            }
20809            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20810        }
20811        e.matmul(&fa.wo, &attn, t)
20812    }
20813
20814    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
20815    /// h_seed = pre-output_norm hidden). Advances cache.pos.
20816    pub(crate) fn gemma4_decode_step_h(
20817        &self,
20818        e: &Engine,
20819        token: u32,
20820        cache: &mut Cache,
20821    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20822        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
20823        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
20824        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
20825        // unsplit rather than guessing a fence.
20826        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
20827            let rt = crate::pp::Pp2Rt::get(e)?;
20828            let _walk = rt.acquire_walk("gemma4_decode_step_h_pp2")?;
20829            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
20830        }
20831        if crate::pp::pp_cuts(self.layers.len()).is_some() {
20832            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
20833        }
20834        let n_embd = self.cfg.n_embd as usize;
20835        let eps = self.cfg.rms_eps;
20836        let pos_d = e.htod_i32(&[cache.pos as i32])?;
20837        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20838        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20839        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
20840        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
20841        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20842        let n_layers = self.layers.len();
20843        for (il, layer) in self.layers.iter().enumerate() {
20844            let (hq, hdq) = match h_carry.take() {
20845                Some(p) => p,
20846                None => {
20847                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
20848                }
20849            };
20850            let Mixer::Full(fa) = &layer.mixer else {
20851                panic!("gemma4 layer {il} not full-attn")
20852            };
20853            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
20854            let next_norm = if il + 1 < n_layers {
20855                Some(self.layers[il + 1].attn_norm.float_data())
20856            } else {
20857                None
20858            };
20859            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20860            x = xn;
20861            h_carry = hn;
20862        }
20863        let mut hn = e.uninit(n_embd)?;
20864        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20865        let h_seed = e.clone_dtod(&x)?;
20866        let mut ld = e.matmul(&self.output, &hn, 1)?;
20867        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20868        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
20869        self.gemma4_suppress(e, &mut ld, 1)?;
20870        let logits = e.dtoh(&ld)?;
20871        cache.pos += 1;
20872        Ok((logits, h_seed))
20873    }
20874
20875    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
20876    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
20877    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
20878    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
20879    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
20880    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
20881    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
20882    fn gemma4_decode_layers(
20883        &self,
20884        e: &Engine,
20885        mut x: CudaSlice<f32>,
20886        lo: usize,
20887        hi: usize,
20888        pos_d: &CudaSlice<i32>,
20889        cache: &mut Cache,
20890    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20891        let n_embd = self.cfg.n_embd as usize;
20892        let eps = self.cfg.rms_eps;
20893        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20894        for il in lo..hi {
20895            let layer = &self.layers[il];
20896            let (hq, hdq) = match h_carry.take() {
20897                Some(p) => p,
20898                // range head: il == lo — norm against THIS layer's attn_norm.
20899                None => {
20900                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
20901                }
20902            };
20903            let Mixer::Full(fa) = &layer.mixer else {
20904                panic!("gemma4 layer {il} not full-attn")
20905            };
20906            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
20907            let next_norm = if il + 1 < hi {
20908                Some(self.layers[il + 1].attn_norm.float_data())
20909            } else {
20910                None
20911            };
20912            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20913            x = xn;
20914            h_carry = hn;
20915        }
20916        Ok(x)
20917    }
20918
20919    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
20920    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
20921    /// boundary handoff — same choreography as the generic arm (decode.rs), same
20922    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
20923    /// stage 1 = layers [split, n) + output_norm + softcapped head.
20924    /// Each stage uploads its own copy of the step's position scalar on its own stream.
20925    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
20926    fn gemma4_decode_step_h_pp2(
20927        &self,
20928        e: &Engine,
20929        token: u32,
20930        cache: &mut Cache,
20931        split: usize,
20932    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20933        if crate::pp::pp2_streams_off() {
20934            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
20935        }
20936        let rt = crate::pp::Pp2Rt::get(e)?;
20937        let e0 = rt.engine(0, e);
20938        let e1 = rt.engine(1, e);
20939        let n_embd = self.cfg.n_embd as usize;
20940        let eps = self.cfg.rms_eps;
20941        let pos = cache.pos as i32;
20942
20943        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
20944        let slot = {
20945            let _st0 = rt.enter(0);
20946            let pos_d = e0.htod_i32(&[pos])?;
20947            #[cfg(debug_assertions)]
20948            crate::debug_assert_tensor_stream_device(
20949                &pos_d,
20950                &e0.stream(),
20951                "gemma4_decode_step_h_pp2.stage0.pos_d",
20952            );
20953            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
20954            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20955            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
20956            rt.tx(0, &x, n_embd)?
20957        };
20958
20959        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
20960        let _st1 = rt.enter(1);
20961        let pos_d = e1.htod_i32(&[pos])?;
20962        #[cfg(debug_assertions)]
20963        crate::debug_assert_tensor_stream_device(
20964            &pos_d,
20965            &e1.stream(),
20966            "gemma4_decode_step_h_pp2.stage1.pos_d",
20967        );
20968        let x = rt.rx(0, slot, n_embd)?;
20969        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
20970
20971        let mut hn = e1.uninit(n_embd)?;
20972        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20973        let h_seed = e1.clone_dtod(&x)?;
20974        let mut ld = e1.matmul(&self.output, &hn, 1)?;
20975        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20976        e1.softcap(&mut ld, cap, self.output.out_features())?;
20977        self.gemma4_suppress(e1, &mut ld, 1)?;
20978        let logits = e1.dtoh(&ld)?;
20979        cache.pos += 1;
20980        Ok((logits, h_seed))
20981    }
20982
20983    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
20984    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
20985    fn gemma4_decode_step_h_pp2_samestream(
20986        &self,
20987        e: &Engine,
20988        token: u32,
20989        cache: &mut Cache,
20990        split: usize,
20991    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20992        let n_embd = self.cfg.n_embd as usize;
20993        let eps = self.cfg.rms_eps;
20994        let pos_d = e.htod_i32(&[cache.pos as i32])?;
20995
20996        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
20997        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20998        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20999        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
21000
21001        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
21002        let boundary_tx = e.clone_dtod(&x)?;
21003        let boundary_rx = e.clone_dtod(&boundary_tx)?;
21004
21005        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
21006        let x =
21007            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
21008
21009        let mut hn = e.uninit(n_embd)?;
21010        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
21011        let h_seed = e.clone_dtod(&x)?;
21012        let mut ld = e.matmul(&self.output, &hn, 1)?;
21013        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
21014        e.softcap(&mut ld, cap, self.output.out_features())?;
21015        self.gemma4_suppress(e, &mut ld, 1)?;
21016        let logits = e.dtoh(&ld)?;
21017        cache.pos += 1;
21018        Ok((logits, h_seed))
21019    }
21020}
21021
21022// ============================ step35 (Step-3.7-Flash) ==================================
21023// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
21024// FAMILY and not a few branches inside the generic `full_attn*` chain:
21025//
21026//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
21027//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
21028//      shapes and the FA head counts would be wrong on 33 of 45 layers.
21029//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
21030//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
21031//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
21032//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
21033//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
21034//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
21035//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
21036//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
21037//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
21038//
21039// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
21040impl HybridModel {
21041    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
21042    /// synthesize a drafter or trunk layer from a neighboring class.
21043    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
21044        let geometry = self
21045            .cfg
21046            .layer_geometry(il as u32)
21047            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
21048        debug_assert_eq!(
21049            geometry.attention_gate,
21050            memra_gguf::config::AttentionGateKind::SeparateHead
21051        );
21052        geometry
21053    }
21054
21055    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
21056    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
21057    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
21058    ///
21059    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
21060    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
21061    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
21062    /// `cache`:
21063    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
21064    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
21065    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
21066    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
21067    ///     contract, lane/chunkinv-flip).
21068    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
21069    ///     q/k/v, no cache side effect.
21070    ///
21071    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
21072    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
21073    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
21074    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
21075    /// still contains must be masked per query. memra's window convention
21076    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
21077    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
21078    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
21079    ///
21080    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
21081    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
21082    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
21083    ///
21084    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
21085    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
21086    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
21087    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
21088    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
21089    /// hidden rows, and the generated text — a function of the chunk size:
21090    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
21091    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
21092    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
21093    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
21094    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
21095    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
21096    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
21097    ///   one-token change in a documented machine-config knob changed the answer.
21098    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
21099    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
21100    /// the same rows moves the logits by ~1.8.
21101    ///
21102    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
21103    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
21104    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
21105    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
21106    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
21107    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
21108    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
21109    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
21110    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
21111    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
21112    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
21113    /// those with t_kv <= win = 512.
21114    #[allow(clippy::too_many_arguments)]
21115    fn step35_attn_pre_wo(
21116        &self,
21117        e: &Engine,
21118        fa: &FullAttnLayer,
21119        mut g3: Vec<CudaSlice<f32>>,
21120        hg: Option<&CudaSlice<f32>>,
21121        gt_pre: Option<&CudaSlice<f32>>,
21122        pos_d: &CudaSlice<i32>,
21123        t: usize,
21124        cache: Option<&mut Cache>,
21125        il: usize,
21126        seq_end: usize,
21127    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21128        let geometry = self.cfg.full_attention_geometry_at(il as u32);
21129        let hd = geometry.head_dim_k as usize;
21130        let nkv = geometry.n_head_kv as usize;
21131        let nh = geometry.n_head as usize;
21132        let rbase = geometry.rope_base;
21133        let scale = geometry.attention_scale();
21134        let swa = geometry.window.is_some();
21135        let eps = self.cfg.rms_eps;
21136        let win = geometry.window.unwrap_or(0) as usize;
21137        let n_rot = geometry.n_rot as usize;
21138
21139        let v = g3.pop().unwrap();
21140        let k0 = g3.pop().unwrap();
21141        let q0 = g3.pop().unwrap();
21142
21143        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
21144        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
21145        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
21146        let mut q = e.uninit(t * nh * hd)?;
21147        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
21148        let mut k = e.uninit(t * nkv * hd)?;
21149        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
21150        let ff = if geometry.rope_factors {
21151            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
21152        } else {
21153            None
21154        };
21155        #[cfg(debug_assertions)]
21156        if let Some(ff) = ff {
21157            crate::debug_assert_tensor_stream_device(
21158                ff,
21159                &e.stream(),
21160                "step35_attn_pre_wo.rope_freqs",
21161            );
21162        }
21163        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
21164
21165        let mut attn = e.uninit(t * nh * hd)?;
21166        match cache {
21167            Some(cache) => {
21168                let base_len = cache.kv[il].as_ref().unwrap().len;
21169                // Read per layer call, never in a measured default.
21170                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
21171                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
21172                let off = if swa {
21173                    let raw = base_len.saturating_sub(win - 1);
21174                    if legacy_tkv || legacy_calllocal {
21175                        raw
21176                    } else {
21177                        raw & !31usize
21178                    }
21179                } else {
21180                    0
21181                };
21182                {
21183                    let kvl = cache.kv[il].as_mut().unwrap();
21184                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
21185                    let write_row = e.prepare_kv_append(kvl, off, t)?;
21186                    e.append_kv_quantized_rows(
21187                        &k,
21188                        &v,
21189                        &mut kvl.k,
21190                        &mut kvl.v,
21191                        write_row,
21192                        t,
21193                        kvl.kv_dim_k,
21194                        kvl.kv_dim_v,
21195                        kvl.k_tok_bytes,
21196                        kvl.v_tok_bytes,
21197                        crate::Engine::kv_fp8_on(),
21198                    )?;
21199                    kvl.len += t;
21200                    let new_len = kvl.len as i32;
21201                    e.set_i32_one(&mut kvl.len_d, new_len)?;
21202                }
21203                let kvl = cache.kv[il].as_ref().unwrap();
21204                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
21205                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
21206                // unaligned view offset here. Both halves are load-bearing for the canaries:
21207                // on the FA default the predicate arms agree bitwise wherever they can differ
21208                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
21209                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
21210                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
21211                // on the current FA path: its tile grid starts at the chunk/call boundary.
21212                // SWA: trim the view to the oldest key any query in this chunk can reach —
21213                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
21214                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
21215                // kernel's online-softmax recurrence groups keys into BK tiles relative to
21216                // the VIEW START — so an unaligned off regroups the same absolute keys into
21217                // different tiles at different chunk sizes = different (m,l) rounding =
21218                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
21219                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
21220                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
21221                // size; the <=31 extra leading keys are older than EVERY query's window
21222                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
21223                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
21224                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
21225                // the floor arm's bits do not move either (gated: G2f, battery 2).
21226                let t_kv = base_len + t - off;
21227                let physical = kvl.physical_rows(off, off + t_kv)?;
21228                let k_view = e.view_u8_range(
21229                    &kvl.k,
21230                    physical.start * kvl.k_tok_bytes,
21231                    physical.end * kvl.k_tok_bytes,
21232                );
21233                let v_view = e.view_u8_range(
21234                    &kvl.v,
21235                    physical.start * kvl.v_tok_bytes,
21236                    physical.end * kvl.v_tok_bytes,
21237                );
21238                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
21239                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
21240                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
21241                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
21242                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
21243                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
21244                // construction, so the invariance assertion MUST break under it (the seam whose
21245                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
21246                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
21247                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
21248                // cached (probes flip it in-process). Never on in a measured default run.
21249                let swa_naive = if legacy_tkv {
21250                    t_kv > win
21251                } else {
21252                    seq_end > win
21253                };
21254                if swa && swa_naive {
21255                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
21256                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
21257                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
21258                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
21259                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
21260                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
21261                    // identically to the unwindowed one modulo the mask, which is the point.
21262                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
21263                    // selected on `seq_end` like every arm here, so the class is uniform for
21264                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
21265                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
21266                    // the f32 floor (the previous numeric config, kept as the A/B seam).
21267                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21268                        e.sdpa_naive_w_quantized_view(
21269                            &q,
21270                            &k_view,
21271                            &v_view,
21272                            &mut attn,
21273                            hd,
21274                            nh,
21275                            nkv,
21276                            t,
21277                            t_kv,
21278                            scale,
21279                            true,
21280                            win,
21281                            kvl.k_tok_bytes,
21282                            kvl.v_tok_bytes,
21283                        )?;
21284                    } else {
21285                        e.fa_prefill_view_ws_w_hd128(
21286                            &q,
21287                            &k_view,
21288                            &v_view,
21289                            &mut attn,
21290                            hd,
21291                            nh,
21292                            nkv,
21293                            t,
21294                            t_kv,
21295                            scale,
21296                            true,
21297                            win,
21298                            kvl.k_tok_bytes,
21299                            kvl.v_tok_bytes,
21300                        )?;
21301                    }
21302                } else if std::env::var("MEMRA_NOFA").is_ok() {
21303                    e.sdpa_naive_quantized_view(
21304                        &q,
21305                        &k_view,
21306                        &v_view,
21307                        &mut attn,
21308                        hd,
21309                        nh,
21310                        nkv,
21311                        t,
21312                        t_kv,
21313                        scale,
21314                        true,
21315                        kvl.k_tok_bytes,
21316                        kvl.v_tok_bytes,
21317                    )?;
21318                } else {
21319                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
21320                    // reach past the window, so the window mask is a no-op under causal and every
21321                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
21322                    // request either way, which is what makes the chunk size arithmetic-free.
21323                    e.fa_prefill_view_ws(
21324                        &q,
21325                        &k_view,
21326                        &v_view,
21327                        &mut attn,
21328                        hd,
21329                        nh,
21330                        nkv,
21331                        t,
21332                        t_kv,
21333                        scale,
21334                        true,
21335                        kvl.k_tok_bytes,
21336                        kvl.v_tok_bytes,
21337                        crate::Engine::kv_fp8_on(),
21338                    )?;
21339                }
21340            }
21341            None => {
21342                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
21343                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
21344                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
21345                // seq_end here too or it re-opens the same door.
21346                debug_assert_eq!(
21347                    seq_end, t,
21348                    "step35 cacheless prefill is monolithic (seq_end == t)"
21349                );
21350                if swa && seq_end > win {
21351                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
21352                } else if std::env::var("MEMRA_NOFA").is_ok() {
21353                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21354                } else {
21355                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21356                }
21357            }
21358        }
21359
21360        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
21361        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
21362        let gw = fa
21363            .attn_gate
21364            .as_ref()
21365            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
21366        let gt_owned = if gt_pre.is_none() {
21367            Some(e.matmul(
21368                gw,
21369                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
21370                t,
21371            )?)
21372        } else {
21373            None
21374        };
21375        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
21376        let mut ag = e.uninit(t * nh * hd)?;
21377        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
21378        Ok(ag)
21379    }
21380
21381    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
21382    /// `forward_last`, t2probe). Post-`wo`.
21383    pub(crate) fn step35_attn(
21384        &self,
21385        e: &Engine,
21386        fa: &FullAttnLayer,
21387        h: &CudaSlice<f32>,
21388        pos_d: &CudaSlice<i32>,
21389        t: usize,
21390        il: usize,
21391    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21392        let g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
21393            Some(g3) => g3,
21394            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21395        };
21396        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
21397        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
21398        self.full_attn_o(e, fa, &ag, t)
21399    }
21400
21401    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
21402    /// resident quantized cache, attend through the cache view). Post-`wo`.
21403    ///
21404    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
21405    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
21406    /// own extent.
21407    #[allow(clippy::too_many_arguments)]
21408    pub(crate) fn step35_attn_prime(
21409        &self,
21410        e: &Engine,
21411        fa: &FullAttnLayer,
21412        h: &CudaSlice<f32>,
21413        hx: Option<&CudaSlice<u8>>,
21414        pos_d: &CudaSlice<i32>,
21415        t: usize,
21416        cache: &mut Cache,
21417        il: usize,
21418        seq_end: usize,
21419    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21420        if step_tp_prefill_enabled()? && fa.step_tp_qkv.is_some() {
21421            if hx.is_some() {
21422                return Err(
21423                    "rank-local Step prefill preserves BF16 activations and refuses the q8_1 \
21424                     pre-quantized prime path"
21425                        .into(),
21426                );
21427            }
21428            return self.step35_tp_prefill_attn_resident(e, fa, il, h, pos_d, t, cache, seq_end);
21429        }
21430        let g3 = if fa.step_tp_qkv.is_some() {
21431            if hx.is_some() {
21432                return Err(
21433                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
21434                     pre-quantized prime path"
21435                        .into(),
21436                );
21437            }
21438            self.full_attn_tp_qkv(e, fa, h, t)?
21439                .expect("Step Q/K/V TP disappeared after the presence check")
21440        } else {
21441            match hx {
21442                Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
21443                None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21444            }
21445        };
21446        let ag =
21447            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
21448        self.full_attn_o(e, fa, &ag, t)
21449    }
21450
21451    fn ensure_step_tp_kv_cache(
21452        &self,
21453        e: &Engine,
21454        fa: &FullAttnLayer,
21455        il: usize,
21456        cache: &mut Cache,
21457    ) -> Result<bool, Box<dyn std::error::Error>> {
21458        let tp = fa
21459            .step_tp_qkv
21460            .as_ref()
21461            .ok_or("Step TP cache hydration lost its resident projections")?;
21462        let geometry = self.cfg.full_attention_geometry_at(il as u32);
21463        let window = geometry.window.map(|window| window as usize);
21464        let ranks = tp.runtime.devices().len();
21465        let head_dim = geometry.head_dim_k as usize;
21466        let kv_heads = geometry.n_head_kv as usize;
21467        let max_ctx = cache.max_ctx;
21468
21469        if cache.tp_kv[il].is_some() {
21470            return Ok(false);
21471        }
21472        let local = cache.kv[il]
21473            .as_ref()
21474            .ok_or_else(|| format!("Step TP layer {il} has no owning-stage KV cache"))?;
21475        if local.kv_dim_k != kv_heads * head_dim || local.kv_dim_v != kv_heads * head_dim {
21476            return Err(format!(
21477                "Step TP layer {il} local KV geometry k={} v={} != {}",
21478                local.kv_dim_k,
21479                local.kv_dim_v,
21480                kv_heads * head_dim
21481            )
21482            .into());
21483        }
21484        let resident_start = window
21485            .map(|window| local.len.saturating_sub(window.saturating_sub(1)) & !31usize)
21486            .unwrap_or(0);
21487        let resident_rows = local.len - resident_start;
21488        let physical = local.physical_rows(resident_start, local.len)?;
21489        let k_rows = if resident_rows == 0 {
21490            Vec::new()
21491        } else {
21492            e.dtoh_u8_view(&e.view_u8_range(
21493                &local.k,
21494                physical.start * local.k_tok_bytes,
21495                physical.end * local.k_tok_bytes,
21496            ))?
21497        };
21498        let v_rows = if resident_rows == 0 {
21499            Vec::new()
21500        } else {
21501            e.dtoh_u8_view(&e.view_u8_range(
21502                &local.v,
21503                physical.start * local.v_tok_bytes,
21504                physical.end * local.v_tok_bytes,
21505            ))?
21506        };
21507        let mut distributed = match window {
21508            Some(window) => tp.runtime.allocate_tp_swa_kv_cache(
21509                kv_heads * head_dim,
21510                kv_heads * head_dim,
21511                max_ctx,
21512                window,
21513            )?,
21514            None => tp.runtime.allocate_tp_kv_cache(
21515                kv_heads * head_dim,
21516                kv_heads * head_dim,
21517                max_ctx,
21518            )?,
21519        };
21520        if distributed.k_tok_bytes() * ranks != local.k_tok_bytes
21521            || distributed.v_tok_bytes() * ranks != local.v_tok_bytes
21522        {
21523            return Err(format!(
21524                "Step TP layer {il} distributed/local KV token bytes disagree: \
21525                 k={}x{ranks}/{} v={}x{ranks}/{}",
21526                distributed.k_tok_bytes(),
21527                local.k_tok_bytes,
21528                distributed.v_tok_bytes(),
21529                local.v_tok_bytes,
21530            )
21531            .into());
21532        }
21533        tp.runtime.hydrate_tp_kv_cache_from(
21534            &mut distributed,
21535            local.len,
21536            resident_start,
21537            &k_rows,
21538            &v_rows,
21539        )?;
21540        cache.tp_kv[il] = Some(distributed);
21541        Ok(true)
21542    }
21543
21544    #[allow(clippy::too_many_arguments)]
21545    #[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
21546    fn step35_tp_prefill_attn_resident(
21547        &self,
21548        e: &Engine,
21549        fa: &FullAttnLayer,
21550        il: usize,
21551        h: &CudaSlice<f32>,
21552        pos_d: &CudaSlice<i32>,
21553        tokens: usize,
21554        cache: &mut Cache,
21555        seq_end: usize,
21556    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21557        let tp = fa
21558            .step_tp_qkv
21559            .as_ref()
21560            .ok_or("Step TP prefill lost its resident projections")?;
21561        let attention = tp
21562            .attention
21563            .as_ref()
21564            .ok_or("Step TP prefill lost its resident attention auxiliaries")?;
21565        let ranks = tp.runtime.devices().len();
21566        if !step_tp_prefill_shape(
21567            true,
21568            tokens,
21569            ranks,
21570            tp.runtime.native_p2p(),
21571            true,
21572            crate::Engine::kv_fp8_on(),
21573        ) {
21574            return Err(format!(
21575                "rank-local Step prefill requires tokens>={PRIME_MIN_T}, TP2/TP4 native P2P, \
21576                 rank-local attention, and q8_0/q5_1 KV; got tokens={tokens} ranks={ranks} \
21577                 native_p2p={} fp8_kv={}",
21578                tp.runtime.native_p2p(),
21579                crate::Engine::kv_fp8_on(),
21580            )
21581            .into());
21582        }
21583        for seam in [
21584            "MEMRA_STEP35_SWA_TKV",
21585            "MEMRA_PRIME_CALLLOCAL",
21586            "MEMRA_PRIME_F32CHUNK0",
21587        ] {
21588            if std::env::var(seam).as_deref() == Ok("1") {
21589                return Err(format!(
21590                    "rank-local Step prefill has not qualified the legacy seam {seam}=1"
21591                )
21592                .into());
21593            }
21594        }
21595
21596        let geometry = self.cfg.full_attention_geometry_at(il as u32);
21597        let window = geometry.window.map(|window| window as usize);
21598        let head_dim = geometry.head_dim_k as usize;
21599        let heads = geometry.n_head as usize;
21600        let kv_heads = geometry.n_head_kv as usize;
21601        if heads % ranks != 0 || kv_heads % ranks != 0 {
21602            return Err(format!(
21603                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
21604            )
21605            .into());
21606        }
21607        let local_heads = heads / ranks;
21608        let local_kv_heads = kv_heads / ranks;
21609        let local_kv_dim = local_kv_heads * head_dim;
21610        let hidden = self.cfg.n_embd as usize;
21611        let expected_input = tokens
21612            .checked_mul(hidden)
21613            .ok_or("Step TP prefill input size overflow")?;
21614        if h.len() < expected_input {
21615            return Err(format!(
21616                "Step TP prefill input {} is shorter than {tokens}x{hidden}",
21617                h.len()
21618            )
21619            .into());
21620        }
21621        let positions = e.dtoh_i32(pos_d)?;
21622        if positions.len() != tokens {
21623            return Err(format!(
21624                "rank-local Step prefill positions {} != tokens {tokens}",
21625                positions.len()
21626            )
21627            .into());
21628        }
21629
21630        let mut active_input = e.uninit(expected_input)?;
21631        e.copy_view_into(
21632            &mut active_input,
21633            0,
21634            &h.slice(0..expected_input),
21635            expected_input,
21636        )?;
21637        let mut input = tp.runtime.allocate_replicated_device_rows(tokens, hidden)?;
21638        // PRODUCER FENCE (2026-08-20 flake fix): active_input was written on the MODEL engine's
21639        // stream; the refresh below reads it from the runtime root engine's stream (same device,
21640        // different stream). Unfenced, the peer read can overtake the in-flight copy — the
21641        // layer-count-amplified arm of the boot flake.
21642        e.stream().synchronize()?;
21643        tp.runtime
21644            .refresh_replicated_device_rows_from_root(&mut input, &active_input)?;
21645        let q_raw = tp
21646            .runtime
21647            .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &input)?;
21648        let k_raw = tp
21649            .runtime
21650            .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &input)?;
21651        let v_raw = tp
21652            .runtime
21653            .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &input)?;
21654        let mut q = Vec::with_capacity(ranks);
21655        let mut k = Vec::with_capacity(ranks);
21656        for rank in 0..ranks {
21657            let engine = tp
21658                .runtime
21659                .rank_engine(rank)
21660                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21661            let _main = engine.gpu.enter_main()?;
21662            let mut q_rank = engine.uninit(tokens * local_heads * head_dim)?;
21663            engine.rms_norm(
21664                &q_raw[rank],
21665                &attention.q_norm[rank],
21666                &mut q_rank,
21667                head_dim,
21668                tokens * local_heads,
21669                self.cfg.rms_eps,
21670            )?;
21671            let mut k_rank = engine.uninit(tokens * local_kv_dim)?;
21672            engine.rms_norm(
21673                &k_raw[rank],
21674                &attention.k_norm[rank],
21675                &mut k_rank,
21676                head_dim,
21677                tokens * local_kv_heads,
21678                self.cfg.rms_eps,
21679            )?;
21680            let position = engine.htod_i32(&positions)?;
21681            let rope_freqs = if geometry.rope_factors {
21682                self.step35_aux
21683                    .as_ref()
21684                    .and_then(|aux| aux.rope_freqs(engine))
21685            } else {
21686                None
21687            };
21688            engine.rope_neox2(
21689                &mut q_rank,
21690                &mut k_rank,
21691                &position,
21692                head_dim,
21693                geometry.n_rot as usize,
21694                local_heads,
21695                local_kv_heads,
21696                tokens,
21697                geometry.rope_base,
21698                1.0,
21699                rope_freqs,
21700            )?;
21701            q.push(q_rank);
21702            k.push(k_rank);
21703        }
21704
21705        let gate_weight = fa
21706            .attn_gate
21707            .as_ref()
21708            .ok_or("step35 layer is missing attn_gate.weight")?;
21709        let gate = e.dtoh(&e.matmul(gate_weight, h, tokens)?)?;
21710        if gate.len() != tokens * heads {
21711            return Err(format!(
21712                "Step TP layer {il} gate output {} != {tokens}x{heads}",
21713                gate.len()
21714            )
21715            .into());
21716        }
21717
21718        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
21719        let base_len = cache.kv[il]
21720            .as_ref()
21721            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
21722            .len;
21723        let distributed = cache.tp_kv[il]
21724            .as_ref()
21725            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
21726        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
21727            return Err(format!(
21728                "Step TP layer {il} cache lengths diverged before prefill: \
21729                 local={base_len} distributed={}/{}",
21730                distributed.committed_len(),
21731                distributed.staged_len()
21732            )
21733            .into());
21734        }
21735        let target_len = base_len
21736            .checked_add(tokens)
21737            .ok_or("Step TP prefill cache length overflow")?;
21738        if target_len > cache.max_ctx {
21739            return Err(format!(
21740                "Step TP layer {il} prefill exceeds cache: {base_len}+{tokens}>{}",
21741                cache.max_ctx
21742            )
21743            .into());
21744        }
21745        if seq_end < target_len {
21746            return Err(format!(
21747                "Step TP layer {il} request end {seq_end} precedes chunk end {target_len}"
21748            )
21749            .into());
21750        }
21751
21752        let transaction = cache.tp_kv[il]
21753            .as_mut()
21754            .expect("distributed cache checked above")
21755            .begin_transaction()?;
21756        if let Err(error) = tp.runtime.append_tp_kv_transaction(
21757            cache.tp_kv[il]
21758                .as_mut()
21759                .expect("distributed cache checked above"),
21760            transaction,
21761            &k,
21762            &v_raw,
21763            tokens,
21764        ) {
21765            let _ = tp.runtime.rollback_tp_kv_transaction(
21766                cache.tp_kv[il]
21767                    .as_mut()
21768                    .expect("distributed cache checked above"),
21769                transaction,
21770            );
21771            return Err(error);
21772        }
21773
21774        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21775            let distributed = cache.tp_kv[il]
21776                .as_ref()
21777                .expect("distributed cache checked above");
21778            let staged_len = distributed.staged_len();
21779            let view_start = window
21780                .map(|window| base_len.saturating_sub(window.saturating_sub(1)) & !31usize)
21781                .unwrap_or(0);
21782            let physical = distributed.physical_range(view_start, staged_len)?;
21783            let t_kv = staged_len - view_start;
21784            let swa_naive = window.is_some_and(|window| seq_end > window);
21785            let mut gated = Vec::with_capacity(ranks);
21786            #[allow(clippy::needless_range_loop)]
21787            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
21788            for rank in 0..ranks {
21789                let engine = tp
21790                    .runtime
21791                    .rank_engine(rank)
21792                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21793                let _main = engine.gpu.enter_main()?;
21794                let rank_cache = distributed
21795                    .rank(rank)
21796                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
21797                let k_view = engine.view_u8_range(
21798                    rank_cache.k(),
21799                    physical.start * distributed.k_tok_bytes(),
21800                    physical.end * distributed.k_tok_bytes(),
21801                );
21802                let v_view = engine.view_u8_range(
21803                    rank_cache.v(),
21804                    physical.start * distributed.v_tok_bytes(),
21805                    physical.end * distributed.v_tok_bytes(),
21806                );
21807                let mut attention_out = engine.uninit(tokens * local_heads * head_dim)?;
21808                if swa_naive {
21809                    let window = window.expect("SWA predicate requires a window");
21810                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21811                        engine.sdpa_naive_w_quantized_view(
21812                            &q[rank],
21813                            &k_view,
21814                            &v_view,
21815                            &mut attention_out,
21816                            head_dim,
21817                            local_heads,
21818                            local_kv_heads,
21819                            tokens,
21820                            t_kv,
21821                            geometry.attention_scale(),
21822                            true,
21823                            window,
21824                            distributed.k_tok_bytes(),
21825                            distributed.v_tok_bytes(),
21826                        )?;
21827                    } else {
21828                        engine.fa_prefill_view_ws_w_hd128(
21829                            &q[rank],
21830                            &k_view,
21831                            &v_view,
21832                            &mut attention_out,
21833                            head_dim,
21834                            local_heads,
21835                            local_kv_heads,
21836                            tokens,
21837                            t_kv,
21838                            geometry.attention_scale(),
21839                            true,
21840                            window,
21841                            distributed.k_tok_bytes(),
21842                            distributed.v_tok_bytes(),
21843                        )?;
21844                    }
21845                } else if std::env::var("MEMRA_NOFA").is_ok() {
21846                    engine.sdpa_naive_quantized_view(
21847                        &q[rank],
21848                        &k_view,
21849                        &v_view,
21850                        &mut attention_out,
21851                        head_dim,
21852                        local_heads,
21853                        local_kv_heads,
21854                        tokens,
21855                        t_kv,
21856                        geometry.attention_scale(),
21857                        true,
21858                        distributed.k_tok_bytes(),
21859                        distributed.v_tok_bytes(),
21860                    )?;
21861                } else {
21862                    engine.fa_prefill_view_ws(
21863                        &q[rank],
21864                        &k_view,
21865                        &v_view,
21866                        &mut attention_out,
21867                        head_dim,
21868                        local_heads,
21869                        local_kv_heads,
21870                        tokens,
21871                        t_kv,
21872                        geometry.attention_scale(),
21873                        true,
21874                        distributed.k_tok_bytes(),
21875                        distributed.v_tok_bytes(),
21876                        false,
21877                    )?;
21878                }
21879
21880                let gate_start = rank * local_heads;
21881                let mut gate_rank = Vec::with_capacity(tokens * local_heads);
21882                for token in 0..tokens {
21883                    let start = token * heads + gate_start;
21884                    gate_rank.extend_from_slice(&gate[start..start + local_heads]);
21885                }
21886                let gate_rank = engine.htod(&gate_rank)?;
21887                let mut gated_rank = engine.uninit(tokens * local_heads * head_dim)?;
21888                engine.attn_head_gate(
21889                    &attention_out,
21890                    &gate_rank,
21891                    &mut gated_rank,
21892                    None,
21893                    head_dim,
21894                    local_heads,
21895                    tokens,
21896                )?;
21897                gated.push(gated_rank);
21898            }
21899            for rank in 1..ranks {
21900                let engine = tp
21901                    .runtime
21902                    .rank_engine(rank)
21903                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21904                let _main = engine.gpu.enter_main()?;
21905                engine.stream().synchronize()?;
21906            }
21907
21908            let (output, k_shadow, v_shadow) = if tp.runtime.bulk_p2p() {
21909                let output = tp
21910                    .runtime
21911                    .step_bf16_row_parallel_resident_root_device(&tp.o, &gated, tokens)?;
21912                let k_shadow =
21913                    tp.runtime
21914                        .gather_native_column_shards_device(&k, tokens, local_kv_dim)?;
21915                let v_shadow =
21916                    tp.runtime
21917                        .gather_native_column_shards_device(&v_raw, tokens, local_kv_dim)?;
21918                let root = tp
21919                    .runtime
21920                    .rank_engine(0)
21921                    .ok_or("Step TP prefill lost its root engine")?;
21922                let _main = root.gpu.enter_main()?;
21923                root.stream().synchronize()?;
21924                (output, k_shadow, v_shadow)
21925            } else {
21926                let attention = tp.runtime.gather_native_column_shards(
21927                    &gated,
21928                    tokens,
21929                    local_heads * head_dim,
21930                )?;
21931                let output = tp
21932                    .runtime
21933                    .step_bf16_row_parallel_resident_native(&tp.o, &attention, tokens)?;
21934                let k_shadow = tp
21935                    .runtime
21936                    .gather_native_column_shards(&k, tokens, local_kv_dim)?;
21937                let v_shadow =
21938                    tp.runtime
21939                        .gather_native_column_shards(&v_raw, tokens, local_kv_dim)?;
21940                (e.htod(&output)?, e.htod(&k_shadow)?, e.htod(&v_shadow)?)
21941            };
21942            let local = cache.kv[il]
21943                .as_mut()
21944                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
21945            if local.len != base_len {
21946                return Err(format!(
21947                    "Step TP layer {il} local cache changed during prefill: \
21948                     len={} base={base_len}",
21949                    local.len
21950                )
21951                .into());
21952            }
21953            let retain_from = window
21954                .map(|window| {
21955                    let staged_retain = staged_len.saturating_sub(window) & !31usize;
21956                    let rollback_retain =
21957                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
21958                    staged_retain.min(rollback_retain)
21959                })
21960                .unwrap_or(0);
21961            let write_row = e.prepare_kv_append(local, retain_from, tokens)?;
21962            e.append_kv_quantized_rows(
21963                &k_shadow,
21964                &v_shadow,
21965                &mut local.k,
21966                &mut local.v,
21967                write_row,
21968                tokens,
21969                local.kv_dim_k,
21970                local.kv_dim_v,
21971                local.k_tok_bytes,
21972                local.v_tok_bytes,
21973                false,
21974            )?;
21975            local.len = staged_len;
21976            e.set_i32_one(&mut local.len_d, staged_len as i32)?;
21977            Ok(output)
21978        })();
21979
21980        let output = match staged {
21981            Ok(output) => output,
21982            Err(error) => {
21983                let _ = tp.runtime.rollback_tp_kv_transaction(
21984                    cache.tp_kv[il]
21985                        .as_mut()
21986                        .expect("distributed cache checked above"),
21987                    transaction,
21988                );
21989                if let Some(local) = cache.kv[il].as_mut() {
21990                    local.len = base_len;
21991                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
21992                }
21993                return Err(error);
21994            }
21995        };
21996        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
21997            cache.tp_kv[il]
21998                .as_mut()
21999                .expect("distributed cache checked above"),
22000            transaction,
22001            tokens,
22002        ) {
22003            let _ = tp.runtime.rollback_tp_kv_transaction(
22004                cache.tp_kv[il]
22005                    .as_mut()
22006                    .expect("distributed cache checked above"),
22007                transaction,
22008            );
22009            let local = cache.kv[il].as_mut().expect("local cache checked above");
22010            local.len = base_len;
22011            e.set_i32_one(&mut local.len_d, base_len as i32)?;
22012            return Err(error);
22013        }
22014
22015        let committed = cache.tp_kv[il]
22016            .as_ref()
22017            .expect("distributed cache checked above")
22018            .committed_len();
22019        let local_len = cache.kv[il]
22020            .as_ref()
22021            .expect("local cache checked above")
22022            .len;
22023        if committed != local_len {
22024            return Err(format!(
22025                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
22026            )
22027            .into());
22028        }
22029        eprintln!(
22030            "[step-tp-prefill-attn] execute layer={} devices={:?} tokens={tokens} \
22031             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
22032             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
22033             attention_scope={} input_path=root-device-replicated gate_tensor_parallel=false \
22034             gate_shards=host-canonical o_tensor_parallel=true local_cache_shadow=true \
22035             cache_commit=chunk transport={} native_p2p=true bulk_p2p={} \
22036             output={} performance_claim=false",
22037            tp.layer,
22038            tp.devices,
22039            hydrated,
22040            if window.is_some() {
22041                "rank-local-swa-ring"
22042            } else {
22043                "rank-local-global"
22044            },
22045            tp.runtime.transport_label(),
22046            tp.runtime.bulk_p2p(),
22047            if tp.runtime.bulk_p2p() {
22048                "root-device"
22049            } else {
22050                "root-readback"
22051            },
22052        );
22053        Ok(output)
22054    }
22055
22056    fn step35_tp_decode_attn_resident(
22057        &self,
22058        e: &Engine,
22059        fa: &FullAttnLayer,
22060        il: usize,
22061        h: &CudaSlice<f32>,
22062        pos_d: &CudaSlice<i32>,
22063        cache: &mut Cache,
22064    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22065        // MEMRA_STEP_TP_TIMING=1: cumulative wall-clock of the rank-local TP attention decode,
22066        // printed every ~10 decode steps' worth of layers — the wall-decomposition twin of the
22067        // nvfp4-dev-routes counter.
22068        static ATTN_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22069        static ATTN_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22070        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
22071        let started = timing.then(std::time::Instant::now);
22072        let result = if crate::tp::step_tp_decode_v2_enabled()? {
22073            self.step35_tp_decode_attn_resident_v2(e, fa, il, h, pos_d, cache)
22074        } else {
22075            self.step35_tp_decode_attn_resident_inner(e, fa, il, h, pos_d, cache)
22076        };
22077        if let Some(started) = started {
22078            use std::sync::atomic::Ordering;
22079            let ns = ATTN_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
22080                + started.elapsed().as_nanos() as u64;
22081            let calls = ATTN_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
22082            if calls.is_multiple_of(430) {
22083                eprintln!(
22084                    "[step-tp-attn-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
22085                    ns as f64 / 1.0e6,
22086                    ns as f64 / calls as f64 / 1.0e3,
22087                );
22088            }
22089        }
22090        result
22091    }
22092
22093    #[allow(clippy::too_many_arguments)]
22094    fn step35_tp_decode_attn_resident_inner(
22095        &self,
22096        e: &Engine,
22097        fa: &FullAttnLayer,
22098        il: usize,
22099        h: &CudaSlice<f32>,
22100        pos_d: &CudaSlice<i32>,
22101        cache: &mut Cache,
22102    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22103        // MEMRA_STEP_TP_TIMING=1 phase decomposition of the 1550us/layer decode wall. Each lap
22104        // drains every stream so queued async work is billed to the phase that queued it — the
22105        // drains perturb absolute wall, but v1 already ends most phases on a host sync, so the
22106        // relative split is honest. Timing OFF is the measured configuration: zero extra syncs.
22107        static T_POS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22108        static T_QKV: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22109        static T_NORMROPE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22110        static T_GATE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22111        static T_APPEND: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22112        static T_ATTN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22113        static T_OPROJ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22114        static T_SHADOW: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22115        static T_PHASE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
22116        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
22117        #[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
22118        fn lap(
22119            runtime: &crate::tp::TpE4m3HostBounce,
22120            e: &Engine,
22121            timer: &std::sync::atomic::AtomicU64,
22122            started: &mut Option<std::time::Instant>,
22123        ) -> Result<(), Box<dyn std::error::Error>> {
22124            let Some(start) = started.as_mut() else {
22125                return Ok(());
22126            };
22127            for rank in 0..runtime.devices().len() {
22128                if let Some(engine) = runtime.rank_engine(rank) {
22129                    let _main = engine.gpu.enter_main()?;
22130                    engine.stream().synchronize()?;
22131                }
22132            }
22133            e.stream().synchronize()?;
22134            timer.fetch_add(
22135                start.elapsed().as_nanos() as u64,
22136                std::sync::atomic::Ordering::Relaxed,
22137            );
22138            *start = std::time::Instant::now();
22139            Ok(())
22140        }
22141        let tp = fa
22142            .step_tp_qkv
22143            .as_ref()
22144            .ok_or("Step TP decode lost its resident projections")?;
22145        let attention = tp
22146            .attention
22147            .as_ref()
22148            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
22149        if !tp.runtime.native_p2p() {
22150            return Err("rank-local Step attention requires native P2P".into());
22151        }
22152        if crate::Engine::kv_fp8_on() {
22153            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
22154        }
22155
22156        let geometry = self.step35_geom(il);
22157        let window = geometry.window.map(|window| window as usize);
22158        let ranks = tp.runtime.devices().len();
22159        let head_dim = geometry.head_dim_k as usize;
22160        let heads = geometry.n_head as usize;
22161        let kv_heads = geometry.n_head_kv as usize;
22162        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
22163            return Err(format!(
22164                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
22165            )
22166            .into());
22167        }
22168        let local_heads = heads / ranks;
22169        let local_kv_heads = kv_heads / ranks;
22170        let local_kv_dim = local_kv_heads * head_dim;
22171        let max_ctx = cache.max_ctx;
22172
22173        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
22174
22175        let base_len = cache.kv[il]
22176            .as_ref()
22177            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
22178            .len;
22179        let distributed = cache.tp_kv[il]
22180            .as_ref()
22181            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
22182        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
22183            return Err(format!(
22184                "Step TP layer {il} cache lengths diverged before decode: \
22185                 local={base_len} distributed={}/{}",
22186                distributed.committed_len(),
22187                distributed.staged_len()
22188            )
22189            .into());
22190        }
22191
22192        let mut lap_start = timing.then(std::time::Instant::now);
22193        let positions = e.dtoh_i32(pos_d)?;
22194        if positions.len() != 1 {
22195            return Err(format!(
22196                "rank-local Step decode requires one position, got {}",
22197                positions.len()
22198            )
22199            .into());
22200        }
22201        lap(&tp.runtime, e, &T_POS, &mut lap_start)?;
22202        let (q_raw, k_raw, v_raw, input_path) = if let Some(decode_input) =
22203            attention.decode_input.as_ref()
22204        {
22205            let mut decode_input = decode_input
22206                .lock()
22207                .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
22208            // PRODUCER FENCE (2026-08-20 flake fix): h is the layer hidden written on the MODEL
22209            // engine's stream; the refresh reads it from the runtime root engine's stream. This
22210            // per-layer-per-token seam is the layer-count-amplified arm of the boot flake.
22211            e.stream().synchronize()?;
22212            tp.runtime
22213                .refresh_replicated_device_rows_from_root(&mut decode_input, h)?;
22214            let q_raw = tp
22215                .runtime
22216                .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &decode_input)?;
22217            let k_raw = tp
22218                .runtime
22219                .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &decode_input)?;
22220            let v_raw = tp
22221                .runtime
22222                .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &decode_input)?;
22223            (q_raw, k_raw, v_raw, "root-device-replicated")
22224        } else {
22225            let activation = e.dtoh(h)?;
22226            let q_raw =
22227                tp.runtime
22228                    .bf16_column_parallel_resident_device_shards(&tp.q, &activation, 1)?;
22229            let k_raw =
22230                tp.runtime
22231                    .bf16_column_parallel_resident_device_shards(&tp.k, &activation, 1)?;
22232            let v_raw =
22233                tp.runtime
22234                    .bf16_column_parallel_resident_device_shards(&tp.v, &activation, 1)?;
22235            (q_raw, k_raw, v_raw, "host-replicated")
22236        };
22237        lap(&tp.runtime, e, &T_QKV, &mut lap_start)?;
22238        let mut q = Vec::with_capacity(ranks);
22239        let mut k = Vec::with_capacity(ranks);
22240        for rank in 0..ranks {
22241            let engine = tp
22242                .runtime
22243                .rank_engine(rank)
22244                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22245            let _main = engine.gpu.enter_main()?;
22246            let mut q_rank = engine.uninit(local_heads * head_dim)?;
22247            engine.rms_norm(
22248                &q_raw[rank],
22249                &attention.q_norm[rank],
22250                &mut q_rank,
22251                head_dim,
22252                local_heads,
22253                self.cfg.rms_eps,
22254            )?;
22255            let mut k_rank = engine.uninit(local_kv_dim)?;
22256            engine.rms_norm(
22257                &k_raw[rank],
22258                &attention.k_norm[rank],
22259                &mut k_rank,
22260                head_dim,
22261                local_kv_heads,
22262                self.cfg.rms_eps,
22263            )?;
22264            let position = engine.htod_i32(&positions)?;
22265            let rope_freqs = if geometry.rope_factors {
22266                self.step35_aux
22267                    .as_ref()
22268                    .and_then(|aux| aux.rope_freqs(engine))
22269            } else {
22270                None
22271            };
22272            engine.rope_neox2(
22273                &mut q_rank,
22274                &mut k_rank,
22275                &position,
22276                head_dim,
22277                geometry.n_rot as usize,
22278                local_heads,
22279                local_kv_heads,
22280                1,
22281                geometry.rope_base,
22282                1.0,
22283                rope_freqs,
22284            )?;
22285            q.push(q_rank);
22286            k.push(k_rank);
22287        }
22288        lap(&tp.runtime, e, &T_NORMROPE, &mut lap_start)?;
22289
22290        let gate_weight = fa
22291            .attn_gate
22292            .as_ref()
22293            .ok_or("step35 layer is missing attn_gate.weight")?;
22294        let gate = e.matmul(gate_weight, h, 1)?;
22295        let gate = e.dtoh(&gate)?;
22296        if gate.len() != heads {
22297            return Err(format!("Step TP layer {il} gate output {} != {heads}", gate.len()).into());
22298        }
22299        lap(&tp.runtime, e, &T_GATE, &mut lap_start)?;
22300
22301        let transaction = cache.tp_kv[il]
22302            .as_mut()
22303            .expect("distributed cache checked above")
22304            .begin_transaction()?;
22305        if let Err(error) = tp.runtime.append_tp_kv_transaction(
22306            cache.tp_kv[il]
22307                .as_mut()
22308                .expect("distributed cache checked above"),
22309            transaction,
22310            &k,
22311            &v_raw,
22312            1,
22313        ) {
22314            let _ = tp.runtime.rollback_tp_kv_transaction(
22315                cache.tp_kv[il]
22316                    .as_mut()
22317                    .expect("distributed cache checked above"),
22318                transaction,
22319            );
22320            return Err(error);
22321        }
22322        lap(&tp.runtime, e, &T_APPEND, &mut lap_start)?;
22323
22324        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22325            let distributed = cache.tp_kv[il]
22326                .as_ref()
22327                .expect("distributed cache checked above");
22328            let staged_len = distributed.staged_len();
22329            let view_start = window
22330                .map(|window| staged_len.saturating_sub(window))
22331                .unwrap_or(0);
22332            let physical = distributed.physical_range(view_start, staged_len)?;
22333            let t_kv = staged_len - view_start;
22334            let mut gated = Vec::with_capacity(ranks);
22335            #[allow(clippy::needless_range_loop)]
22336            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
22337            for rank in 0..ranks {
22338                let engine = tp
22339                    .runtime
22340                    .rank_engine(rank)
22341                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22342                let _main = engine.gpu.enter_main()?;
22343                let rank_cache = distributed
22344                    .rank(rank)
22345                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
22346                let k_view = engine.view_u8_range(
22347                    rank_cache.k(),
22348                    physical.start * distributed.k_tok_bytes(),
22349                    physical.end * distributed.k_tok_bytes(),
22350                );
22351                let v_view = engine.view_u8_range(
22352                    rank_cache.v(),
22353                    physical.start * distributed.v_tok_bytes(),
22354                    physical.end * distributed.v_tok_bytes(),
22355                );
22356                let mut attention_out = engine.uninit(local_heads * head_dim)?;
22357                engine.fa_decode_kvmod(
22358                    &q[rank],
22359                    &k_view,
22360                    &v_view,
22361                    &mut attention_out,
22362                    head_dim,
22363                    local_heads,
22364                    local_kv_heads,
22365                    t_kv,
22366                    geometry.attention_scale(),
22367                    distributed.k_tok_bytes(),
22368                    distributed.v_tok_bytes(),
22369                    false,
22370                )?;
22371                let gate_start = rank * local_heads;
22372                let gate_rank = engine.htod(&gate[gate_start..gate_start + local_heads])?;
22373                let mut gated_rank = engine.uninit(local_heads * head_dim)?;
22374                engine.attn_head_gate(
22375                    &attention_out,
22376                    &gate_rank,
22377                    &mut gated_rank,
22378                    None,
22379                    head_dim,
22380                    local_heads,
22381                    1,
22382                )?;
22383                gated.push(gated_rank);
22384            }
22385            lap(&tp.runtime, e, &T_ATTN, &mut lap_start)?;
22386
22387            let gathered =
22388                tp.runtime
22389                    .gather_native_column_shards(&gated, 1, local_heads * head_dim)?;
22390            let output = tp
22391                .runtime
22392                .step_bf16_row_parallel_resident_native(&tp.o, &gathered, 1)?;
22393            let output = e.htod(&output)?;
22394            lap(&tp.runtime, e, &T_OPROJ, &mut lap_start)?;
22395
22396            let k_shadow = tp
22397                .runtime
22398                .gather_native_column_shards(&k, 1, local_kv_dim)?;
22399            let v_shadow = tp
22400                .runtime
22401                .gather_native_column_shards(&v_raw, 1, local_kv_dim)?;
22402            let k_shadow = e.htod(&k_shadow)?;
22403            let v_shadow = e.htod(&v_shadow)?;
22404            let local = cache.kv[il]
22405                .as_mut()
22406                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
22407            if local.len != base_len || base_len + 1 > max_ctx {
22408                return Err(format!(
22409                    "Step TP layer {il} local cache changed during decode: \
22410                     len={} base={base_len} max={max_ctx}",
22411                    local.len
22412                )
22413                .into());
22414            }
22415            let retain_from = window
22416                .map(|window| {
22417                    let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
22418                    let rollback_retain =
22419                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
22420                    staged_retain.min(rollback_retain)
22421                })
22422                .unwrap_or(0);
22423            let write_row = e.prepare_kv_append(local, retain_from, 1)?;
22424            e.append_kv_quantized(
22425                &k_shadow,
22426                &v_shadow,
22427                &mut local.k,
22428                &mut local.v,
22429                write_row,
22430                local.kv_dim_k,
22431                local.kv_dim_v,
22432                local.k_tok_bytes,
22433                local.v_tok_bytes,
22434                false,
22435            )?;
22436            local.len = base_len + 1;
22437            e.set_i32_one(&mut local.len_d, local.len as i32)?;
22438            Ok(output)
22439        })();
22440
22441        let output = match staged {
22442            Ok(output) => output,
22443            Err(error) => {
22444                let _ = tp.runtime.rollback_tp_kv_transaction(
22445                    cache.tp_kv[il]
22446                        .as_mut()
22447                        .expect("distributed cache checked above"),
22448                    transaction,
22449                );
22450                if let Some(local) = cache.kv[il].as_mut() {
22451                    local.len = base_len;
22452                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
22453                }
22454                return Err(error);
22455            }
22456        };
22457        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
22458            cache.tp_kv[il]
22459                .as_mut()
22460                .expect("distributed cache checked above"),
22461            transaction,
22462            1,
22463        ) {
22464            let _ = tp.runtime.rollback_tp_kv_transaction(
22465                cache.tp_kv[il]
22466                    .as_mut()
22467                    .expect("distributed cache checked above"),
22468                transaction,
22469            );
22470            let local = cache.kv[il].as_mut().expect("local cache checked above");
22471            local.len = base_len;
22472            e.set_i32_one(&mut local.len_d, base_len as i32)?;
22473            return Err(error);
22474        }
22475
22476        let committed = cache.tp_kv[il]
22477            .as_ref()
22478            .expect("distributed cache checked above")
22479            .committed_len();
22480        let local_len = cache.kv[il]
22481            .as_ref()
22482            .expect("local cache checked above")
22483            .len;
22484        if committed != local_len {
22485            return Err(format!(
22486                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
22487            )
22488            .into());
22489        }
22490        lap(&tp.runtime, e, &T_SHADOW, &mut lap_start)?;
22491        if timing {
22492            use std::sync::atomic::Ordering;
22493            let calls = T_PHASE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
22494            if calls.is_multiple_of(430) {
22495                let avg = |t: &std::sync::atomic::AtomicU64| {
22496                    t.load(Ordering::Relaxed) as f64 / calls as f64 / 1.0e3
22497                };
22498                eprintln!(
22499                    "[step-tp-attn-phase] calls={calls} avg_us pos={:.1} qkv={:.1} \
22500                     normrope={:.1} gate={:.1} append={:.1} attn={:.1} oproj={:.1} shadow={:.1}",
22501                    avg(&T_POS),
22502                    avg(&T_QKV),
22503                    avg(&T_NORMROPE),
22504                    avg(&T_GATE),
22505                    avg(&T_APPEND),
22506                    avg(&T_ATTN),
22507                    avg(&T_OPROJ),
22508                    avg(&T_SHADOW),
22509                );
22510            }
22511        }
22512        eprintln!(
22513            "[step-tp-attn] execute layer={} devices={:?} tokens=1 \
22514             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
22515             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
22516             attention_scope={} input_path={} kv_physical_rows={} \
22517             gate_tensor_parallel=false gate_shards=host-canonical o_tensor_parallel=true \
22518             local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
22519             bulk_p2p={} output=root-readback performance_claim=false",
22520            tp.layer,
22521            tp.devices,
22522            hydrated,
22523            if window.is_some() {
22524                "rank-local-swa-ring"
22525            } else {
22526                "rank-local-global"
22527            },
22528            input_path,
22529            cache.tp_kv[il]
22530                .as_ref()
22531                .expect("distributed cache checked above")
22532                .physical_capacity(),
22533            tp.runtime.transport_label(),
22534            tp.runtime.bulk_p2p(),
22535        );
22536        Ok(output)
22537    }
22538
22539    /// v2 rank-local decode attention (MEMRA_STEP_TP_DECODE_V2): the same kernels, operand
22540    /// values, and canonical reduction order as `step35_tp_decode_attn_resident_inner`,
22541    /// restructured onto a persistent per-runtime workspace with evented cross-stream ordering.
22542    /// Per layer per token this path performs exactly one cuMemAlloc (the returned e-context
22543    /// output row), no host round-trip, and no host stream synchronize — the phase timers
22544    /// measured v1 spending 81% of its 1550us/layer wall on those three classes.
22545    #[allow(clippy::too_many_arguments)]
22546    /// T-COLUMN verify precompute for layer `il`: weight-amortized QKV(+gate) over the T
22547    /// verify columns into the ws slabs (per-column rope/append/fa run later through the
22548    /// unmodified t=1 program via the col-select door). Ok(false) when the layer is not on
22549    /// the resident fused TP2 class (caller falls back to the per-row walk).
22550    pub(crate) fn step35_verify_qkv_precompute(
22551        &self,
22552        e: &Engine,
22553        il: usize,
22554        h_t: &CudaSlice<f32>,
22555        t: usize,
22556    ) -> Result<bool, Box<dyn std::error::Error>> {
22557        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22558            return Ok(false);
22559        };
22560        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22561            return Ok(false);
22562        };
22563        let Some(attention) = tp.attention.as_ref() else {
22564            return Ok(false);
22565        };
22566        if !tp.runtime.native_p2p() || !crate::tp::step_tp_qkv_fused_enabled()? {
22567            return Ok(false);
22568        }
22569        let geometry = self.step35_geom(il);
22570        let heads = geometry.n_head as usize;
22571        let ws_index = tp
22572            .runtime
22573            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22574        let gate_shards = attention
22575            .gate_shards_bf16
22576            .as_deref()
22577            .map(crate::tp::StepTpGateShards::Bf16);
22578        tp.runtime.decode_v2_input_qkv_tcol(
22579            ws_index,
22580            e,
22581            h_t,
22582            t,
22583            &tp.q,
22584            &tp.k,
22585            &tp.v,
22586            gate_shards,
22587        )?;
22588        Ok(true)
22589    }
22590
22591    /// MEMRA_TCOL_OPROJ join for the verify walk: after every column of layer `il`
22592    /// stashed its `gated` rows, produce the [t, o_out] `mixed` slab on `e` via the
22593    /// weight-amortized b4_tcol + slab join. Callers only reach this after the stash
22594    /// flag confirmed the defer engaged for every column.
22595    pub(crate) fn step35_verify_oproj_tcol(
22596        &self,
22597        e: &Engine,
22598        il: usize,
22599        t: usize,
22600    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22601        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22602            return Err("tcol o_proj join expects full attention".into());
22603        };
22604        let tp = fa
22605            .step_tp_qkv
22606            .as_ref()
22607            .ok_or("tcol o_proj join lost its resident projections")?;
22608        let heads = self.step35_geom(il).n_head as usize;
22609        let ws_index = tp
22610            .runtime
22611            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22612        tp.runtime.decode_v2_oproj_tcol(ws_index, e, &tp.o, t)
22613    }
22614
22615    /// MEMRA_SPEC_FA2 precheck: decide BEFORE arming the defer whether both verify
22616    /// columns of layer `il` will take the dcw arm AND the T=2 launch is bit-safe —
22617    /// stashing is unrecoverable (no per-column output exists), so every dynamic input
22618    /// to the engine-side dcw decision is evaluated here, plus the equal-partition
22619    /// guard fa_decode_dcw2's contract requires. Boundary rounds return false and the
22620    /// walk runs the ordinary per-column program.
22621    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
22622    pub(crate) fn step35_spec_fa2_precheck(
22623        &self,
22624        cache: &Cache,
22625        il: usize,
22626        pos0: usize,
22627    ) -> Result<bool, Box<dyn std::error::Error>> {
22628        // MEMRA_SPEC_FA2_DEBUG=1: print the first failing clause once per clause id —
22629        // a silently-vacuous door is indistinguishable from a slow one without this.
22630        fn nope(clause: &str, il: usize, pos0: usize) -> bool {
22631            static DBG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22632            static SEEN: std::sync::Mutex<Vec<&'static str>> = std::sync::Mutex::new(Vec::new());
22633            if *DBG.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1")) {
22634                let mut seen = SEEN.lock().unwrap();
22635                if !seen.contains(&clause) {
22636                    // leak: bounded by the clause-id set
22637                    seen.push(Box::leak(clause.to_string().into_boxed_str()));
22638                    eprintln!("[spec-fa2] precheck FAIL clause={clause} il={il} pos0={pos0}");
22639                }
22640            }
22641            false
22642        }
22643        // MEMRA_SPEC_FA2_LAYER=<il>: engage on ONE layer only (divergence bisection).
22644        static ONLY: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
22645        if let Some(only) =
22646            ONLY.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_LAYER").ok()?.parse().ok())
22647            && *only != il
22648        {
22649            return Ok(false);
22650        }
22651        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22652            return Ok(nope("mixer", il, pos0));
22653        };
22654        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22655            return Ok(nope("step_tp", il, pos0));
22656        };
22657        let Some(attention) = tp.attention.as_ref() else {
22658            return Ok(nope("attention", il, pos0));
22659        };
22660        if !tp.runtime.native_p2p()
22661            || crate::Engine::kv_fp8_on()
22662            || !crate::tp::step_tp_dcw_enabled()?
22663            || !crate::tp::step_tp_qkv_fused_enabled()?
22664            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22665        {
22666            return Ok(nope("runtime-doors", il, pos0));
22667        }
22668        let geometry = self.step35_geom(il);
22669        let head_dim = geometry.head_dim_k as usize;
22670        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22671            return Ok(nope("fa-class", il, pos0));
22672        }
22673        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22674            return Ok(nope("tp-kv", il, pos0));
22675        };
22676        if distributed.staged_len() != pos0 {
22677            return Ok(nope("staged-len", il, pos0));
22678        }
22679        // Both appends must land without a ring rebase (rebase columns take the
22680        // host-row path, which cannot stash).
22681        let (_, would_rebase) = distributed.peek_append_ring(2)?;
22682        if would_rebase {
22683            return Ok(nope("rebase", il, pos0));
22684        }
22685        let window = geometry.window.map(|w| w as usize);
22686        // Capped SWA is REDUCTION-CLASS in the joined kernel (the two rows' windows
22687        // shift by one key, so one shared tile grid cannot reproduce both rows'
22688        // per-column FP grouping) — and drifted verify logits change accept decisions,
22689        // breaking the spec==target contract. Engage only when BOTH rows' views start
22690        // at 0 (global, or SWA still inside its window): bitwise per row under the
22691        // partition guard below. At agentic ctx this keeps the global layers — ~3/4 of
22692        // the per-key fa work — and leaves capped-SWA layers on the per-column program.
22693        if let Some(w) = window
22694            && pos0 + 2 > w
22695        {
22696            return Ok(nope("swa-capped", il, pos0));
22697        }
22698        // Row r's own per-column launch sees the POST-append view: T_r = pos0 + 1 + r
22699        // (kernel T_kv = len_dev - lstart; the host bucket matches it — the one-partition
22700        // law). Both dcw eligibility (t_kv_eff >= 96) and the vec floor key off T0.
22701        let (t0, t1) = (pos0 + 1, pos0 + 2);
22702        if t0 < 96 {
22703            return Ok(nope("dcw-floor", il, pos0));
22704        }
22705        if std::env::var("MEMRA_NO_FA_VEC").is_ok() || t0 < crate::fa_vec_min_tkv() {
22706            return Ok(nope("vec-floor", il, pos0));
22707        }
22708        // Equal-partition guard, on the KERNEL's derivation: split width (sp), effective
22709        // count (ns = ceil(T/sp)) and stride (per = ceil(T/ns)) must all match between
22710        // the two rows' own launches — the joined kernel derives one grid from T1 and
22711        // row0 inherits it, so any difference shifts row0's split boundaries and changes
22712        // the combine's merge rounding. Boundary rounds fall back per column.
22713        let ranks = tp.runtime.devices().len();
22714        let local_kv_heads = (geometry.n_head_kv as usize / ranks).max(1);
22715        let sp0 = crate::fa_split_keys_pub(t0, local_kv_heads);
22716        let sp1 = crate::fa_split_keys_pub(t1, local_kv_heads);
22717        if sp0 != sp1 {
22718            return Ok(nope("partition-sp", il, pos0));
22719        }
22720        let (ns0, ns1) = (t0.div_ceil(sp0), t1.div_ceil(sp1));
22721        if ns0 != ns1 {
22722            return Ok(nope("partition-ns", il, pos0));
22723        }
22724        if t0.div_ceil(ns0) != t1.div_ceil(ns1) {
22725            return Ok(nope("partition-per", il, pos0));
22726        }
22727        Ok(true)
22728    }
22729
22730    /// T-ROW fa precheck (the rows kernel supersedes the dcw2 pair-join): every dynamic
22731    /// input of the engine-side dcw decision must hold for EVERY row — stashing is
22732    /// unrecoverable — plus the rows-launcher guards (big-rig ladder, no env split
22733    /// overrides). No partition or capped-SWA clauses: each row derives its OWN geometry.
22734    pub(crate) fn step35_fa_rows_precheck(
22735        &self,
22736        cache: &Cache,
22737        il: usize,
22738        pos0: usize,
22739        t: usize,
22740    ) -> Result<bool, Box<dyn std::error::Error>> {
22741        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22742            return Ok(false);
22743        };
22744        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22745            return Ok(false);
22746        };
22747        let Some(attention) = tp.attention.as_ref() else {
22748            return Ok(false);
22749        };
22750        if !tp.runtime.native_p2p()
22751            || crate::Engine::kv_fp8_on()
22752            || !crate::tp::step_tp_dcw_enabled()?
22753            || !crate::tp::step_tp_qkv_fused_enabled()?
22754            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22755        {
22756            return Ok(false);
22757        }
22758        let geometry = self.step35_geom(il);
22759        let head_dim = geometry.head_dim_k as usize;
22760        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22761            return Ok(false);
22762        }
22763        if crate::fa_sm_count() < 128
22764            || std::env::var("MEMRA_FA_SPLIT").is_ok()
22765            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22766            || std::env::var("MEMRA_FA_SP16").is_ok()
22767            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22768        {
22769            return Ok(false);
22770        }
22771        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22772            return Ok(false);
22773        };
22774        if distributed.staged_len() != pos0 {
22775            return Ok(false);
22776        }
22777        let (_, would_rebase) = distributed.peek_append_ring(t)?;
22778        if would_rebase {
22779            return Ok(false);
22780        }
22781        // Row 0 sees the smallest view: its post-append effective t_kv must clear both
22782        // the dcw floor and the vec-class floor (later rows only grow).
22783        let window = geometry.window.map(|w| w as usize);
22784        let t0 = window.map(|w| (pos0 + 1).min(w)).unwrap_or(pos0 + 1);
22785        if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22786            return Ok(false);
22787        }
22788        Ok(true)
22789    }
22790
22791    /// T-ROW fa join for the verify walk (same-session rows: shared ring/len with
22792    /// len_back = t-1-r). Tables stage once per (layer, rank, ring, t) and live on the
22793    /// owning rank.
22794    pub(crate) fn step35_verify_fa_rows_join(
22795        &self,
22796        e: &Engine,
22797        il: usize,
22798        cache: &Cache,
22799        pos0: usize,
22800        t: usize,
22801    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22802        use cudarc::driver::DevicePtr;
22803        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22804            return Err("fa rows join expects full attention".into());
22805        };
22806        let tp = fa
22807            .step_tp_qkv
22808            .as_ref()
22809            .ok_or("fa rows join lost its resident projections")?;
22810        let geometry = self.step35_geom(il);
22811        let heads = geometry.n_head as usize;
22812        let head_dim = geometry.head_dim_k as usize;
22813        let window = geometry.window.map(|w| w as usize);
22814        let distributed = cache.tp_kv[il]
22815            .as_ref()
22816            .ok_or("fa rows join lost its distributed KV cache")?;
22817        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
22818        // Host mirror of the kernel's big-rig ladder (launcher-guarded identical).
22819        let ladder = |t_kv: usize| -> usize {
22820            if t_kv <= 2048 {
22821                16
22822            } else if t_kv <= 16384 {
22823                64
22824            } else {
22825                128
22826            }
22827        };
22828        let mut max_ns = 1usize;
22829        for r in 0..t {
22830            let t_kv = window
22831                .map(|w| (pos0 + r + 1).min(w))
22832                .unwrap_or(pos0 + r + 1);
22833            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
22834        }
22835        // Rebuild the tiny raw-pointer table from the live distributed cache immediately
22836        // before launch. A process-lifetime map cannot prove allocation generation: CUDA may
22837        // recycle len/base independently of the large K/V rings, making a pointer-key cache
22838        // hit refer to another session (Hermes `11339f5cd3c132a3`).
22839        let ranks = tp.runtime.devices().len();
22840        let mut tables = Vec::with_capacity(ranks);
22841        for rank in 0..ranks {
22842            let engine = tp
22843                .runtime
22844                .rank_engine(rank)
22845                .ok_or("fa rows join lost a rank engine")?;
22846            let rank_cache = distributed
22847                .rank(rank)
22848                .ok_or("fa rows join lost a KV cache rank")?;
22849            let _main = engine.gpu.enter_main()?;
22850            let s = engine.stream();
22851            let (kp, _g0) = rank_cache.k().device_ptr(&s);
22852            let (vp, _g1) = rank_cache.v().device_ptr(&s);
22853            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
22854            let bp = match rank_cache.base_d() {
22855                Some(b) => {
22856                    let (p, _g) = b.device_ptr(&s);
22857                    p
22858                }
22859                None => 0u64,
22860            };
22861            let mut host = Vec::with_capacity(t * 6);
22862            for r in 0..t {
22863                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, (t - 1 - r) as u64]);
22864            }
22865            tables.push(engine.stream().clone_htod(&host)?);
22866        }
22867        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
22868        let ws_index = tp
22869            .runtime
22870            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22871        tp.runtime.decode_v2_fa_rows_join(
22872            ws_index,
22873            e,
22874            &tp.o,
22875            &tabs,
22876            t,
22877            head_dim,
22878            window.unwrap_or(0),
22879            max_ns,
22880            geometry.attention_scale(),
22881            k_tok_bytes,
22882            v_tok_bytes,
22883        )
22884    }
22885
22886    /// Multi-session t-row fa precheck (the batched serving walk): the static doors of
22887    /// the rows kernel plus per-SESSION dynamic checks — every row's own cache must be
22888    /// hydrated, in sync, rebase-free and above both floors.
22889    pub(crate) fn step35_batch_fa_rows_precheck(
22890        &self,
22891        caches: &[&mut Cache],
22892        row_to_cache: impl Fn(usize) -> usize,
22893        positions: &[i32],
22894        il: usize,
22895    ) -> Result<bool, Box<dyn std::error::Error>> {
22896        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22897            return Ok(false);
22898        };
22899        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22900            return Ok(false);
22901        };
22902        let Some(attention) = tp.attention.as_ref() else {
22903            return Ok(false);
22904        };
22905        if !tp.runtime.native_p2p()
22906            || crate::Engine::kv_fp8_on()
22907            || !crate::tp::step_tp_dcw_enabled()?
22908            || !crate::tp::step_tp_qkv_fused_enabled()?
22909            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22910        {
22911            return Ok(false);
22912        }
22913        let geometry = self.step35_geom(il);
22914        let head_dim = geometry.head_dim_k as usize;
22915        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22916            return Ok(false);
22917        }
22918        if crate::fa_sm_count() < 128
22919            || std::env::var("MEMRA_FA_SPLIT").is_ok()
22920            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22921            || std::env::var("MEMRA_FA_SP16").is_ok()
22922            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22923        {
22924            return Ok(false);
22925        }
22926        let window = geometry.window.map(|w| w as usize);
22927        for (r, &pos) in positions.iter().enumerate() {
22928            let cache = &caches[row_to_cache(r)];
22929            let Some(distributed) = cache.tp_kv[il].as_ref() else {
22930                return Ok(false);
22931            };
22932            if distributed.staged_len() != pos as usize {
22933                return Ok(false);
22934            }
22935            if distributed.peek_append_ring(1)?.1 {
22936                return Ok(false);
22937            }
22938            let t0 = window
22939                .map(|w| (pos as usize + 1).min(w))
22940                .unwrap_or(pos as usize + 1);
22941            if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22942                return Ok(false);
22943            }
22944        }
22945        Ok(true)
22946    }
22947
22948    /// FULL t-row attention pass for the VERIFY walk (same-session rows): rope/append +
22949    /// fa + combine + o_proj join in 3 launches/rank/layer. Row r appends at slot
22950    /// len-base+r and one last block advances len by t; the fa rows read len_back =
22951    /// t-1-r. Returns None when the fused-rope class does not hold (the walk keeps the
22952    /// per-column stash flow). Caller has passed `step35_fa_rows_precheck`.
22953    pub(crate) fn step35_verify_rope_fa_pass(
22954        &self,
22955        e: &Engine,
22956        il: usize,
22957        cache: &Cache,
22958        pos0: usize,
22959        t: usize,
22960        stage_pos: bool,
22961    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22962        use cudarc::driver::DevicePtr;
22963        if !crate::tp::fuse_rope_append_on() {
22964            return Ok(None);
22965        }
22966        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22967            return Ok(None);
22968        };
22969        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22970            return Ok(None);
22971        };
22972        let Some(attention) = tp.attention.as_ref() else {
22973            return Ok(None);
22974        };
22975        let geometry = self.step35_geom(il);
22976        let head_dim = geometry.head_dim_k as usize;
22977        if head_dim != 128 {
22978            return Ok(None);
22979        }
22980        let heads = geometry.n_head as usize;
22981        let window = geometry.window.map(|w| w as usize);
22982        let ranks = tp.runtime.devices().len();
22983        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22984            return Ok(None);
22985        };
22986        if distributed.kv_dim_k() != distributed.kv_dim_v() {
22987            return Ok(None);
22988        }
22989        {
22990            let rank0 = distributed.rank(0).ok_or("verify rope pass lost rank 0")?;
22991            if rank0.base_d().is_none()
22992                && distributed.staged_len() + t > distributed.physical_capacity()
22993            {
22994                return Ok(None);
22995            }
22996        }
22997        let mut rope_freqs = Vec::with_capacity(ranks);
22998        for rank in 0..ranks {
22999            let engine = tp
23000                .runtime
23001                .rank_engine(rank)
23002                .ok_or("verify rope pass lost a rank engine")?;
23003            rope_freqs.push(if geometry.rope_factors {
23004                match self
23005                    .step35_aux
23006                    .as_ref()
23007                    .and_then(|aux| aux.rope_freqs(engine))
23008                {
23009                    Some(f) => Some(f),
23010                    None => return Ok(None),
23011                }
23012            } else {
23013                None
23014            });
23015        }
23016        let ladder = |t_kv: usize| -> usize {
23017            if t_kv <= 2048 {
23018                16
23019            } else if t_kv <= 16384 {
23020                64
23021            } else {
23022                128
23023            }
23024        };
23025        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23026        let mut max_ns = 1usize;
23027        let mut positions = Vec::with_capacity(t);
23028        for r in 0..t {
23029            positions.push((pos0 + r) as i32);
23030            let t_kv = window
23031                .map(|w| (pos0 + r + 1).min(w))
23032                .unwrap_or(pos0 + r + 1);
23033            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23034        }
23035        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
23036        let mut tab_keys = vec![0u64; ranks];
23037        for rank in 0..ranks {
23038            let engine = tp
23039                .runtime
23040                .rank_engine(rank)
23041                .ok_or("verify rope pass lost a rank engine")?;
23042            let rank_cache = distributed
23043                .rank(rank)
23044                .ok_or("verify rope pass lost a KV cache rank")?;
23045            let _main = engine.gpu.enter_main()?;
23046            let s = engine.stream();
23047            let (kp, _g0) = rank_cache.k().device_ptr(&s);
23048            let (vp, _g1) = rank_cache.v().device_ptr(&s);
23049            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23050            let bp = match rank_cache.base_d() {
23051                Some(b) => {
23052                    let (p, _g) = b.device_ptr(&s);
23053                    p
23054                }
23055                None => 0u64,
23056            };
23057            tab_keys[rank] = kp
23058                .rotate_left(17)
23059                .wrapping_add(bp)
23060                .wrapping_add((il as u64) << 32)
23061                .wrapping_add(t as u64)
23062                .wrapping_add(1 << 63);
23063            for _r in 0..t {
23064                session_parts[rank].push([kp, vp, lp, bp]);
23065            }
23066        }
23067        let ws_index = tp
23068            .runtime
23069            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23070        tp.runtime
23071            .decode_v2_rope_fa_rows(
23072                ws_index,
23073                e,
23074                &tp.o,
23075                &session_parts,
23076                &tab_keys,
23077                &positions,
23078                stage_pos,
23079                true,
23080                &attention.q_norm,
23081                &attention.k_norm,
23082                &rope_freqs,
23083                t,
23084                head_dim,
23085                geometry.n_rot as usize,
23086                window.unwrap_or(0),
23087                max_ns,
23088                geometry.attention_scale(),
23089                k_tok_bytes,
23090                v_tok_bytes,
23091                self.cfg.rms_eps,
23092                geometry.rope_base,
23093            )
23094            .map(Some)
23095    }
23096
23097    /// FULL t-row attention pass for the batched walk (rope/append + fa + combine +
23098    /// o_proj join, 3 launches/rank/layer): returns None when the fused-rope class does
23099    /// not hold — the caller falls back to the per-row stash flow. The caller has
23100    /// already passed `step35_batch_fa_rows_precheck`.
23101    #[allow(clippy::too_many_arguments)]
23102    pub(crate) fn step35_batch_rope_fa_pass(
23103        &self,
23104        e: &Engine,
23105        il: usize,
23106        caches: &[&mut Cache],
23107        row_to_cache: impl Fn(usize) -> usize,
23108        positions: &[i32],
23109        t: usize,
23110        stage_pos: bool,
23111    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
23112        use cudarc::driver::DevicePtr;
23113        if !crate::tp::fuse_rope_append_on() {
23114            return Ok(None);
23115        }
23116        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23117            return Ok(None);
23118        };
23119        let Some(tp) = fa.step_tp_qkv.as_ref() else {
23120            return Ok(None);
23121        };
23122        let Some(attention) = tp.attention.as_ref() else {
23123            return Ok(None);
23124        };
23125        let geometry = self.step35_geom(il);
23126        let head_dim = geometry.head_dim_k as usize;
23127        if head_dim != 128 {
23128            return Ok(None);
23129        }
23130        let heads = geometry.n_head as usize;
23131        let window = geometry.window.map(|w| w as usize);
23132        let ranks = tp.runtime.devices().len();
23133        // The rows kernels never arm base_d; refuse once a ring could have rebased
23134        // without an armed base (the table would read base=0 after a real rebase).
23135        for r in 0..t {
23136            let cache = &caches[row_to_cache(r)];
23137            let Some(distributed) = cache.tp_kv[il].as_ref() else {
23138                return Ok(None);
23139            };
23140            if distributed.kv_dim_k() != distributed.kv_dim_v() {
23141                return Ok(None);
23142            }
23143            let rank0 = distributed.rank(0).ok_or("rope fa pass lost rank 0")?;
23144            if rank0.base_d().is_none()
23145                && distributed.staged_len() + t > distributed.physical_capacity()
23146            {
23147                return Ok(None);
23148            }
23149        }
23150        let mut rope_freqs = Vec::with_capacity(ranks);
23151        for rank in 0..ranks {
23152            let engine = tp
23153                .runtime
23154                .rank_engine(rank)
23155                .ok_or("rope fa pass lost a rank engine")?;
23156            rope_freqs.push(if geometry.rope_factors {
23157                match self
23158                    .step35_aux
23159                    .as_ref()
23160                    .and_then(|aux| aux.rope_freqs(engine))
23161                {
23162                    Some(f) => Some(f),
23163                    None => return Ok(None),
23164                }
23165            } else {
23166                None
23167            });
23168        }
23169        let ladder = |t_kv: usize| -> usize {
23170            if t_kv <= 2048 {
23171                16
23172            } else if t_kv <= 16384 {
23173                64
23174            } else {
23175                128
23176            }
23177        };
23178        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23179        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
23180        let mut tab_keys = vec![0u64; ranks];
23181        for (r, &pos) in positions.iter().enumerate().take(t) {
23182            let cache = &caches[row_to_cache(r)];
23183            let distributed = cache.tp_kv[il]
23184                .as_ref()
23185                .ok_or("rope fa pass lost a distributed KV cache")?;
23186            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23187            let t_kv = window
23188                .map(|w| (pos as usize + 1).min(w))
23189                .unwrap_or(pos as usize + 1);
23190            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23191            for rank in 0..ranks {
23192                let engine = tp
23193                    .runtime
23194                    .rank_engine(rank)
23195                    .ok_or("rope fa pass lost a rank engine")?;
23196                let rank_cache = distributed
23197                    .rank(rank)
23198                    .ok_or("rope fa pass lost a KV cache rank")?;
23199                let _main = engine.gpu.enter_main()?;
23200                let s = engine.stream();
23201                let (kp, _g0) = rank_cache.k().device_ptr(&s);
23202                let (vp, _g1) = rank_cache.v().device_ptr(&s);
23203                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23204                let bp = match rank_cache.base_d() {
23205                    Some(b) => {
23206                        let (p, _g) = b.device_ptr(&s);
23207                        p
23208                    }
23209                    None => 0u64,
23210                };
23211                tab_keys[rank] = tab_keys[rank]
23212                    .rotate_left(9)
23213                    .wrapping_add(kp)
23214                    .wrapping_add(bp)
23215                    .wrapping_add(il as u64);
23216                session_parts[rank].push([kp, vp, lp, bp]);
23217            }
23218        }
23219        let ws_index = tp
23220            .runtime
23221            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23222        tp.runtime
23223            .decode_v2_rope_fa_rows(
23224                ws_index,
23225                e,
23226                &tp.o,
23227                &session_parts,
23228                &tab_keys,
23229                positions,
23230                stage_pos,
23231                false,
23232                &attention.q_norm,
23233                &attention.k_norm,
23234                &rope_freqs,
23235                t,
23236                head_dim,
23237                geometry.n_rot as usize,
23238                window.unwrap_or(0),
23239                max_ns,
23240                geometry.attention_scale(),
23241                k_tok_bytes,
23242                v_tok_bytes,
23243                self.cfg.rms_eps,
23244                geometry.rope_base,
23245            )
23246            .map(Some)
23247    }
23248
23249    /// Multi-session t-row fa join (batched serving): per-row table entries point at
23250    /// each row's OWN session rings/counters (len_back = 0 — every session appended
23251    /// exactly its one row). Tables stage once per (layer, rank, session-set, t).
23252    #[allow(clippy::too_many_arguments)]
23253    pub(crate) fn step35_batch_fa_rows_join(
23254        &self,
23255        e: &Engine,
23256        il: usize,
23257        caches: &[&mut Cache],
23258        row_to_cache: impl Fn(usize) -> usize,
23259        positions: &[i32],
23260        t: usize,
23261    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23262        use cudarc::driver::DevicePtr;
23263        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23264            return Err("batch fa rows join expects full attention".into());
23265        };
23266        let tp = fa
23267            .step_tp_qkv
23268            .as_ref()
23269            .ok_or("batch fa rows join lost its resident projections")?;
23270        let geometry = self.step35_geom(il);
23271        let heads = geometry.n_head as usize;
23272        let head_dim = geometry.head_dim_k as usize;
23273        let window = geometry.window.map(|w| w as usize);
23274        let ladder = |t_kv: usize| -> usize {
23275            if t_kv <= 2048 {
23276                16
23277            } else if t_kv <= 16384 {
23278                64
23279            } else {
23280                128
23281            }
23282        };
23283        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23284        for (r, &pos) in positions.iter().enumerate() {
23285            let cache = &caches[row_to_cache(r)];
23286            let distributed = cache.tp_kv[il]
23287                .as_ref()
23288                .ok_or("batch fa rows join lost a distributed KV cache")?;
23289            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23290            let t_kv = window
23291                .map(|w| (pos as usize + 1).min(w))
23292                .unwrap_or(pos as usize + 1);
23293            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23294        }
23295        // Multi-session tables also rebuild from every live K/V/len/base tuple. Keeping a
23296        // process-lifetime raw-pointer cache here omitted V and len identity and had no
23297        // allocation generation, so allocator reuse could bind one request to another.
23298        let ranks = tp.runtime.devices().len();
23299        let mut tables = Vec::with_capacity(ranks);
23300        for rank in 0..ranks {
23301            let engine = tp
23302                .runtime
23303                .rank_engine(rank)
23304                .ok_or("batch fa rows join lost a rank engine")?;
23305            let _main = engine.gpu.enter_main()?;
23306            let s = engine.stream();
23307            let mut host = Vec::with_capacity(t * 6);
23308            for r in 0..t {
23309                let cache = &caches[row_to_cache(r)];
23310                let distributed = cache.tp_kv[il]
23311                    .as_ref()
23312                    .ok_or("batch fa rows join lost a distributed KV cache")?;
23313                let rank_cache = distributed
23314                    .rank(rank)
23315                    .ok_or("batch fa rows join lost a KV cache rank")?;
23316                let (kp, _g0) = rank_cache.k().device_ptr(&s);
23317                let (vp, _g1) = rank_cache.v().device_ptr(&s);
23318                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23319                let bp = match rank_cache.base_d() {
23320                    Some(b) => {
23321                        let (p, _g) = b.device_ptr(&s);
23322                        p
23323                    }
23324                    None => 0u64,
23325                };
23326                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, 0u64]);
23327            }
23328            tables.push(engine.stream().clone_htod(&host)?);
23329        }
23330        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
23331        let ws_index = tp
23332            .runtime
23333            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23334        tp.runtime.decode_v2_fa_rows_join(
23335            ws_index,
23336            e,
23337            &tp.o,
23338            &tabs,
23339            t,
23340            head_dim,
23341            window.unwrap_or(0),
23342            max_ns,
23343            geometry.attention_scale(),
23344            k_tok_bytes,
23345            v_tok_bytes,
23346        )
23347    }
23348
23349    /// MEMRA_SPEC_FA2 join for the verify walk: both columns stashed; one shared-KV T=2
23350    /// fa per rank + the weight-amortized o_proj join produce the [2, o_out] `mixed`
23351    /// slab on `e`.
23352    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
23353    pub(crate) fn step35_verify_spec_fa2_join(
23354        &self,
23355        e: &Engine,
23356        il: usize,
23357        cache: &Cache,
23358        pos0: usize,
23359    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23360        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23361            return Err("spec fa2 join expects full attention".into());
23362        };
23363        let tp = fa
23364            .step_tp_qkv
23365            .as_ref()
23366            .ok_or("spec fa2 join lost its resident projections")?;
23367        let geometry = self.step35_geom(il);
23368        let heads = geometry.n_head as usize;
23369        let head_dim = geometry.head_dim_k as usize;
23370        let window = geometry.window.map(|w| w as usize);
23371        // POST-append view of the second row (kernel T1 = len - lstart with len =
23372        // pos0 + 2): sp/ns derive from it, and the precheck proved row0 shares them.
23373        let bucket = window.map(|w| (pos0 + 2).min(w)).unwrap_or(pos0 + 2);
23374        let distributed = cache.tp_kv[il]
23375            .as_ref()
23376            .ok_or("spec fa2 join lost its distributed KV cache")?;
23377        let ws_index = tp
23378            .runtime
23379            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23380        tp.runtime.decode_v2_spec_fa2_join(
23381            ws_index,
23382            e,
23383            &tp.o,
23384            distributed,
23385            head_dim,
23386            window.unwrap_or(0),
23387            bucket,
23388            geometry.attention_scale(),
23389        )
23390    }
23391
23392    pub(crate) fn step35_tp_decode_attn_resident_v2(
23393        &self,
23394        e: &Engine,
23395        fa: &FullAttnLayer,
23396        il: usize,
23397        h: &CudaSlice<f32>,
23398        pos_d: &CudaSlice<i32>,
23399        cache: &mut Cache,
23400    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23401        let tp = fa
23402            .step_tp_qkv
23403            .as_ref()
23404            .ok_or("Step TP decode lost its resident projections")?;
23405        let attention = tp
23406            .attention
23407            .as_ref()
23408            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
23409        if !tp.runtime.native_p2p() {
23410            return Err("rank-local Step attention requires native P2P".into());
23411        }
23412        if crate::Engine::kv_fp8_on() {
23413            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
23414        }
23415
23416        let geometry = self.cfg.full_attention_geometry_at(il as u32);
23417        let window = geometry.window.map(|window| window as usize);
23418        let ranks = tp.runtime.devices().len();
23419        let head_dim = geometry.head_dim_k as usize;
23420        let heads = geometry.n_head as usize;
23421        let kv_heads = geometry.n_head_kv as usize;
23422        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
23423            return Err(format!(
23424                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
23425            )
23426            .into());
23427        }
23428        let local_heads = heads / ranks;
23429        let local_kv_heads = kv_heads / ranks;
23430        let max_ctx = cache.max_ctx;
23431
23432        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
23433
23434        let base_len = cache.kv[il]
23435            .as_ref()
23436            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
23437            .len;
23438        {
23439            let distributed = cache.tp_kv[il]
23440                .as_ref()
23441                .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
23442            if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
23443                return Err(format!(
23444                    "Step TP layer {il} cache lengths diverged before decode: \
23445                     local={base_len} distributed={}/{}",
23446                    distributed.committed_len(),
23447                    distributed.staged_len()
23448                )
23449                .into());
23450            }
23451        }
23452        if pos_d.len() != 1 {
23453            return Err(format!(
23454                "rank-local Step decode requires one position, got {}",
23455                pos_d.len()
23456            )
23457            .into());
23458        }
23459
23460        let decode_input = attention
23461            .decode_input
23462            .as_ref()
23463            .ok_or("Step TP decode v2 requires the replicated decode input")?;
23464        let mut decode_input = decode_input
23465            .lock()
23466            .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
23467
23468        let has_gate = fa.attn_gate.is_some();
23469        // Gate: with per-rank shards loaded (fused door), the fused QKV+gate kernel computes
23470        // it rank-locally and the model-engine matmul (and its staging copies) disappears.
23471        // Otherwise it queues on e's stream BEFORE decode_v2_input_qkv records the entry
23472        // event, so the rank-stream reads of the staged gate are ordered without a host sync.
23473        let use_gate_shards = has_gate
23474            && (attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some())
23475            && crate::tp::step_tp_qkv_fused_enabled()?;
23476        let gate_raw = if !has_gate || use_gate_shards {
23477            None
23478        } else {
23479            let gate_weight = fa
23480                .attn_gate
23481                .as_ref()
23482                .ok_or("step35 layer is missing attn_gate.weight")?;
23483            let gate_raw = e.matmul(gate_weight, h, 1)?;
23484            if gate_raw.len() != heads {
23485                return Err(format!(
23486                    "Step TP layer {il} gate output {} != {heads}",
23487                    gate_raw.len()
23488                )
23489                .into());
23490            }
23491            Some(gate_raw)
23492        };
23493
23494        let ws_index = tp
23495            .runtime
23496            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23497        let mut ws_guard = tp
23498            .runtime
23499            .decode_v2_workspace()
23500            .lock()
23501            .map_err(|_| "Step TP decode v2 workspace lock is poisoned")?;
23502        let ws = ws_guard
23503            .get_mut(ws_index)
23504            .ok_or("Step TP decode v2 workspace missing after ensure")?;
23505
23506        let mut rope_freqs = Vec::with_capacity(ranks);
23507        for rank in 0..ranks {
23508            let engine = tp
23509                .runtime
23510                .rank_engine(rank)
23511                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23512            rope_freqs.push(if geometry.rope_factors {
23513                self.step35_aux
23514                    .as_ref()
23515                    .and_then(|aux| aux.rope_freqs(engine))
23516            } else {
23517                None
23518            });
23519        }
23520        // DCW arm (MEMRA_STEP_TP_DCW=1): device-counter append + counter-derived fa — the
23521        // exact captured-child content, run eagerly. bucket = effective t_kv keeps the fa
23522        // split geometry identical to the kvmod path (one-partition law) -> bit-identical.
23523        // Rebase tokens and sub-vec-floor contexts take the host-row path below.
23524        // (Eligibility computed BEFORE input_qkv so FUSION #1 can defer the norm+rope into
23525        // the fused rope+append+inc launch on dcw tokens.)
23526        let staged_next = base_len + 1;
23527        let t_kv_eff = window
23528            .map(|window| staged_next.min(window))
23529            .unwrap_or(staged_next);
23530        let dcw = crate::tp::step_tp_dcw_enabled()?
23531            && (use_gate_shards || (!has_gate && crate::tp::step_tp_qkv_fused_enabled()?))
23532            && t_kv_eff >= 96
23533            && {
23534                let (write_row, would_rebase) = cache.tp_kv[il]
23535                    .as_ref()
23536                    .expect("distributed cache checked above")
23537                    .peek_append_ring(1)?;
23538                if !would_rebase {
23539                    // Arm the base mirrors on first use: base = logical staged - physical row.
23540                    let base = (base_len - write_row) as i32;
23541                    let distributed = cache.tp_kv[il]
23542                        .as_mut()
23543                        .expect("distributed cache checked above");
23544                    for rank in 0..ranks {
23545                        let engine = tp.runtime.rank_engine(rank).ok_or_else(|| {
23546                            format!("Step TP layer {il} has no engine for rank {rank}")
23547                        })?;
23548                        let _main = engine.gpu.enter_main()?;
23549                        let rank_cache = distributed.rank_mut(rank).ok_or_else(|| {
23550                            format!("Step TP layer {il} has no KV cache rank {rank}")
23551                        })?;
23552                        if rank_cache.base_d().is_none() {
23553                            rank_cache.arm_base_d(engine.htod_i32(&[base])?);
23554                        }
23555                    }
23556                }
23557                !would_rebase
23558            };
23559        let fuse_rope = dcw
23560            && crate::tp::fuse_rope_append_on()
23561            && head_dim == 128
23562            && cache.tp_kv[il]
23563                .as_ref()
23564                .map(|d| d.kv_dim_k() == d.kv_dim_v() && d.kv_dim_k() == local_kv_heads * head_dim)
23565                .unwrap_or(false);
23566
23567        let tcol_col = crate::tp::take_verify_tcol();
23568        // MEMRA_SPEC_FA2 defer: the verify walk armed this column for the shared-KV T=2
23569        // attention. On dcw tokens the per-rank pass still norms/ropes/APPENDS (cache
23570        // state must advance per column) but skips the fa+gate launch; post-rope q and
23571        // gate rows are stashed instead, and ONE fa_decode_dcw2 per rank joins both
23572        // columns after the second append. Non-dcw tokens ignore the defer (the fa runs
23573        // normally and the walk consumes the real output — stash flag stays unset).
23574        let fa2_col = crate::tp::take_spec_fa2_defer();
23575        tp.runtime.decode_v2_input_qkv(
23576            ws,
23577            e,
23578            h,
23579            pos_d,
23580            gate_raw.as_ref(),
23581            if !use_gate_shards {
23582                None
23583            } else if let Some(shards) = attention.gate_shards.as_deref() {
23584                Some(crate::tp::StepTpGateShards::F32(shards))
23585            } else {
23586                attention
23587                    .gate_shards_bf16
23588                    .as_deref()
23589                    .map(crate::tp::StepTpGateShards::Bf16)
23590            },
23591            &mut decode_input,
23592            &tp.q,
23593            &tp.k,
23594            &tp.v,
23595            &attention.q_norm,
23596            &attention.k_norm,
23597            head_dim,
23598            geometry.n_rot as usize,
23599            geometry.rope_base,
23600            &rope_freqs,
23601            self.cfg.rms_eps,
23602            has_gate,
23603            fuse_rope,
23604            tcol_col,
23605        )?;
23606
23607        let transaction = cache.tp_kv[il]
23608            .as_mut()
23609            .expect("distributed cache checked above")
23610            .begin_transaction()?;
23611        let append_result = tp.runtime.append_tp_kv_transaction_inner(
23612            cache.tp_kv[il]
23613                .as_mut()
23614                .expect("distributed cache checked above"),
23615            transaction,
23616            &ws.k,
23617            &ws.v_raw,
23618            1,
23619            dcw,
23620        );
23621        if let Err(error) = append_result {
23622            let _ = tp.runtime.rollback_tp_kv_transaction(
23623                cache.tp_kv[il]
23624                    .as_mut()
23625                    .expect("distributed cache checked above"),
23626                transaction,
23627            );
23628            return Err(error);
23629        }
23630
23631        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23632            // Copy the view scalars out and DROP the shared borrow — the merged dcw arm
23633            // reborrows the cache mutably per rank.
23634            let (staged_len, physical, k_tok_bytes_c, v_tok_bytes_c, capacity) = {
23635                let distributed = cache.tp_kv[il]
23636                    .as_ref()
23637                    .expect("distributed cache checked above");
23638                let staged_len = distributed.staged_len();
23639                let view_start = window
23640                    .map(|window| staged_len.saturating_sub(window))
23641                    .unwrap_or(0);
23642                (
23643                    staged_len,
23644                    distributed.physical_range(view_start, staged_len)?,
23645                    distributed.k_tok_bytes(),
23646                    distributed.v_tok_bytes(),
23647                    distributed.physical_capacity(),
23648                )
23649            };
23650            let view_start = window
23651                .map(|window| staged_len.saturating_sub(window))
23652                .unwrap_or(0);
23653            let t_kv = staged_len - view_start;
23654            for rank in 0..ranks {
23655                let engine = tp
23656                    .runtime
23657                    .rank_engine(rank)
23658                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23659                let _main = engine.gpu.enter_main()?;
23660                if dcw {
23661                    // MERGED per-rank pass (the capture unit): append + inc + fa + gate on ONE
23662                    // stream visit. distributed is borrowed shared here; the planes need mut —
23663                    // reborrow through the cache Option (the closure holds cache mutably).
23664                    {
23665                        let distributed_mut = cache.tp_kv[il]
23666                            .as_mut()
23667                            .expect("distributed cache checked above");
23668                        let (kv_dim_k, kv_dim_v) =
23669                            (distributed_mut.kv_dim_k(), distributed_mut.kv_dim_v());
23670                        let (k_tok_bytes, v_tok_bytes) =
23671                            (distributed_mut.k_tok_bytes(), distributed_mut.v_tok_bytes());
23672                        let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23673                            format!("Step TP layer {il} has no KV cache rank {rank}")
23674                        })?;
23675                        let (k_plane, v_plane, len_d, base_d) =
23676                            rank_cache.planes_and_counters_mut();
23677                        if fuse_rope {
23678                            // FUSION #1: norms + rope (deferred out of input_qkv) + append
23679                            // + last-block len inc in ONE launch. Bit-identical bodies.
23680                            let same_dev = engine.ctx().ordinal() == e.ctx().ordinal();
23681                            let crate::tp::StepTpDecodeV2Ws {
23682                                q_raw,
23683                                k_raw,
23684                                v_raw,
23685                                q,
23686                                k,
23687                                pos,
23688                                pos_stage,
23689                                fuse_ctr,
23690                                ..
23691                            } = &mut *ws;
23692                            // Same-device rank: the staged-copy elision leaves pos[rank]
23693                            // stale — read the e-context pos stage directly (mirrors the
23694                            // rope elision in input_qkv_rank).
23695                            let pos_ref: &CudaSlice<i32> = if same_dev {
23696                                pos_stage
23697                                    .as_ref()
23698                                    .ok_or("step TP decode v2 pos stage not armed")?
23699                            } else {
23700                                &pos[rank]
23701                            };
23702                            engine.qk_norm_rope_append_inc_dcw(
23703                                &q_raw[rank],
23704                                &k_raw[rank],
23705                                &v_raw[rank],
23706                                &attention.q_norm[rank],
23707                                &attention.k_norm[rank],
23708                                &mut q[rank],
23709                                &mut k[rank],
23710                                pos_ref,
23711                                k_plane,
23712                                v_plane,
23713                                len_d,
23714                                base_d,
23715                                &mut fuse_ctr[rank],
23716                                kv_dim_k,
23717                                kv_dim_v,
23718                                k_tok_bytes,
23719                                v_tok_bytes,
23720                                head_dim,
23721                                geometry.n_rot as usize,
23722                                local_heads,
23723                                local_kv_heads,
23724                                self.cfg.rms_eps,
23725                                geometry.rope_base,
23726                                1.0,
23727                                rope_freqs[rank],
23728                            )?;
23729                        } else {
23730                            engine.append_kv_quantized_dcw(
23731                                &ws.k[rank],
23732                                &ws.v_raw[rank],
23733                                k_plane,
23734                                v_plane,
23735                                len_d,
23736                                base_d,
23737                                kv_dim_k,
23738                                kv_dim_v,
23739                                k_tok_bytes,
23740                                v_tok_bytes,
23741                            )?;
23742                        }
23743                        if !fuse_rope {
23744                            let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23745                                format!("Step TP layer {il} has no KV cache rank {rank}")
23746                            })?;
23747                            engine.inc_i32(rank_cache.len_d_mut())?;
23748                        }
23749                    }
23750                    let distributed = cache.tp_kv[il]
23751                        .as_ref()
23752                        .expect("distributed cache checked above");
23753                    let rank_cache = distributed
23754                        .rank(rank)
23755                        .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23756                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * k_tok_bytes_c);
23757                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * v_tok_bytes_c);
23758                    if fa2_col.is_some() {
23759                        // SPEC_FA2 defer: append landed above; the fa for this column
23760                        // runs in the T=2 joined launch after the pair's second append.
23761                        continue;
23762                    }
23763                    {
23764                        // FUSION #2d: combine + head gate in the dcw tail — `gated` receives
23765                        // the gated output directly (bit-identical; one launch saved).
23766                        let crate::tp::StepTpDecodeV2Ws { q, gate, gated, .. } = &mut *ws;
23767                        engine.fa_decode_dcw(
23768                            &q[rank],
23769                            &k_ring,
23770                            &v_ring,
23771                            &mut gated[rank],
23772                            head_dim,
23773                            local_heads,
23774                            local_kv_heads,
23775                            rank_cache.len_d(),
23776                            rank_cache.base_d(),
23777                            window.unwrap_or(0),
23778                            t_kv,
23779                            geometry.attention_scale(),
23780                            k_tok_bytes_c,
23781                            v_tok_bytes_c,
23782                            has_gate.then_some(&gate[rank]),
23783                        )?;
23784                    }
23785                    continue;
23786                }
23787                let distributed = cache.tp_kv[il]
23788                    .as_ref()
23789                    .expect("distributed cache checked above");
23790                let rank_cache = distributed
23791                    .rank(rank)
23792                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23793                let k_view = engine.view_u8_range(
23794                    rank_cache.k(),
23795                    physical.start * k_tok_bytes_c,
23796                    physical.end * k_tok_bytes_c,
23797                );
23798                let v_view = engine.view_u8_range(
23799                    rank_cache.v(),
23800                    physical.start * v_tok_bytes_c,
23801                    physical.end * v_tok_bytes_c,
23802                );
23803                if has_gate {
23804                    engine.fa_decode_kvmod(
23805                        &ws.q[rank],
23806                        &k_view,
23807                        &v_view,
23808                        &mut ws.attn_out[rank],
23809                        head_dim,
23810                        local_heads,
23811                        local_kv_heads,
23812                        t_kv,
23813                        geometry.attention_scale(),
23814                        k_tok_bytes_c,
23815                        v_tok_bytes_c,
23816                        false,
23817                    )?;
23818                    engine.attn_head_gate(
23819                        &ws.attn_out[rank],
23820                        &ws.gate[rank],
23821                        &mut ws.gated[rank],
23822                        None,
23823                        head_dim,
23824                        local_heads,
23825                        1,
23826                    )?;
23827                } else {
23828                    engine.fa_decode_kvmod(
23829                        &ws.q[rank],
23830                        &k_view,
23831                        &v_view,
23832                        &mut ws.gated[rank],
23833                        head_dim,
23834                        local_heads,
23835                        local_kv_heads,
23836                        t_kv,
23837                        geometry.attention_scale(),
23838                        k_tok_bytes_c,
23839                        v_tok_bytes_c,
23840                        false,
23841                    )?;
23842                }
23843            }
23844
23845            // MEMRA_TCOL_OPROJ defer: the verify driver armed a column — stash this
23846            // column's `gated` rows and skip the per-column finish choreography entirely
23847            // (the batched b4_tcol + join runs after every column). The returned buffer
23848            // is UNWRITTEN in that mode (oproj-tail precedent); the driver reads the
23849            // stashed flag, never this buffer. Ineligible configs fall back to the
23850            // normal finish and the driver consumes the real `mixed` per column.
23851            let output = if let Some(col) = fa2_col.filter(|_| dcw) {
23852                // SPEC_FA2 stash: q + gate rows to the fa2 slabs; fa, o_proj and the
23853                // finish all run in the joined pass. Returned buffer is UNWRITTEN
23854                // (oproj-defer precedent — the walk reads the stash flag, never this).
23855                tp.runtime.decode_v2_stash_fa2(ws, e, col)?;
23856                crate::tp::set_spec_fa2_stashed();
23857                e.uninit(ws.o_out)?
23858            } else if let Some(col) = crate::tp::take_tcol_oproj_defer() {
23859                if tp.runtime.decode_v2_oproj_tcol_eligible(ws, &tp.o) {
23860                    tp.runtime.decode_v2_stash_gated(ws, e, col)?;
23861                    crate::tp::set_tcol_oproj_stashed();
23862                    e.uninit(ws.o_out)?
23863                } else {
23864                    tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23865                }
23866            } else {
23867                tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23868            };
23869
23870            // Local shadow append: reads ws.k_shadow/ws.v_shadow on e's stream, which
23871            // decode_v2_finish ordered behind the root event. Same math and cache state
23872            // transitions as v1.
23873            let local = cache.kv[il]
23874                .as_mut()
23875                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
23876            if local.len != base_len || base_len + 1 > max_ctx {
23877                return Err(format!(
23878                    "Step TP layer {il} local cache changed during decode: \
23879                     len={} base={base_len} max={max_ctx}",
23880                    local.len
23881                )
23882                .into());
23883            }
23884            if crate::tp::no_local_shadow_on() {
23885                // Lengths advance, contents stay stale (graph-door precedent: decode reads
23886                // only the distributed TP caches; local contents feed spec/MTP scratch).
23887                local.len = base_len + 1;
23888                // MEMRA_LEN_MIRROR_LAZY=1: skip the 4B in-stream htod — nothing reads the
23889                // LOCAL device mirror in TP decode (the dcw fa reads the RANK counters),
23890                // and each tiny copy costs a compute->copy engine turnaround mid-layer.
23891                if !crate::tp::len_mirror_lazy_on() {
23892                    e.set_i32_one(&mut local.len_d, local.len as i32)?;
23893                }
23894            } else {
23895                let retain_from = window
23896                    .map(|window| {
23897                        let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
23898                        let rollback_retain =
23899                            base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
23900                        staged_retain.min(rollback_retain)
23901                    })
23902                    .unwrap_or(0);
23903                let write_row = e.prepare_kv_append(local, retain_from, 1)?;
23904                e.append_kv_quantized(
23905                    &ws.k_shadow,
23906                    &ws.v_shadow,
23907                    &mut local.k,
23908                    &mut local.v,
23909                    write_row,
23910                    local.kv_dim_k,
23911                    local.kv_dim_v,
23912                    local.k_tok_bytes,
23913                    local.v_tok_bytes,
23914                    false,
23915                )?;
23916                local.len = base_len + 1;
23917                e.set_i32_one(&mut local.len_d, local.len as i32)?;
23918            }
23919            Ok(output)
23920        })();
23921
23922        let output = match staged {
23923            Ok(output) => output,
23924            Err(error) => {
23925                let _ = tp.runtime.rollback_tp_kv_transaction(
23926                    cache.tp_kv[il]
23927                        .as_mut()
23928                        .expect("distributed cache checked above"),
23929                    transaction,
23930                );
23931                if let Some(local) = cache.kv[il].as_mut() {
23932                    local.len = base_len;
23933                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
23934                }
23935                return Err(error);
23936            }
23937        };
23938        // MEMRA_LEN_MIRROR_LAZY under FUSE_ROPE_APPEND: the fused append atomicInc owns
23939        // the rank counters (same value as the absolute re-set on full accept), so commit
23940        // host bookkeeping only — kills two 4B in-stream htods per layer. Non-fused dcw
23941        // keeps the absolute set (its appends do NOT inc).
23942        let lazy_commit = fuse_rope && crate::tp::len_mirror_lazy_on();
23943        if lazy_commit {
23944            if let Err(error) = tp.runtime.commit_tp_kv_transaction_external(
23945                cache.tp_kv[il]
23946                    .as_mut()
23947                    .expect("distributed cache checked above"),
23948                transaction,
23949                1,
23950            ) {
23951                let _ = tp.runtime.rollback_tp_kv_transaction(
23952                    cache.tp_kv[il]
23953                        .as_mut()
23954                        .expect("distributed cache checked above"),
23955                    transaction,
23956                );
23957                let local = cache.kv[il].as_mut().expect("local cache checked above");
23958                local.len = base_len;
23959                e.set_i32_one(&mut local.len_d, base_len as i32)?;
23960                return Err(error);
23961            }
23962        } else if let Err(error) = tp.runtime.commit_tp_kv_transaction(
23963            cache.tp_kv[il]
23964                .as_mut()
23965                .expect("distributed cache checked above"),
23966            transaction,
23967            1,
23968        ) {
23969            let _ = tp.runtime.rollback_tp_kv_transaction(
23970                cache.tp_kv[il]
23971                    .as_mut()
23972                    .expect("distributed cache checked above"),
23973                transaction,
23974            );
23975            let local = cache.kv[il].as_mut().expect("local cache checked above");
23976            local.len = base_len;
23977            e.set_i32_one(&mut local.len_d, base_len as i32)?;
23978            return Err(error);
23979        }
23980
23981        let committed = cache.tp_kv[il]
23982            .as_ref()
23983            .expect("distributed cache checked above")
23984            .committed_len();
23985        let local_len = cache.kv[il]
23986            .as_ref()
23987            .expect("local cache checked above")
23988            .len;
23989        if committed != local_len {
23990            return Err(format!(
23991                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
23992            )
23993            .into());
23994        }
23995        static V2_LOGGED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
23996        if !V2_LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) {
23997            eprintln!(
23998                "[step-tp-attn-v2] execute layer={} devices={:?} tokens=1 driver=v2 \
23999                 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
24000                 kv_cache_distributed=true kv_cache_hydrated={hydrated} \
24001                 attention_tensor_parallel=true attention_scope={} \
24002                 input_path=root-device-replicated gate={} gate_tensor_parallel={} \
24003                 gate_shards={} o_tensor_parallel=true o_reduce=root-device \
24004                 local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
24005                 bulk_p2p={} workspace=persistent ordering=evented output=e-device \
24006                 performance_claim=false (logged once; every decode layer runs this driver)",
24007                tp.layer,
24008                tp.devices,
24009                if window.is_some() {
24010                    "rank-local-swa-ring"
24011                } else {
24012                    "rank-local-global"
24013                },
24014                has_gate,
24015                use_gate_shards,
24016                if use_gate_shards {
24017                    "device-staged"
24018                } else if has_gate {
24019                    "root-staged"
24020                } else {
24021                    "none"
24022                },
24023                tp.runtime.transport_label(),
24024                tp.runtime.bulk_p2p(),
24025            );
24026        }
24027        Ok(output)
24028    }
24029
24030    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
24031    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
24032    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
24033    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
24034    /// requiring `attn_gate`).
24035    ///
24036    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
24037    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
24038    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
24039    #[allow(clippy::too_many_arguments)]
24040    pub(crate) fn step35_decode_attn(
24041        &self,
24042        e: &Engine,
24043        fa: &FullAttnLayer,
24044        il: usize,
24045        h: &CudaSlice<f32>,
24046        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
24047        pos_d: &CudaSlice<i32>,
24048        cache: &mut Cache,
24049    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24050        if fa
24051            .step_tp_qkv
24052            .as_ref()
24053            .is_some_and(|tp| tp.attention.is_some())
24054        {
24055            if pre_q.is_some() {
24056                return Err(
24057                    "rank-local Step attention preserves BF16 activations and refuses the q8_1 \
24058                     pre-quantized decode path"
24059                        .into(),
24060                );
24061            }
24062            return self.step35_tp_decode_attn_resident(e, fa, il, h, pos_d, cache);
24063        }
24064
24065        let geometry = self.step35_geom(il);
24066        let hd = geometry.head_dim_k as usize;
24067        let nkv = geometry.n_head_kv as usize;
24068        let nh = geometry.n_head as usize;
24069        let rbase = geometry.rope_base;
24070        let scale = geometry.attention_scale();
24071        let swa = geometry.window.is_some();
24072        let eps = self.cfg.rms_eps;
24073        let win = geometry.window.unwrap_or(0) as usize;
24074        let n_rot = geometry.n_rot as usize;
24075        let n_embd = self.cfg.n_embd as usize;
24076        let gw = fa
24077            .attn_gate
24078            .as_ref()
24079            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
24080
24081        let tp_qkv = if fa.step_tp_qkv.is_some() {
24082            if pre_q.is_some() {
24083                return Err(
24084                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
24085                     pre-quantized decode path"
24086                        .into(),
24087                );
24088            }
24089            self.full_attn_tp_qkv(e, fa, h, 1)?
24090        } else {
24091            None
24092        };
24093
24094        let (q0, k0, v0, gt) = match tp_qkv {
24095            Some(mut g3) => {
24096                let v = g3.pop().unwrap();
24097                let k = g3.pop().unwrap();
24098                let q = g3.pop().unwrap();
24099                let gt = e.matmul(gw, h, 1)?;
24100                (q, k, v, gt)
24101            }
24102            None => match pre_q {
24103                Some((hq, hdq)) => {
24104                    debug_assert!(
24105                        e.uses_q8_1_fast(gw),
24106                        "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
24107                         (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
24108                    );
24109                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
24110                        Some(t3) => t3,
24111                        None => (
24112                            e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
24113                            e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
24114                            e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
24115                        ),
24116                    };
24117                    let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
24118                    (a, b, c, gt)
24119                }
24120                None => {
24121                    if e.uses_q8_1_fast(&fa.wq)
24122                        && e.uses_q8_1_fast(&fa.wk)
24123                        && e.uses_q8_1_fast(&fa.wv)
24124                        && e.uses_q8_1_fast(gw)
24125                    {
24126                        let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
24127                        let (a, b, c) =
24128                            match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
24129                                Some(t3) => t3,
24130                                None => (
24131                                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
24132                                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
24133                                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
24134                                ),
24135                            };
24136                        let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
24137                        (a, b, c, gt)
24138                    } else {
24139                        (
24140                            e.matmul(&fa.wq, h, 1)?,
24141                            e.matmul(&fa.wk, h, 1)?,
24142                            e.matmul(&fa.wv, h, 1)?,
24143                            e.matmul(gw, h, 1)?,
24144                        )
24145                    }
24146                }
24147            },
24148        };
24149
24150        let mut q = e.uninit(nh * hd)?;
24151        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
24152        let mut k = e.uninit(nkv * hd)?;
24153        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
24154        let ff = if swa {
24155            None
24156        } else {
24157            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
24158        };
24159        #[cfg(debug_assertions)]
24160        if let Some(ff) = ff {
24161            crate::debug_assert_tensor_stream_device(
24162                ff,
24163                &e.stream(),
24164                "step35_decode_attn.rope_freqs",
24165            );
24166        }
24167        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
24168
24169        if std::env::var("MEMRA_NOFA").is_ok() {
24170            return Err(
24171                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
24172                        cache; unset MEMRA_NOFA to use fa_decode"
24173                    .into(),
24174            );
24175        }
24176        let kvl = cache.kv[il].as_mut().unwrap();
24177        let next_len = kvl.len + 1;
24178        let (off, t_kv) = if swa && next_len > win {
24179            (next_len - win, win)
24180        } else {
24181            (0, next_len)
24182        };
24183        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
24184        e.append_kv_quantized(
24185            &k,
24186            &v0,
24187            &mut kvl.k,
24188            &mut kvl.v,
24189            write_row,
24190            kvl.kv_dim_k,
24191            kvl.kv_dim_v,
24192            kvl.k_tok_bytes,
24193            kvl.v_tok_bytes,
24194            crate::Engine::kv_fp8_on(),
24195        )?;
24196        kvl.len = next_len;
24197        let physical = kvl.physical_rows(off, off + t_kv)?;
24198        let k_view = e.view_u8_range(
24199            &kvl.k,
24200            physical.start * kvl.k_tok_bytes,
24201            physical.end * kvl.k_tok_bytes,
24202        );
24203        let v_view = e.view_u8_range(
24204            &kvl.v,
24205            physical.start * kvl.v_tok_bytes,
24206            physical.end * kvl.v_tok_bytes,
24207        );
24208        let mut attn = e.uninit(nh * hd)?;
24209        e.fa_decode_kvmod(
24210            &q,
24211            &k_view,
24212            &v_view,
24213            &mut attn,
24214            hd,
24215            nh,
24216            nkv,
24217            t_kv,
24218            scale,
24219            kvl.k_tok_bytes,
24220            kvl.v_tok_bytes,
24221            crate::Engine::kv_fp8_on(),
24222        )?;
24223
24224        let mut ag = e.uninit(nh * hd)?;
24225        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
24226        self.full_attn_o(e, fa, &ag, 1)
24227    }
24228}
24229
24230// ===================================================================================== //
24231//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
24232//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
24233//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
24234//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
24235//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
24236//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
24237// ===================================================================================== //
24238impl HybridModel {
24239    pub fn is_gemma4_e4b(&self) -> bool {
24240        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
24241    }
24242
24243    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
24244    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
24245    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
24246    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
24247        let g = self.cfg.gemma4.as_ref().unwrap();
24248        let swa = g.swa_pattern[il];
24249        let hd = if swa {
24250            g.key_length_swa
24251        } else {
24252            g.key_length_global
24253        } as usize;
24254        let Mixer::Full(fa) = &self.layers[il].mixer else {
24255            panic!("e4b layer {il} not full-attn")
24256        };
24257        let nh = fa.wq.out_features() / hd;
24258        let nkv = fa.wk.out_features() / hd;
24259        (
24260            hd,
24261            nkv,
24262            nh,
24263            if swa {
24264                g.rope_base_swa
24265            } else {
24266                g.rope_base_global
24267            },
24268            1.0,
24269            swa,
24270        )
24271    }
24272
24273    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
24274    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
24275        self.layers[il]
24276            .gemma4
24277            .as_ref()
24278            .and_then(|b| b.e4b.as_ref())
24279            .and_then(|e4| e4.kv_share.map(|t| t as usize))
24280    }
24281
24282    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
24283    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
24284    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
24285    ///     (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
24286    fn gemma4_e4b_inp_pl(
24287        &self,
24288        e: &Engine,
24289        tokens: &[u32],
24290        x_scaled: &CudaSlice<f32>,
24291        t: usize,
24292    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24293        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
24294        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
24295    }
24296
24297    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
24298    fn gemma4_e4b_inp_pl_dev(
24299        &self,
24300        e: &Engine,
24301        tok_d: &CudaSlice<u32>,
24302        x_scaled: &CudaSlice<f32>,
24303        t: usize,
24304    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24305        let aux = self.gemma4_aux.as_ref().unwrap();
24306        let m = aux.e4b.as_ref().unwrap();
24307        let n_embd = self.cfg.n_embd as usize;
24308        let n_layer = self.layers.len();
24309        let width = m.n_epl * n_layer;
24310        let tbl = m.tok_tbl_gpu.get_or_init(|| {
24311            e.upload_u8(&m.tok_embd_bytes)
24312                .expect("e4b per-layer token table upload")
24313        });
24314        let mut a =
24315            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
24316        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
24317        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
24318        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
24319        let mut pn = e.uninit(t * width)?;
24320        e.rms_norm(
24321            &p,
24322            m.proj_norm.float_data(),
24323            &mut pn,
24324            m.n_epl,
24325            t * n_layer,
24326            self.cfg.rms_eps,
24327        )?;
24328        let mut out = e.uninit(t * width)?;
24329        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
24330        Ok(out)
24331    }
24332
24333    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
24334    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
24335    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
24336    /// already holds this forward's rows — the target runs earlier in the stack).
24337    #[allow(clippy::too_many_arguments)]
24338    fn gemma4_e4b_attn(
24339        &self,
24340        e: &Engine,
24341        il: usize,
24342        hq: &CudaSlice<i8>,
24343        hdq: &CudaSlice<f32>,
24344        pos_d: &CudaSlice<i32>,
24345        t: usize,
24346        cache: &mut Cache,
24347        dc_bucket: Option<usize>,
24348    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24349        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
24350        let eps = self.cfg.rms_eps;
24351        let aux = self.gemma4_aux.as_ref().unwrap();
24352        let ones = aux.ones(e);
24353        #[cfg(debug_assertions)]
24354        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
24355        let Mixer::Full(fa) = &self.layers[il].mixer else {
24356            unreachable!()
24357        };
24358        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
24359        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
24360        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
24361        let h0 = e.zeros(0)?;
24362        let h = &h0;
24363
24364        let ff = if swa {
24365            None
24366        } else {
24367            Some(
24368                aux.rope_freqs(e)
24369                    .expect("e4b global rope needs rope_freqs.weight"),
24370            )
24371        };
24372        #[cfg(debug_assertions)]
24373        if let Some(ff) = ff {
24374            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
24375        }
24376        let share = self.gemma4_e4b_kv_target(il);
24377        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
24378        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
24379        let mut q;
24380        if let Some(_tgt) = share {
24381            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
24382            q = e.uninit(t * nh * hd)?;
24383            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
24384            // empty; q0 stands in for the unused k/v pointers).
24385            let mut kdummy = e.uninit(1)?;
24386            let mut vdummy = e.uninit(1)?;
24387            e.rms_norm_qkv_rope(
24388                &q0,
24389                &q0,
24390                &q0,
24391                fa.q_norm.float_data(),
24392                fa.q_norm.float_data(),
24393                ones,
24394                &mut q,
24395                &mut kdummy,
24396                &mut vdummy,
24397                hd,
24398                self.gemma4_rope_dims(il),
24399                nh * t,
24400                0,
24401                pos_d,
24402                nh,
24403                1,
24404                base,
24405                1.0,
24406                ff,
24407                eps,
24408            )?;
24409        } else {
24410            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
24411            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
24412            // q|k|v rows — the cat norm+rope twin consumes it directly.
24413            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
24414            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
24415            q = e.uninit(t * nh * hd)?;
24416            let mut k = e.uninit(t * nkv * hd)?;
24417            let mut v = e.uninit(t * nkv * hd)?;
24418            if t == 1 && cat.is_some() {
24419                #[allow(clippy::unnecessary_unwrap)]
24420                // allow: the Some-guard sits in a multi-clause regime gate; if-let would reshape the arm structure
24421                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
24422                e.rms_norm_qkv_rope_cat(
24423                    &qkv0,
24424                    fa.q_norm.float_data(),
24425                    fa.k_norm.float_data(),
24426                    ones,
24427                    &mut q,
24428                    &mut k,
24429                    &mut v,
24430                    hd,
24431                    self.gemma4_rope_dims(il),
24432                    nh,
24433                    nkv,
24434                    pos_d,
24435                    nh,
24436                    nkv,
24437                    base,
24438                    1.0,
24439                    ff,
24440                    eps,
24441                )?;
24442            } else {
24443                let (q0, k0, v0) = match if t == 1 {
24444                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
24445                } else {
24446                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
24447                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
24448                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24449                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
24450                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
24451                    } else {
24452                        None
24453                    }
24454                } {
24455                    Some(triple) => triple,
24456                    None => (
24457                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
24458                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
24459                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
24460                    ), // E4B: real v (K != V)
24461                };
24462                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
24463                // the normed rows; V ones-rms, never roped).
24464                e.rms_norm_qkv_rope(
24465                    &q0,
24466                    &k0,
24467                    &v0,
24468                    fa.q_norm.float_data(),
24469                    fa.k_norm.float_data(),
24470                    ones,
24471                    &mut q,
24472                    &mut k,
24473                    &mut v,
24474                    hd,
24475                    self.gemma4_rope_dims(il),
24476                    nh * t,
24477                    nkv * t,
24478                    pos_d,
24479                    nh,
24480                    nkv,
24481                    base,
24482                    1.0,
24483                    ff,
24484                    eps,
24485                )?;
24486            }
24487            let kvl = cache.kv[il].as_mut().unwrap();
24488            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
24489            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
24490            // degenerate tok-0 stream, 2026-07-12).
24491            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24492            if dc_bucket.is_some() {
24493                // DC arm (graph serving): append at the len_d slot, advance the counter
24494                // in-stream — replay-correct, no host len in the launch args. Host mirrors
24495                // are NOT touched here (the replay loop owns them; a bump at capture-record
24496                // time would double-count the capture iteration).
24497                debug_assert!(t == 1);
24498                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
24499                e.append_kv_quantized_row_dc_inc(
24500                    &k,
24501                    &v,
24502                    &mut kvl.k,
24503                    &mut kvl.v,
24504                    &mut kvl.len_d,
24505                    kvl.kv_dim_k,
24506                    kvl.kv_dim_v,
24507                    kvl.k_tok_bytes,
24508                    kvl.v_tok_bytes,
24509                    cls,
24510                )?;
24511            } else {
24512                e.append_kv_quantized_rows(
24513                    &k,
24514                    &v,
24515                    &mut kvl.k,
24516                    &mut kvl.v,
24517                    kvl.len,
24518                    t,
24519                    kvl.kv_dim_k,
24520                    kvl.kv_dim_v,
24521                    kvl.k_tok_bytes,
24522                    kvl.v_tok_bytes,
24523                    cls,
24524                )?;
24525                kvl.len += t;
24526            }
24527            kv_f32 = Some((k, v));
24528        }
24529        // attention: per-row causal fa over the (own or target) quantized cache. The cache
24530        // already contains this forward's rows in both arms; row i attends [.., base+i].
24531        let kvl_idx = share.unwrap_or(il);
24532        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
24533        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
24534        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
24535        let mut attn = e.uninit(t * nh * hd)?;
24536        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
24537        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
24538        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
24539        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
24540        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
24541        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
24542        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
24543        //     rows (the T=K verify kernel; the target appended this forward's rows already).
24544        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
24545        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
24546        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
24547        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
24548            if let Some((kf, vf)) = &kv_f32 {
24549                if hd == 256 && t <= win {
24550                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24551                    return e.matmul(&fa.wo, &attn, t);
24552                }
24553                if hd == 256 && swa && t > win {
24554                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24555                    return e.matmul(&fa.wo, &attn, t);
24556                }
24557                if hd == 512 && !swa {
24558                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24559                    return e.matmul(&fa.wo, &attn, t);
24560                }
24561            } else if share.is_some() {
24562                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24563                let k_view = e.view_u8(&kvl.k, kvl.k.len());
24564                let v_view = e.view_u8(&kvl.v, kvl.v.len());
24565                if hd == 256 && (!swa || t <= win) {
24566                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
24567                    e.fa_prefill_view(
24568                        &q,
24569                        &k_view,
24570                        &v_view,
24571                        &mut attn,
24572                        hd,
24573                        nh,
24574                        nkv,
24575                        t,
24576                        t,
24577                        scale,
24578                        true,
24579                        kvl.k_tok_bytes,
24580                        kvl.v_tok_bytes,
24581                        g,
24582                    )?;
24583                    return e.matmul(&fa.wo, &attn, t);
24584                }
24585                // remaining shared classes (swa above the window; hd512 globals): dequant
24586                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
24587                let kv_dim = nkv * hd;
24588                let mut kf = e.uninit(t * kv_dim)?;
24589                let mut vf = e.uninit(t * kv_dim)?;
24590                e.fa_dequant_kv_view_f32(
24591                    &k_view,
24592                    &v_view,
24593                    &mut kf,
24594                    &mut vf,
24595                    kv_dim,
24596                    kv_dim,
24597                    t,
24598                    kvl.k_tok_bytes,
24599                    kvl.v_tok_bytes,
24600                    g,
24601                )?;
24602                if hd == 512 {
24603                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24604                } else {
24605                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24606                }
24607                return e.matmul(&fa.wo, &attn, t);
24608            }
24609        }
24610        if let Some(bucket) = dc_bucket {
24611            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
24612            // fa_decode_dc over the live counter. len_d already advanced past this token
24613            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
24614            // counter (advanced when the target ran earlier in the stack).
24615            assert!(t == 1);
24616            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
24617            // and under the window every live t_kv sits below it — cap the capture bucket
24618            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
24619            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
24620            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
24621            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
24622                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
24623            } else {
24624                bucket
24625            };
24626            let k_view = e.view_u8(&kvl.k, kvl.k.len());
24627            let v_view = e.view_u8(&kvl.v, kvl.v.len());
24628            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24629            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
24630            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
24631            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
24632            // captured into the dc graph like any other launch. Extending the cascade to
24633            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
24634            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
24635            // MEMRA_WPF=0 rollback seam.
24636            if crate::Engine::wpf_level() >= 1 {
24637                e.prefetch_weight_l2(&fa.wo)?;
24638            }
24639            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
24640            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
24641            if e.uses_q8_1_fast(&fa.wo) {
24642                let mut oq = e.alloc_i8_uninit(nh * hd)?;
24643                let mut od = e.zeros(nh * hd / 32)?;
24644                e.fa_decode_dc_q8(
24645                    &q,
24646                    &k_view,
24647                    &v_view,
24648                    &mut attn,
24649                    hd,
24650                    nh,
24651                    nkv,
24652                    &kvl.len_d,
24653                    bucket,
24654                    scale,
24655                    kvl.k_tok_bytes,
24656                    kvl.v_tok_bytes,
24657                    g,
24658                    Some((&mut oq, &mut od)),
24659                )?;
24660                return e.matmul_pre(&fa.wo, &oq, &od, &attn, t);
24661            }
24662            e.fa_decode_dc(
24663                &q,
24664                &k_view,
24665                &v_view,
24666                &mut attn,
24667                hd,
24668                nh,
24669                nkv,
24670                &kvl.len_d,
24671                bucket,
24672                scale,
24673                kvl.k_tok_bytes,
24674                kvl.v_tok_bytes,
24675                g,
24676            )?;
24677            return e.matmul(&fa.wo, &attn, t);
24678        }
24679        for i in 0..t {
24680            let avail = base_len + i + 1;
24681            let (off_tok, t_kv) = if swa && avail > win {
24682                (avail - win, win)
24683            } else {
24684                (0, avail)
24685            };
24686            let k_view = e.view_u8_range(
24687                &kvl.k,
24688                off_tok * kvl.k_tok_bytes,
24689                (off_tok + t_kv) * kvl.k_tok_bytes,
24690            );
24691            let v_view = e.view_u8_range(
24692                &kvl.v,
24693                off_tok * kvl.v_tok_bytes,
24694                (off_tok + t_kv) * kvl.v_tok_bytes,
24695            );
24696            let qv = e.view(&q, t * nh * hd);
24697            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
24698            let mut q_one = e.uninit(nh * hd)?;
24699            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
24700            let mut a_one = e.uninit(nh * hd)?;
24701            // read class MUST match the append class (globals are e4m3 under gkv): the
24702            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
24703            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
24704            e.fa_decode_kvmod(
24705                &q_one,
24706                &k_view,
24707                &v_view,
24708                &mut a_one,
24709                hd,
24710                nh,
24711                nkv,
24712                t_kv,
24713                scale,
24714                kvl.k_tok_bytes,
24715                kvl.v_tok_bytes,
24716                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
24717            )?;
24718            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
24719        }
24720        e.matmul(&fa.wo, &attn, t)
24721    }
24722
24723    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
24724    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
24725    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
24726    /// layer; does NOT advance cache.pos (caller owns pos).
24727    fn gemma4_e4b_trunk(
24728        &self,
24729        e: &Engine,
24730        tokens: &[u32],
24731        pos0: usize,
24732        cache: &mut Cache,
24733        head_last: bool,
24734    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24735        let n_embd = self.cfg.n_embd as usize;
24736        let t = tokens.len();
24737        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24738        let pos_d = e.htod_i32(&pos)?;
24739        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
24740        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24741        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
24742        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
24743    }
24744
24745    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
24746    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
24747    /// eager chain by construction: SAME functions, not twins).
24748    #[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
24749    fn gemma4_e4b_trunk_core(
24750        &self,
24751        e: &Engine,
24752        x_in: CudaSlice<f32>,
24753        inp_pl: CudaSlice<f32>,
24754        pos_d: &CudaSlice<i32>,
24755        t: usize,
24756        cache: &mut Cache,
24757        dc_bucket: Option<usize>,
24758        cap_logits: bool,
24759        head_last: bool,
24760    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24761        let n_embd = self.cfg.n_embd as usize;
24762        let eps = self.cfg.rms_eps;
24763        let n_layer = self.layers.len();
24764        let mut x = x_in;
24765        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
24766        let n_epl = aux_e4b.n_epl;
24767
24768        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
24769        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
24770        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
24771        // head rides matmul_pre too. First layer's pair comes from a standalone fused
24772        // norm+quant.
24773        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
24774        for il in 0..n_layer {
24775            let layer = &self.layers[il];
24776            let (hq, hdq) = match h_carry.take() {
24777                Some(p) => p,
24778                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
24779            };
24780            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
24781            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
24782            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
24783            let bits = layer.gemma4.as_ref().unwrap();
24784            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
24785            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
24786            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
24787            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
24788            // the fused single-phase reduction is NOT FP-order-identical to the unfused
24789            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
24790            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
24791            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
24792            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
24793            // gate dropped, decode AND verify ride the same fused chain — parity by
24794            // construction, VERIFY-GATE 0.000e0.
24795            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
24796            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
24797                e,
24798                layer,
24799                &o,
24800                &x,
24801                t,
24802                Some(layer.post_attn_norm.float_data()),
24803                fuse_exit,
24804            )?;
24805            let mut resid = e.uninit(t * n_embd)?;
24806            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
24807            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
24808            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
24809            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
24810            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
24811            let g = if fuse_exit {
24812                // sn here = RAW f0 (post_ffw deferred).
24813                let (rq, rd) = e.rms_pre_add_q8_1(
24814                    &sn,
24815                    bits.post_ffw_norm.float_data(),
24816                    &attn_out,
24817                    &mut resid,
24818                    n_embd,
24819                    t,
24820                    self.cfg.rms_eps,
24821                )?;
24822                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
24823            } else {
24824                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
24825                e.matmul(&e4b.inp_gate, &resid, t)?
24826            };
24827            let mut act = e.uninit(t * n_epl)?;
24828            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
24829                let ipv = e.view(&inp_pl, n_epl * n_layer);
24830                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
24831                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
24832                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
24833            } else {
24834                let mut inp_this = e.uninit(t * n_epl)?;
24835                e.copy_rows_strided(
24836                    &inp_pl,
24837                    &mut inp_this,
24838                    n_epl,
24839                    t,
24840                    n_epl * n_layer,
24841                    il * n_epl,
24842                )?;
24843                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
24844                e.matmul(&e4b.proj, &act, t)?
24845            };
24846            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
24847            // ONE launch (glue-fusion lane; last layer emits through output_norm).
24848            let next_norm = if il + 1 < n_layer {
24849                self.layers[il + 1].attn_norm.float_data()
24850            } else {
24851                self.output_norm.float_data()
24852            };
24853            let mut xn = e.uninit(t * n_embd)?;
24854            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
24855                &y,
24856                e4b.post_norm.float_data(),
24857                &resid,
24858                bits.layer_scale,
24859                next_norm,
24860                &mut xn,
24861                n_embd,
24862                t,
24863                eps,
24864            )?;
24865            h_carry = Some(pair);
24866            x = xn;
24867        }
24868        // the head consumes the last layer's fused (output_norm) emit. head_last callers
24869        // (prime, last_only forward) need only the final row's logits — the all-T head is
24870        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
24871        let (oq, odq) = h_carry.take().unwrap();
24872        let h0 = e.zeros(0)?;
24873        let hm = if head_last { 1 } else { t };
24874        let (hq, hd) = if head_last && t > 1 {
24875            let mut q1 = e.uninit_i8(n_embd)?;
24876            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
24877            let nb = n_embd / 32;
24878            let mut d1 = e.uninit(nb)?;
24879            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
24880            (q1, d1)
24881        } else {
24882            (oq, odq)
24883        };
24884        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
24885        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
24886        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
24887        // Logit-returning callers (host logits / spec prime) keep the capped emit.
24888        if cap_logits {
24889            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
24890            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
24891        }
24892        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
24893        Ok((ld, x))
24894    }
24895
24896    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
24897    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
24898    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
24899    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
24900    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
24901    /// covers exactly the layers that appended).
24902    pub fn gemma4_e4b_decode_step_t_am_dev(
24903        &self,
24904        e: &Engine,
24905        tok_d: &CudaSlice<u32>,
24906        t: usize,
24907        pos0: usize,
24908        cache: &mut Cache,
24909    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24910        let n_embd = self.cfg.n_embd as usize;
24911        let eps = self.cfg.rms_eps;
24912        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24913        let pos_d = e.htod_i32(&pos)?;
24914        let embd_gpu = self
24915            .embd_gpu
24916            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
24917        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
24918        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
24919        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24920        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
24921        let (ld, xp) =
24922            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
24923        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
24924        // emit is already capped, matching the eager chain bit-for-bit).
24925        let n_vocab = self.output.out_features();
24926        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
24927        for i in 0..t {
24928            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
24929        }
24930        let mut hn = e.uninit(t * n_embd)?;
24931        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24932        cache.pos += t;
24933        Ok((vam, hn))
24934    }
24935
24936    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
24937    /// prime path — mirror of `gemma4_decode_step_t_h`).
24938    pub(crate) fn gemma4_e4b_decode_step_t_h(
24939        &self,
24940        e: &Engine,
24941        tokens: &[u32],
24942        pos0: usize,
24943        cache: &mut Cache,
24944    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24945        let n_embd = self.cfg.n_embd as usize;
24946        let eps = self.cfg.rms_eps;
24947        let t = tokens.len();
24948        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
24949        let mut hn = e.uninit(t * n_embd)?;
24950        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24951        cache.pos += t;
24952        Ok((e.dtoh(&ld)?, hn))
24953    }
24954
24955    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
24956    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
24957    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
24958    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
24959    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
24960    #[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
24961    pub fn gemma4_e4b_decode_step_dcg(
24962        &self,
24963        e: &Engine,
24964        token_d: &mut CudaSlice<u32>,
24965        pos_d: &mut CudaSlice<i32>,
24966        embd_gpu: &CudaSlice<u8>,
24967        embd_qt: i32,
24968        embd_rb: usize,
24969        cache: &mut Cache,
24970        n_vocab: usize,
24971        bucket: usize,
24972    ) -> Result<(), Box<dyn std::error::Error>> {
24973        let n_embd = self.cfg.n_embd as usize;
24974        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
24975        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
24976        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
24977        let (ld, _x) =
24978            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
24979        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
24980        e.inc_seqlen(pos_d)?;
24981        Ok(())
24982    }
24983
24984    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
24985    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
24986    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
24987    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
24988    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
24989    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
24990    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
24991    #[allow(clippy::too_many_arguments)]
24992    pub fn gemma4_e4b_decode_step_dc(
24993        &self,
24994        e: &Engine,
24995        token_d: &CudaSlice<u32>,
24996        pos_d: &mut CudaSlice<i32>,
24997        embd_gpu: &CudaSlice<u8>,
24998        embd_qt: i32,
24999        embd_rb: usize,
25000        cache: &mut Cache,
25001        n_vocab: usize,
25002    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
25003        let n_embd = self.cfg.n_embd as usize;
25004        let eps = self.cfg.rms_eps;
25005        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
25006        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
25007        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
25008        let (ld, _x) =
25009            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
25010        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
25011        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
25012        e.inc_seqlen(pos_d)?;
25013        cache.pos += 1;
25014        let _ = eps;
25015        Ok(tok_out)
25016    }
25017
25018    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
25019    /// pre-output_norm hidden). Advances cache.pos.
25020    pub(crate) fn gemma4_e4b_decode_step_h(
25021        &self,
25022        e: &Engine,
25023        token: u32,
25024        cache: &mut Cache,
25025    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25026        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
25027        let logits = e.dtoh(&ld)?;
25028        cache.pos += 1;
25029        Ok((logits, x))
25030    }
25031
25032    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
25033    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
25034    /// fast; the prefill fa arms come later.
25035    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
25036    pub(crate) fn gemma4_e4b_prime(
25037        &self,
25038        e: &Engine,
25039        tokens: &[u32],
25040        cache: &mut Cache,
25041    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25042        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
25043        // process-kill as gemma4_prime — refuse per-request.
25044        if cache.pos != 0 {
25045            return Err(
25046                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
25047                        call or decode tokenwise"
25048                    .into(),
25049            );
25050        }
25051        let n_embd = self.cfg.n_embd as usize;
25052        let t = tokens.len();
25053        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
25054        cache.pos += t;
25055        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
25056        let xv = e.view(&x, t * n_embd);
25057        let row = xv.slice((t - 1) * n_embd..t * n_embd);
25058        let mut h_seed = e.uninit(n_embd)?;
25059        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
25060        Ok((last, h_seed, x))
25061    }
25062
25063    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
25064    pub(crate) fn gemma4_e4b_forward(
25065        &self,
25066        e: &Engine,
25067        tokens: &[u32],
25068        last_only: bool,
25069    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
25070        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
25071        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
25072        e.dtoh(&ld) // head_last already reduced to the final row when last_only
25073    }
25074}
25075
25076#[cfg(test)]
25077mod prime_chunk_schedule_tests {
25078    use super::{
25079        CUDA_GRID_YZ_MAX, PRIME_CHUNK_LAUNCH_CAP, PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, PrimePpSignal,
25080        PrimePpStageChannels, PrimePpWaveCredits, PrimePpWaveSlot, active_matrix_values,
25081        align_prime_ranges_to_gdn, dynamic_prime_chunk_ranges, explicit_prime_chunk,
25082        fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring, move_prime_cache_layers,
25083        parse_step_ep_grouped_prefill, parse_step_tp_prefill, prime_cache_stage_for_layer,
25084        recv_prime_pp_signal, restore_prime_cache_layers, step_grouped_decode_shape,
25085        step_grouped_prefill_shape, step_tp_prefill_shape, validate_step_prime_batch_modes,
25086    };
25087
25088    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
25089        ranges.iter().map(|(start, end)| end - start).collect()
25090    }
25091
25092    #[allow(clippy::manual_clamp)] // allow: the min/max chain mirrors the reference arithmetic order in pinned sizing/quant math
25093    fn auto_chunk(t: usize) -> usize {
25094        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
25095    }
25096
25097    /// The ornith cold-long 66k defect (darklanes research/ornith-move-20260829 F2,
25098    /// re-hit on prod 2026-09-01): MEMRA_PRIME_CHUNK=0 ("monolithic") must not schedule
25099    /// a prime range wider than the CUDA grid.y limit, and must stay byte-identical
25100    /// (single monolithic range) for every prompt the limit allows.
25101    #[test]
25102    // assertions_on_constants: the constant relation (cap + fold headroom fits the CUDA
25103    // grid wall) IS the invariant under test; a red here means someone moved a constant.
25104    #[allow(clippy::assertions_on_constants)]
25105    fn monolithic_prime_chunk_caps_at_the_cuda_launch_wall() {
25106        // Ring OFF (ornith serving shape): 0 maps to the launch cap, larger explicit
25107        // values cap there too, workable explicit values pass through untouched.
25108        assert_eq!(explicit_prime_chunk(0, false), PRIME_CHUNK_LAUNCH_CAP);
25109        assert_eq!(explicit_prime_chunk(100_000, false), PRIME_CHUNK_LAUNCH_CAP);
25110        assert_eq!(explicit_prime_chunk(4096, false), 4096);
25111        assert_eq!(
25112            explicit_prime_chunk(PRIME_CHUNK_LAUNCH_CAP, false),
25113            PRIME_CHUNK_LAUNCH_CAP
25114        );
25115        // Ring ON keeps the historical PRIME_CHUNK_MAX_TOKENS clamp exactly.
25116        assert_eq!(
25117            explicit_prime_chunk(0, true),
25118            crate::cache::PRIME_CHUNK_MAX_TOKENS
25119        );
25120        assert_eq!(
25121            explicit_prime_chunk(100_000, true),
25122            crate::cache::PRIME_CHUNK_MAX_TOKENS
25123        );
25124        assert_eq!(explicit_prime_chunk(512, true), 512);
25125        // The fold headroom the cap exists for.
25126        assert!(PRIME_CHUNK_LAUNCH_CAP + PRIME_MIN_T - 1 <= CUDA_GRID_YZ_MAX);
25127    }
25128
25129    #[test]
25130    fn capped_monolithic_ranges_are_identical_below_the_wall_and_legal_above() {
25131        let chunk = explicit_prime_chunk(0, false);
25132        // Every t the CUDA limit allows keeps the exact pre-fix monolithic schedule:
25133        // one range covering the whole prompt (t <= chunk directly, or the < PRIME_MIN_T
25134        // tail folds the split back into a single range).
25135        for t in [
25136            PRIME_MIN_T,
25137            4096,
25138            61_000,
25139            64_984,
25140            PRIME_CHUNK_LAUNCH_CAP,
25141            PRIME_CHUNK_LAUNCH_CAP + 1,
25142            CUDA_GRID_YZ_MAX,
25143        ] {
25144            assert_eq!(
25145                fixed_prime_chunk_ranges_for_ring(t, chunk, false),
25146                vec![(0, t)],
25147                "t={t} must stay a single monolithic range"
25148            );
25149        }
25150        // Above the wall — the sizes the campaign measured failing, the F2 bracket's
25151        // first FAIL, and the boundary — every scheduled range must be launch-legal,
25152        // contiguous, and full-coverage.
25153        for t in [65_536, 65_643, 66_045, 79_717, 82_440, 262_144] {
25154            let ranges = fixed_prime_chunk_ranges_for_ring(t, chunk, false);
25155            assert!(ranges.len() >= 2, "t={t} must chunk");
25156            let mut cursor = 0usize;
25157            for &(start, end) in &ranges {
25158                assert_eq!(start, cursor, "t={t}: ranges must be contiguous");
25159                assert!(
25160                    end - start <= CUDA_GRID_YZ_MAX,
25161                    "t={t}: range width {} exceeds the CUDA grid.y limit",
25162                    end - start
25163                );
25164                assert!(
25165                    end - start >= PRIME_MIN_T,
25166                    "t={t}: range width {} below PRIME_MIN_T",
25167                    end - start
25168                );
25169                cursor = end;
25170            }
25171            assert_eq!(cursor, t, "t={t}: ranges must cover the prompt");
25172        }
25173        // Dense sweep across the boundary band: no width may ever exceed the limit.
25174        for t in (CUDA_GRID_YZ_MAX - 64)..=(CUDA_GRID_YZ_MAX + 2 * PRIME_MIN_T + 64) {
25175            for &(start, end) in &fixed_prime_chunk_ranges_for_ring(t, chunk, false) {
25176                assert!(
25177                    end - start <= CUDA_GRID_YZ_MAX,
25178                    "t={t} width {}",
25179                    end - start
25180                );
25181            }
25182        }
25183    }
25184
25185    #[test]
25186    fn ppn_prime_cache_partition_moves_and_restores_every_layer() {
25187        let round_trip = |fence: &[usize], layers: usize| {
25188            let original: Vec<Option<usize>> = (0..layers).map(Some).collect();
25189            let mut parent = original.clone();
25190            let mut stages: Vec<Vec<Option<usize>>> =
25191                (0..fence.len() - 1).map(|_| vec![None; layers]).collect();
25192
25193            move_prime_cache_layers(&mut parent, &mut stages, fence);
25194            assert!(parent.iter().all(Option::is_none));
25195            for layer in 0..layers {
25196                let owner = prime_cache_stage_for_layer(fence, layer);
25197                for (stage, values) in stages.iter().enumerate() {
25198                    assert_eq!(values[layer], (stage == owner).then_some(layer));
25199                }
25200            }
25201
25202            restore_prime_cache_layers(&mut parent, &mut stages, fence);
25203            assert_eq!(parent, original);
25204            assert!(stages.iter().flatten().all(Option::is_none));
25205        };
25206
25207        // Layers beyond the trunk fence end model MTP/tail state and remain last-stage owned.
25208        round_trip(&[0, 5, 8], 10);
25209        round_trip(&[0, 2, 5, 8], 10);
25210        round_trip(&[0, 1, 3, 6, 8], 10);
25211    }
25212
25213    #[test]
25214    fn ppn_prime_wave_credit_requires_the_exact_oldest_wave_and_slot() {
25215        let mut credits = PrimePpWaveCredits::default();
25216        let wave0 = PrimePpWaveSlot { wave: 0, slot: 1 };
25217        let wave1 = PrimePpWaveSlot { wave: 1, slot: 0 };
25218        credits.record_send(wave0).unwrap();
25219        assert_eq!(credits.release_required(), None);
25220        credits.record_send(wave1).unwrap();
25221        assert_eq!(credits.release_required(), Some(wave0));
25222
25223        assert!(
25224            credits
25225                .record_release(PrimePpWaveSlot { wave: 0, slot: 0 })
25226                .unwrap_err()
25227                .contains("does not match oldest pending")
25228        );
25229        assert_eq!(credits.release_required(), Some(wave0));
25230        credits.record_release(wave0).unwrap();
25231        credits
25232            .record_send(PrimePpWaveSlot { wave: 2, slot: 1 })
25233            .unwrap();
25234        assert!(
25235            credits
25236                .record_send(PrimePpWaveSlot { wave: 4, slot: 0 })
25237                .unwrap_err()
25238                .contains("while wave 3 was next")
25239        );
25240        assert!(
25241            credits
25242                .record_send(PrimePpWaveSlot { wave: 3, slot: 1 })
25243                .unwrap_err()
25244                .contains("reused slot 1")
25245        );
25246    }
25247
25248    #[test]
25249    fn ppn_prime_wave_signal_reports_order_error_injected_error_and_closure() {
25250        let expected = PrimePpWaveSlot { wave: 2, slot: 1 };
25251
25252        let (sender, receiver) = std::sync::mpsc::channel();
25253        sender.send(PrimePpSignal::Slot(expected)).unwrap();
25254        assert_eq!(
25255            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap(),
25256            expected
25257        );
25258
25259        let (sender, receiver) = std::sync::mpsc::channel();
25260        sender
25261            .send(PrimePpSignal::Slot(PrimePpWaveSlot { wave: 3, slot: 1 }))
25262            .unwrap();
25263        assert!(
25264            recv_prime_pp_signal(&receiver, expected, true, "test")
25265                .unwrap_err()
25266                .contains("expected wave/slot")
25267        );
25268
25269        let (sender, receiver) = std::sync::mpsc::channel();
25270        sender
25271            .send(PrimePpSignal::Error("injected stage failure".into()))
25272            .unwrap();
25273        assert_eq!(
25274            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap_err(),
25275            "injected stage failure"
25276        );
25277
25278        let (upstream_sender, upstream_receiver) = std::sync::mpsc::channel();
25279        let (outgoing_sender, outgoing_receiver) = std::sync::mpsc::channel();
25280        let (_release_sender, released_downstream) = std::sync::mpsc::channel();
25281        PrimePpStageChannels {
25282            incoming: None,
25283            release_upstream: Some(upstream_sender),
25284            outgoing: outgoing_sender,
25285            released_downstream,
25286        }
25287        .notify_failure("injected worker error");
25288        assert_eq!(
25289            recv_prime_pp_signal(&upstream_receiver, expected, false, "test").unwrap_err(),
25290            "injected worker error"
25291        );
25292        assert_eq!(
25293            recv_prime_pp_signal(&outgoing_receiver, expected, false, "test").unwrap_err(),
25294            "injected worker error"
25295        );
25296
25297        let (sender, receiver) = std::sync::mpsc::channel::<PrimePpSignal>();
25298        drop(sender);
25299        assert!(
25300            recv_prime_pp_signal(&receiver, expected, true, "test")
25301                .unwrap_err()
25302                .contains("channel closed while waiting for wave 2")
25303        );
25304    }
25305
25306    /// TOOTH for the PP-auto-ranges GDN grid law (lane/hermes-perf-fixes, 2026-08-23;
25307    /// primegrid pattern): the AUTO schedules put internal prime-call boundaries OFF the
25308    /// WY-chunk grid — the broken arm must be demonstrably off-grid, and the aligned twin
25309    /// must land every boundary on it without changing coverage.
25310    #[test]
25311    fn auto_prime_ranges_align_to_the_gdn_grid() {
25312        let c = 32usize; // shipped MEMRA_GDN_CHUNK default/clamp floor
25313        let assert_covers = |ranges: &[(usize, usize)], t: usize| {
25314            assert_eq!(ranges.first().map(|&(s, _)| s), Some(0));
25315            assert_eq!(ranges.last().map(|&(_, e)| e), Some(t));
25316            for w in ranges.windows(2) {
25317                assert_eq!(w[0].1, w[1].0, "ranges must stay contiguous");
25318            }
25319            assert!(ranges.iter().all(|&(s, e)| e > s), "no empty range");
25320        };
25321
25322        // The PP-2 auto geometry at a real agentic length: t=9510 -> fill = 1189 (div_ceil
25323        // by 8), every internal boundary off the 32 grid — the falsified-identity arm.
25324        let t = 9510usize;
25325        let fill = auto_chunk(t);
25326        let fixed = fixed_prime_chunk_ranges(t, fill);
25327        assert!(
25328            fixed[..fixed.len() - 1].iter().any(|&(_, e)| e % c != 0),
25329            "broken arm vanished: fixed auto boundaries all landed on-grid"
25330        );
25331        let dynamic = dynamic_prime_chunk_ranges(t, fill, &fixed);
25332        assert!(
25333            dynamic[..dynamic.len() - 1]
25334                .iter()
25335                .any(|&(_, e)| e % c != 0),
25336            "broken arm vanished: dynamic auto boundaries all landed on-grid"
25337        );
25338
25339        for ranges in [&fixed, &dynamic] {
25340            let aligned = align_prime_ranges_to_gdn(ranges, t, c);
25341            assert_covers(&aligned, t);
25342            for &(_, e) in &aligned[..aligned.len() - 1] {
25343                assert_eq!(e % c, 0, "internal boundary {e} off the {c}-grid");
25344            }
25345            // boundaries only move DOWN, at most c-1 tokens.
25346            for (&(_, a), &(_, b)) in aligned.iter().zip(ranges.iter()) {
25347                assert!(a <= b && b - a < c);
25348            }
25349        }
25350
25351        // Collapse/merge: boundaries inside one grid cell fuse instead of emitting an
25352        // empty range; the schedule survives degenerate short fills.
25353        let tight = vec![(0usize, 33usize), (33, 40), (40, 200)];
25354        let aligned = align_prime_ranges_to_gdn(&tight, 200, c);
25355        assert_covers(&aligned, 200);
25356        assert_eq!(aligned, vec![(0, 32), (32, 200)]);
25357
25358        // No-ops: single range, c=0 (grid off), already-aligned schedules.
25359        assert_eq!(align_prime_ranges_to_gdn(&[(0, 200)], 200, c), [(0, 200)]);
25360        assert_eq!(align_prime_ranges_to_gdn(&tight, 200, 0), tight.as_slice());
25361        let on_grid = vec![(0usize, 128usize), (128, 256), (256, 300)];
25362        assert_eq!(
25363            align_prime_ranges_to_gdn(&on_grid, 300, c),
25364            on_grid.as_slice()
25365        );
25366    }
25367
25368    #[test]
25369    fn active_matrix_prefix_scopes_reused_prime_slabs() {
25370        assert_eq!(
25371            active_matrix_values(40 * 4096, 29, 4096, "activation").unwrap(),
25372            29 * 4096
25373        );
25374        assert_eq!(
25375            active_matrix_values(29 * 4096, 29, 4096, "activation").unwrap(),
25376            29 * 4096
25377        );
25378        assert_eq!(
25379            active_matrix_values(29 * 4096, 24, 4096, "activation").unwrap(),
25380            24 * 4096
25381        );
25382        assert!(active_matrix_values(28 * 4096, 29, 4096, "activation").is_err());
25383        assert!(active_matrix_values(usize::MAX, usize::MAX, 2, "activation").is_err());
25384    }
25385
25386    #[test]
25387    fn step_tp_prefill_batch_refuses_before_scheduler_fallback() {
25388        assert!(validate_step_prime_batch_modes(false, false).is_ok());
25389
25390        let grouped_without_tp = validate_step_prime_batch_modes(false, true).unwrap_err();
25391        assert!(grouped_without_tp.contains("requires MEMRA_STEP_TP_PREFILL=1"));
25392
25393        for grouped in [false, true] {
25394            let err = validate_step_prime_batch_modes(true, grouped).unwrap_err();
25395            assert!(err.contains("did not clear the live-server performance gate"));
25396            assert!(err.contains("per-session grouped prefill"));
25397        }
25398    }
25399
25400    #[test]
25401    fn step_grouped_path_is_eager_single_token_only() {
25402        assert!(step_grouped_decode_shape(false, 1));
25403        assert!(!step_grouped_decode_shape(true, 1));
25404        assert!(!step_grouped_decode_shape(false, 2));
25405        assert!(!step_grouped_decode_shape(true, 2));
25406    }
25407
25408    #[test]
25409    fn step_grouped_prefill_door_is_strict_and_capacity_bounded() {
25410        assert!(!parse_step_ep_grouped_prefill(None).unwrap());
25411        assert!(!parse_step_ep_grouped_prefill(Some("")).unwrap());
25412        assert!(!parse_step_ep_grouped_prefill(Some("0")).unwrap());
25413        assert!(parse_step_ep_grouped_prefill(Some("1")).unwrap());
25414        assert!(parse_step_ep_grouped_prefill(Some("true")).is_err());
25415        assert!(parse_step_ep_grouped_prefill(Some("2")).is_err());
25416
25417        assert!(step_grouped_prefill_shape(true, true, PRIME_MIN_T));
25418        assert!(step_grouped_prefill_shape(
25419            true,
25420            true,
25421            crate::cache::PRIME_CHUNK_MAX_TOKENS,
25422        ));
25423        assert!(!step_grouped_prefill_shape(true, true, PRIME_MIN_T - 1,));
25424        assert!(!step_grouped_prefill_shape(
25425            true,
25426            true,
25427            crate::cache::PRIME_CHUNK_MAX_TOKENS + 1,
25428        ));
25429        assert!(!step_grouped_prefill_shape(false, true, PRIME_MIN_T));
25430        assert!(!step_grouped_prefill_shape(true, false, PRIME_MIN_T));
25431    }
25432
25433    #[test]
25434    fn step_tp_prefill_door_is_strict_and_default_off() {
25435        assert!(!parse_step_tp_prefill(None).unwrap());
25436        assert!(!parse_step_tp_prefill(Some("")).unwrap());
25437        assert!(!parse_step_tp_prefill(Some("0")).unwrap());
25438        assert!(parse_step_tp_prefill(Some("1")).unwrap());
25439        assert!(parse_step_tp_prefill(Some("true")).is_err());
25440        assert!(parse_step_tp_prefill(Some("2")).is_err());
25441    }
25442
25443    #[test]
25444    fn step_tp_prefill_requires_a_qualified_even_rank_shape() {
25445        assert!(step_tp_prefill_shape(
25446            true,
25447            PRIME_MIN_T,
25448            4,
25449            true,
25450            true,
25451            false,
25452        ));
25453        assert!(!step_tp_prefill_shape(
25454            false,
25455            PRIME_MIN_T,
25456            4,
25457            true,
25458            true,
25459            false,
25460        ));
25461        assert!(!step_tp_prefill_shape(
25462            true,
25463            PRIME_MIN_T - 1,
25464            4,
25465            true,
25466            true,
25467            false,
25468        ));
25469        // TP2 admits (2026-08-25); odd/1-card placements still refuse.
25470        assert!(step_tp_prefill_shape(
25471            true,
25472            PRIME_MIN_T,
25473            2,
25474            true,
25475            true,
25476            false
25477        ));
25478        assert!(!step_tp_prefill_shape(
25479            true,
25480            PRIME_MIN_T,
25481            1,
25482            true,
25483            true,
25484            false
25485        ));
25486        assert!(!step_tp_prefill_shape(
25487            true,
25488            PRIME_MIN_T,
25489            3,
25490            true,
25491            true,
25492            false
25493        ));
25494        assert!(!step_tp_prefill_shape(
25495            true,
25496            PRIME_MIN_T,
25497            4,
25498            false,
25499            true,
25500            false,
25501        ));
25502        assert!(!step_tp_prefill_shape(
25503            true,
25504            PRIME_MIN_T,
25505            4,
25506            true,
25507            false,
25508            false,
25509        ));
25510        assert!(!step_tp_prefill_shape(
25511            true,
25512            PRIME_MIN_T,
25513            4,
25514            true,
25515            true,
25516            true,
25517        ));
25518    }
25519
25520    #[test]
25521    fn fixed_schedule_retains_measured_geometry() {
25522        assert_eq!(
25523            sizes(&fixed_prime_chunk_ranges(461, 128)),
25524            vec![128, 128, 128, 77]
25525        );
25526        assert_eq!(
25527            sizes(&fixed_prime_chunk_ranges(1833, 230)),
25528            vec![230, 230, 230, 230, 230, 230, 230, 223]
25529        );
25530        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
25531        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
25532        assert_eq!(capped, vec![4096, 4088, 16]);
25533        assert!(capped.iter().all(|&rows| rows <= 4096));
25534        assert_eq!(
25535            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
25536            vec![4100],
25537            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
25538        );
25539    }
25540
25541    #[test]
25542    fn dynamic_schedule_matches_registered_shapes() {
25543        let cases = [
25544            (461, vec![64, 141, 132, 124]),
25545            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
25546            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
25547        ];
25548        for (t, expected) in cases {
25549            let chunk = auto_chunk(t);
25550            let fixed = fixed_prime_chunk_ranges(t, chunk);
25551            assert_eq!(
25552                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
25553                expected
25554            );
25555        }
25556    }
25557
25558    #[test]
25559    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
25560        for t in 256..=8192 {
25561            let chunk = auto_chunk(t);
25562            let fixed = fixed_prime_chunk_ranges(t, chunk);
25563            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
25564            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
25565            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
25566            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
25567            for pair in dynamic.windows(2) {
25568                assert_eq!(pair[0].1, pair[1].0, "T={t}");
25569            }
25570            assert!(
25571                dynamic
25572                    .iter()
25573                    .all(|(start, end)| end - start >= PRIME_MIN_T),
25574                "T={t} sizes={:?}",
25575                sizes(&dynamic)
25576            );
25577            if dynamic.len() >= 3 {
25578                let chunk_sizes = sizes(&dynamic);
25579                assert!(
25580                    chunk_sizes[0] < chunk_sizes[1],
25581                    "T={t} sizes={chunk_sizes:?}"
25582                );
25583                assert!(
25584                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
25585                    "T={t} sizes={chunk_sizes:?}"
25586                );
25587            }
25588        }
25589    }
25590}
25591
25592#[cfg(test)]
25593mod page_prefetch_tests {
25594    use super::{
25595        grouped_worker_prefetch_position, page_prefetch_positions,
25596        page_prefetch_window_from_values, worker_prefetch_positions,
25597    };
25598
25599    #[test]
25600    fn page_prefetch_window_keeps_existing_opt_in_default() {
25601        assert_eq!(page_prefetch_window_from_values(false, None), 0);
25602        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
25603        assert_eq!(page_prefetch_window_from_values(true, None), 1);
25604        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
25605        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
25606        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
25607    }
25608
25609    #[test]
25610    fn rolling_page_prefetch_advises_each_future_expert_once() {
25611        let advised: Vec<_> = (0..7)
25612            .flat_map(|position| page_prefetch_positions(position, 7, 3))
25613            .collect();
25614        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
25615
25616        let one_ahead: Vec<_> = (0..4)
25617            .flat_map(|position| page_prefetch_positions(position, 4, 1))
25618            .collect();
25619        assert_eq!(one_ahead, vec![1, 2, 3]);
25620        assert!(page_prefetch_positions(0, 4, 0).is_empty());
25621    }
25622
25623    #[test]
25624    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
25625        assert_eq!(grouped_worker_prefetch_position(0, None), None);
25626        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
25627            .chain(
25628                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
25629            )
25630            .collect();
25631        assert_eq!(positions, vec![0, 1, 2, 3]);
25632        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
25633    }
25634
25635    #[test]
25636    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
25637        let queued: Vec<_> = (0..8)
25638            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
25639            .collect();
25640        assert_eq!(queued, (0..8).collect::<Vec<_>>());
25641
25642        let one_at_a_time: Vec<_> = (0..4)
25643            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
25644            .collect();
25645        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
25646        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
25647    }
25648}
25649
25650pub struct G4DcSlots {
25651    x: CudaSlice<f32>,
25652    xn: CudaSlice<f32>,
25653    cur: CudaSlice<f32>,
25654    hq: CudaSlice<i8>,
25655    hd_: CudaSlice<f32>,
25656    q0: CudaSlice<f32>,
25657    k0: CudaSlice<f32>,
25658    v0: CudaSlice<f32>,
25659    q: CudaSlice<f32>,
25660    k: CudaSlice<f32>,
25661    v: CudaSlice<f32>,
25662    attn: CudaSlice<f32>,
25663    o: CudaSlice<f32>,
25664    attn_out: CudaSlice<f32>,
25665    zsh: CudaSlice<f32>,
25666    zq: CudaSlice<i8>,
25667    zd: CudaSlice<f32>,
25668    gate: CudaSlice<f32>,
25669    up: CudaSlice<f32>,
25670    act: CudaSlice<f32>,
25671    actq: CudaSlice<i8>,
25672    actd: CudaSlice<f32>,
25673    f0: CudaSlice<f32>,
25674    sn: CudaSlice<f32>,
25675    hn: CudaSlice<f32>,
25676    logits: CudaSlice<f32>,
25677}
25678
25679/// Whole-token decode graph state (step TP graph increment B). One stitched multi-device
25680/// parent per fa bucket, plus the persistent host->graph plumbing: the device token id the
25681/// in-graph embed gathers, the device position the ropes read (advanced in-graph), and the
25682/// fixed logits stage the head writes.
25683pub struct Step35TokenGraphState {
25684    /// (bucket_max, graph) — bucket keyed by the fa split geometry (fa_geom_eager).
25685    pub graphs: Vec<(usize, crate::tp::TokenGraph)>,
25686    pub token_d: cudarc::driver::CudaSlice<u32>,
25687    pub pos_d: cudarc::driver::CudaSlice<i32>,
25688    pub logits_stage: cudarc::driver::CudaSlice<f32>,
25689    /// Cross-child intermediates MUST live at fixed addresses (graph mem nodes remap at
25690    /// launch, so an alloc made inside one captured child is not referable from another):
25691    /// the running residual, the post-attention pair, the shared-expert row, and the
25692    /// e-context mirrors of the root-produced attention output and K/V shadow rows.
25693    pub x: cudarc::driver::CudaSlice<f32>,
25694    pub x1: cudarc::driver::CudaSlice<f32>,
25695    pub mixed_stage: cudarc::driver::CudaSlice<f32>,
25696    pub sh_stage: cudarc::driver::CudaSlice<f32>,
25697    pub k_shadow_stage: cudarc::driver::CudaSlice<f32>,
25698    pub v_shadow_stage: cudarc::driver::CudaSlice<f32>,
25699    /// Alloc-free e-section scratch (child graphs cannot contain mem nodes): router logits,
25700    /// shared-expert gate/up/act rows + sigmoid scalar, dense-FFN z/gate/up/act, head hidden.
25701    pub router_logits: cudarc::driver::CudaSlice<f32>,
25702    pub shexp_gate: cudarc::driver::CudaSlice<f32>,
25703    pub shexp_up: cudarc::driver::CudaSlice<f32>,
25704    pub shexp_act: cudarc::driver::CudaSlice<f32>,
25705    pub gate_sig: cudarc::driver::CudaSlice<f32>,
25706    pub dense_z: cudarc::driver::CudaSlice<f32>,
25707    pub dense_gate: cudarc::driver::CudaSlice<f32>,
25708    pub dense_up: cudarc::driver::CudaSlice<f32>,
25709    pub dense_act: cudarc::driver::CudaSlice<f32>,
25710    pub hn: cudarc::driver::CudaSlice<f32>,
25711    /// MEMRA_TG_PROBE_LAYER diagnostics: capture-time copies of layer K's attention output
25712    /// and post-FFN residual, dumped after replay for graph-vs-eager layer bisection.
25713    pub probe_mixed: cudarc::driver::CudaSlice<f32>,
25714    pub probe_x: cudarc::driver::CudaSlice<f32>,
25715    /// Chunk-loop (F-lite): device token history ring + its device write index, filled by
25716    /// the in-graph tail argmax chain; host reads the ring once per chunk.
25717    pub token_hist: cudarc::driver::CudaSlice<u32>,
25718    pub hist_idx: cudarc::driver::CudaSlice<i32>,
25719}
25720
25721impl HybridModel {
25722    /// Whole-token decode graph (step TP graph increment B): ONE stitched multi-device parent
25723    /// replays the entire 45-layer token — the launch-collapse the per-layer minis could not
25724    /// reach. Returns Some(logits) when the graph handled the token, None for eager fallback
25725    /// (door off, ineligible class, any layer's ring would rebase, or a bucket boundary that
25726    /// needs a rebuild this token).
25727    ///
25728    /// v1 SCOPE (diagnostic door, default OFF): the LOCAL shadow caches advance their lengths
25729    /// but not their contents under this door (the TP rank caches are fully maintained
25730    /// in-graph via the dcw counters). Sessions relying on shadow CONTENT (save/rollback)
25731    /// must not run with the door on until the local-dcw twin lands.
25732    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
25733    pub(crate) fn step35_token_graph_step(
25734        &self,
25735        e: &Engine,
25736        token: u32,
25737        cache: &mut Cache,
25738    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
25739        if !self.uses_sliding_gated_moe_program()
25740            || !crate::tp::step_tp_graph_enabled()?
25741            || !crate::tp::step_tp_dcw_enabled()?
25742            || !crate::tp::step_tp_qkv_fused_enabled()?
25743            || !crate::tp::step_tp_dev_router_enabled()?
25744            || !crate::tp::step_nvfp4_dev_routes_enabled()?
25745        {
25746            return Ok(None);
25747        }
25748        // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): the eager
25749        // token step is this route's byte-identical twin — warmup and rebase tokens
25750        // already ride it — so below the driver-free floor the token goes eager
25751        // (`Ok(None)` = the caller's eager fallback) instead of feeding cuGraphLaunch
25752        // an exhausted card (lane/graph-launch-guard-sweep-20260831).
25753        if !crate::spec::graph_launch_headroom_ok(e) {
25754            static NOTED: std::sync::Once = std::sync::Once::new();
25755            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
25756            return Ok(None);
25757        }
25758        let n_embd = self.cfg.n_embd as usize;
25759        let n_vocab = self.cfg.n_vocab as usize;
25760        let n_layers = self.layers.len();
25761        let pos = cache.pos;
25762        let staged_next = pos + 1;
25763        if staged_next < 96 {
25764            return Ok(None); // sub-vec-floor contexts keep eager (fa kernel-class boundary)
25765        }
25766
25767        // Per-layer eligibility: every TP layer contiguous-appends this token (any rebase ->
25768        // eager fallback for the whole token; the host path also updates base_d there).
25769        for il in 0..n_layers {
25770            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
25771                return Ok(None); // caches not hydrated yet — eager warms them
25772            };
25773            if tp_kv.peek_append_ring(1)?.1 {
25774                return Ok(None);
25775            }
25776        }
25777
25778        // Bucket key: the global layers' fa split geometry at this depth (SWA layers cap at
25779        // their window and share one bucket forever after ctx > window).
25780        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
25781        if !fa_vec {
25782            return Ok(None);
25783        }
25784        let sp = crate::fa_split_keys(staged_next, 8);
25785        let bucket_max = (n_splits * sp).max(staged_next);
25786
25787        let mut state_guard = self
25788            .step35_token_graph
25789            .lock()
25790            .map_err(|_| "step35 token graph lock is poisoned")?;
25791        if state_guard.is_none() {
25792            let _main = e.gpu.enter_main()?;
25793            let n_expert = self
25794                .cfg
25795                .moe
25796                .as_ref()
25797                .map(|m| m.expert_count as usize)
25798                .unwrap_or(0);
25799            let n_ff_sh = self
25800                .layers
25801                .iter()
25802                .find_map(|l| match &l.ffn {
25803                    crate::hybrid::Ffn::Moe(m) => m.gate_shexp.as_ref().map(|g| g.out_features()),
25804                    _ => None,
25805                })
25806                .unwrap_or(0);
25807            let n_ff_dense = self
25808                .layers
25809                .iter()
25810                .find_map(|l| match &l.ffn {
25811                    crate::hybrid::Ffn::Dense { ffn_gate, .. } => Some(ffn_gate.out_features()),
25812                    _ => None,
25813                })
25814                .unwrap_or(0);
25815            *state_guard = Some(Step35TokenGraphState {
25816                graphs: Vec::new(),
25817                token_d: e.stream().clone_htod(&[0u32])?,
25818                pos_d: e.htod_i32(&[pos as i32])?,
25819                logits_stage: e.htod(&vec![0.0f32; n_vocab])?,
25820                x: e.htod(&vec![0.0f32; n_embd])?,
25821                x1: e.htod(&vec![0.0f32; n_embd])?,
25822                mixed_stage: e.htod(&vec![0.0f32; n_embd])?,
25823                sh_stage: e.htod(&vec![0.0f32; n_embd])?,
25824                k_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25825                v_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25826                router_logits: e.htod(&vec![0.0f32; n_expert.max(1)])?,
25827                shexp_gate: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25828                shexp_up: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25829                shexp_act: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25830                gate_sig: e.htod(&[1.0f32; 1])?,
25831                dense_z: e.htod(&vec![0.0f32; n_embd])?,
25832                dense_gate: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25833                dense_up: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25834                dense_act: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25835                hn: e.htod(&vec![0.0f32; n_embd])?,
25836                probe_mixed: e.htod(&vec![0.0f32; n_embd])?,
25837                probe_x: e.htod(&vec![0.0f32; n_embd])?,
25838                token_hist: e.stream().clone_htod(&[0u32; 16])?,
25839                hist_idx: e.htod_i32(&[0])?,
25840            });
25841        }
25842        let state = state_guard.as_mut().expect("state armed above");
25843        // Pre-arm the argmax partials OUTSIDE any capture: the launcher allocates them on
25844        // first use, and an alloc inside a captured section is a mem node (child graphs
25845        // reject those — the tail argmax chain needs them already resident).
25846        {
25847            let _main = e.gpu.enter_main()?;
25848            let Step35TokenGraphState {
25849                logits_stage,
25850                token_d,
25851                ..
25852            } = &mut *state;
25853            e.argmax_token_device_into(logits_stage, token_d, n_vocab)?;
25854        }
25855
25856        // ONE graph, retargeted per bucket (increment C): the per-16-token whole rebuild was
25857        // ~55ms (3.4ms/token persistent); the M1 exec update path moves nsp/ski/gridDimY and
25858        // the partial-pool memset widths in ~1ms. The partial pool is pre-grown to the run
25859        // ceiling at build so the baked pointers never move.
25860        if state.graphs.is_empty() {
25861            // Build the parent at this bucket. Capture executes nothing; correctness is
25862            // pinned at replay by the token-identity gate.
25863            self.step35_token_graph_build(e, cache, state, bucket_max)?;
25864        }
25865        {
25866            let (b, g) = state.graphs.first_mut().expect("graph built above");
25867            if *b != bucket_max {
25868                g.retarget_bucket(bucket_max)?;
25869                *b = bucket_max;
25870            }
25871        }
25872        let graph = state
25873            .graphs
25874            .first()
25875            .map(|(_, g)| g)
25876            .expect("graph built above");
25877
25878        let tg_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
25879        let t_fence = tg_timing.then(std::time::Instant::now);
25880        // Rank-stream fence: an eager token (warmup, rebase fallback) leaves len-mirror sets
25881        // queued on the rank streams, and graph children carry no ordering edge to those
25882        // streams. Graph tokens themselves enqueue nothing there (external commit), so this
25883        // sync is a no-op between consecutive replays.
25884        {
25885            let fa0 = match &self.layers[0].mixer {
25886                Mixer::Full(fa) => fa,
25887                _ => return Err("step35 token graph expects full-attention layers".into()),
25888            };
25889            let tp0 = fa0
25890                .step_tp_qkv
25891                .as_ref()
25892                .ok_or("step35 token graph lost its TP state")?;
25893            for rank in 0..tp0.runtime.devices().len() {
25894                let engine = tp0
25895                    .runtime
25896                    .rank_engine(rank)
25897                    .ok_or("step35 token graph lost a rank engine")?;
25898                let _main = engine.gpu.enter_main()?;
25899                engine.stream().synchronize()?;
25900            }
25901        }
25902
25903        // Replay: feed the token, launch, read the logits, mirror the host bookkeeping.
25904        {
25905            let _main = e.gpu.enter_main()?;
25906            e.set_u32_one(&mut state.token_d, token)?;
25907            e.set_i32_one(&mut state.pos_d, pos as i32)?;
25908        }
25909        let t_launch = tg_timing.then(std::time::Instant::now);
25910        graph.launch(e)?;
25911        let t_book = tg_timing.then(std::time::Instant::now);
25912        // Host bookkeeping OVERLAPS the replay (nsys 2026-08-21: the 45-layer txn loop was
25913        // 5.4ms/token of inter-token gap when it ran after the sync). Host-only work except
25914        // the local len_d set, which is stream-ordered AFTER the graph on e's stream. On a
25915        // replay error the counters are already advanced — acceptable: the decode aborts.
25916        for il in 0..n_layers {
25917            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
25918            let transaction = tp_kv.begin_transaction()?;
25919            let fa = match &self.layers[il].mixer {
25920                Mixer::Full(fa) => fa,
25921                _ => return Err("step35 token graph expects full-attention layers".into()),
25922            };
25923            let tp = fa
25924                .step_tp_qkv
25925                .as_ref()
25926                .ok_or("step35 token graph lost its TP state")?;
25927            // Bookkeeping-only txn (external appends; mirror sets skipped — the in-graph
25928            // incs own the counters). Shards unused.
25929            let empty: [CudaSlice<f32>; 0] = [];
25930            tp.runtime.append_tp_kv_transaction_inner(
25931                tp_kv,
25932                transaction,
25933                &empty,
25934                &empty,
25935                1,
25936                true,
25937            )?;
25938            tp.runtime
25939                .commit_tp_kv_transaction_external(tp_kv, transaction, 1)?;
25940            // Local shadow: lengths advance (v1 keeps contents stale under the door).
25941            if let Some(local) = cache.kv[il].as_mut() {
25942                local.len = pos + 1;
25943                let _main = e.gpu.enter_main()?;
25944                e.set_i32_one(&mut local.len_d, (pos + 1) as i32)?;
25945            }
25946        }
25947        cache.pos = pos + 1;
25948        let t_sync = tg_timing.then(std::time::Instant::now);
25949        let (logits, h_seed) = {
25950            let _main = e.gpu.enter_main()?;
25951            e.stream().synchronize()?;
25952            (e.dtoh(&state.logits_stage)?, e.clone_dtod(&state.x)?)
25953        };
25954        if let (Some(f), Some(l), Some(b), Some(sy)) = (t_fence, t_launch, t_book, t_sync) {
25955            use std::sync::atomic::{AtomicU64, Ordering};
25956            static NS: [AtomicU64; 5] = [
25957                AtomicU64::new(0),
25958                AtomicU64::new(0),
25959                AtomicU64::new(0),
25960                AtomicU64::new(0),
25961                AtomicU64::new(0),
25962            ];
25963            static CALLS: AtomicU64 = AtomicU64::new(0);
25964            let now = std::time::Instant::now();
25965            NS[0].fetch_add((l - f).as_nanos() as u64, Ordering::Relaxed); // fence+set
25966            NS[1].fetch_add((b - l).as_nanos() as u64, Ordering::Relaxed); // launch call
25967            NS[2].fetch_add((sy - b).as_nanos() as u64, Ordering::Relaxed); // bookkeeping
25968            NS[3].fetch_add((now - sy).as_nanos() as u64, Ordering::Relaxed); // sync+dtoh
25969            NS[4].fetch_add((now - f).as_nanos() as u64, Ordering::Relaxed); // total
25970            let calls = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
25971            if calls.is_multiple_of(100) {
25972                let avg = |i: usize| NS[i].load(Ordering::Relaxed) as f64 / calls as f64 / 1e3;
25973                eprintln!(
25974                    "[tg-timing] calls={calls} fence_us={:.0} launch_us={:.0} book_us={:.0} \
25975                     syncdtoh_us={:.0} total_us={:.0}",
25976                    avg(0),
25977                    avg(1),
25978                    avg(2),
25979                    avg(3),
25980                    avg(4)
25981                );
25982            }
25983        }
25984        // MEMRA_TG_PROBE_LAYER diagnostics: append the captured layer-K probes.
25985        if std::env::var("MEMRA_TG_PROBE_LAYER").is_ok() {
25986            use std::io::Write;
25987            let (pm, px) = {
25988                let _main = e.gpu.enter_main()?;
25989                (e.dtoh(&state.probe_mixed)?, e.dtoh(&state.probe_x)?)
25990            };
25991            for (path, data) in [
25992                ("/root/tg-probe-mixed.bin", &pm),
25993                ("/root/tg-probe-x.bin", &px),
25994            ] {
25995                let mut fo = std::fs::OpenOptions::new()
25996                    .create(true)
25997                    .append(true)
25998                    .open(path)?;
25999                for v in data {
26000                    fo.write_all(&v.to_le_bytes())?;
26001                }
26002            }
26003        }
26004        // MEMRA_DUMP_HN twin of the eager tail's dump (same format: appended raw LE f32 rows)
26005        // so a graph arm and an eager arm produce position-aligned pre-head hidden streams.
26006        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
26007            let hh = {
26008                let _main = e.gpu.enter_main()?;
26009                e.dtoh(&state.hn)?
26010            };
26011            use std::io::Write;
26012            let mut fo = std::fs::OpenOptions::new()
26013                .create(true)
26014                .append(true)
26015                .open(path)?;
26016            for v in &hh {
26017                fo.write_all(&v.to_le_bytes())?;
26018            }
26019        }
26020        // MEMRA_STEP_TP_GRAPH_DEBUG=1: per-token device-counter dump (drift hunts). One dtoh
26021        // per rank per token; diagnostics only.
26022        if std::env::var("MEMRA_STEP_TP_GRAPH_DEBUG").as_deref() == Ok("1") {
26023            for il in [0usize, 1, 44] {
26024                let tp_kv = cache.tp_kv[il].as_ref().expect("eligibility checked above");
26025                let host_len = tp_kv.staged_len();
26026                let fa = match &self.layers[il].mixer {
26027                    Mixer::Full(fa) => fa,
26028                    _ => continue,
26029                };
26030                let tp = fa
26031                    .step_tp_qkv
26032                    .as_ref()
26033                    .ok_or("step35 token graph lost its TP state")?;
26034                for rank in 0..tp.runtime.devices().len() {
26035                    let engine = tp
26036                        .runtime
26037                        .rank_engine(rank)
26038                        .ok_or("step35 token graph lost a rank engine")?;
26039                    let rank_cache = tp_kv.rank(rank).ok_or("debug rank cache missing")?;
26040                    let _main = engine.gpu.enter_main()?;
26041                    engine.stream().synchronize()?;
26042                    let len_d = engine.dtoh_i32_one(rank_cache.len_d())?;
26043                    let base_d = match rank_cache.base_d() {
26044                        Some(b) => engine.dtoh_i32_one(b)?,
26045                        None => -1,
26046                    };
26047                    eprintln!(
26048                        "[graph-debug] pos={pos} il={il} rank={rank} host_len={host_len} \
26049                         len_d={len_d} base_d={base_d}"
26050                    );
26051                }
26052            }
26053        }
26054        Ok(Some((logits, h_seed)))
26055    }
26056
26057    /// MEMRA_HEAD_SPLIT worker: dev0 computes lm-head rows [0, half), rank1 computes
26058    /// [half, n_vocab) concurrently on its otherwise-idle tail, and the halves concatenate
26059    /// on e — bit-identical per logit to the single-device matvec. Process-static workspace
26060    /// (SHEXP_WS pattern) pinned by the head tensor pointer; rank1 holds a one-time 0.5GB
26061    /// replica of its row half. Returns None when ineligible (no bf16 head / no rank1).
26062    pub(crate) fn head_split_matvec(
26063        &self,
26064        e: &Engine,
26065        hn: &CudaSlice<f32>,
26066    ) -> Result<Option<Vec<f32>>, Box<dyn std::error::Error>> {
26067        if self.head_split_fill_device(e, hn)?.is_none() {
26068            return Ok(None);
26069        }
26070        let guard = HEAD_SPLIT_WS
26071            .lock()
26072            .map_err(|_| "head split lock is poisoned")?;
26073        let ws = guard.as_ref().expect("filled above");
26074        let _main = e.gpu.enter_main()?;
26075        Ok(Some(e.dtoh(&ws.logits_e)?))
26076    }
26077
26078    /// Compute body of the split head: arms the replica + staging on first use, then fills
26079    /// the persistent full-logits row (e's half by view matvec, rank1's half by raw P2P
26080    /// push) and orders e's stream behind it. None = ineligible.
26081    fn head_split_fill_device(
26082        &self,
26083        e: &Engine,
26084        hn: &CudaSlice<f32>,
26085    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
26086        use cudarc::driver::DevicePtr;
26087        let crate::model::GpuTensor::FloatBf16 { data: head, .. } = &self.output else {
26088            return Ok(None);
26089        };
26090        let Some(rank1) = self.layers.first().and_then(|l| match &l.mixer {
26091            Mixer::Full(fa) => fa
26092                .step_tp_qkv
26093                .as_ref()
26094                .and_then(|tp| tp.runtime.rank_engine(1)),
26095            _ => None,
26096        }) else {
26097            return Ok(None);
26098        };
26099        let n_embd = self.cfg.n_embd as usize;
26100        let n_vocab = self.cfg.n_vocab as usize;
26101        let half = n_vocab / 2;
26102        let mut guard = HEAD_SPLIT_WS
26103            .lock()
26104            .map_err(|_| "head split lock is poisoned")?;
26105        let pin = {
26106            let _main = e.gpu.enter_main()?;
26107            let stream = e.stream();
26108            let (ptr, _g) = head.device_ptr(&stream);
26109            ptr
26110        };
26111        if guard.as_ref().is_none_or(|ws| ws.pin != pin) {
26112            // One-time: upload rank1's row half + persistent staging.
26113            let hi_rows = n_vocab - half;
26114            let (w1, hn1, y1, ev_done) = {
26115                let _r1 = rank1.gpu.enter_main()?;
26116                (
26117                    rank1.alloc_u8_uninit(hi_rows * n_embd * 2)?,
26118                    rank1.htod(&vec![0.0f32; n_embd])?,
26119                    rank1.htod(&vec![0.0f32; hi_rows])?,
26120                    rank1.ctx().new_event(None)?,
26121                )
26122            };
26123            {
26124                use cudarc::driver::sys;
26125                let src = pin + (half * n_embd * 2) as u64;
26126                let dst = {
26127                    let _r1 = rank1.gpu.enter_main()?;
26128                    let rstream = rank1.stream();
26129                    let (d, _g) = w1.device_ptr(&rstream);
26130                    d
26131                };
26132                let _r1 = rank1.gpu.enter_main()?;
26133                let r = unsafe {
26134                    sys::cuMemcpyAsync(
26135                        dst as sys::CUdeviceptr,
26136                        src as sys::CUdeviceptr,
26137                        hi_rows * n_embd * 2,
26138                        rank1.stream().cu_stream() as sys::CUstream,
26139                    )
26140                };
26141                if r != sys::CUresult::CUDA_SUCCESS {
26142                    return Err(format!("head split replica upload: {r:?}").into());
26143                }
26144                rank1.stream().synchronize()?;
26145            }
26146            let (logits_e, ev_hn) = {
26147                let _main = e.gpu.enter_main()?;
26148                (e.htod(&vec![0.0f32; n_vocab])?, e.ctx().new_event(None)?)
26149            };
26150            let (raw_hn1, raw_y1) = {
26151                let _r1 = rank1.gpu.enter_main()?;
26152                let rstream = rank1.stream();
26153                let (a, _g0) = hn1.device_ptr(&rstream);
26154                let (b, _g1) = y1.device_ptr(&rstream);
26155                (a, b)
26156            };
26157            let raw_logits_hi = {
26158                let _main = e.gpu.enter_main()?;
26159                let stream = e.stream();
26160                let (l, _g) = logits_e.device_ptr(&stream);
26161                l + (half * 4) as u64
26162            };
26163            *guard = Some(HeadSplit {
26164                pin,
26165                w1,
26166                hn1,
26167                y1,
26168                logits_e,
26169                ev_hn,
26170                ev_done,
26171                raw_hn1,
26172                raw_y1,
26173                raw_logits_hi,
26174                samp: None,
26175            });
26176        }
26177        let ws = guard.as_mut().expect("armed above");
26178        let hi_rows = n_vocab - half;
26179        // e: signal hn ready; rank1: pull hn, matvec its half, push the logits half back.
26180        let raw_hn = {
26181            let _main = e.gpu.enter_main()?;
26182            let stream = e.stream();
26183            let (h, _g) = hn.device_ptr(&stream);
26184            ws.ev_hn.record(&stream)?;
26185            h
26186        };
26187        {
26188            let _r1 = rank1.gpu.enter_main()?;
26189            rank1.stream().wait(&ws.ev_hn)?;
26190            crate::tp::raw_copy_bytes(ws.raw_hn1, raw_hn, n_embd * 4, rank1)?;
26191            let HeadSplit { w1, hn1, y1, .. } = &mut *ws;
26192            rank1.matvec_bf16_into(w1, hn1, y1, n_embd, hi_rows)?;
26193            crate::tp::raw_copy_bytes(ws.raw_logits_hi, ws.raw_y1, hi_rows * 4, rank1)?;
26194            ws.ev_done.record(&rank1.stream())?;
26195        }
26196        {
26197            let _main = e.gpu.enter_main()?;
26198            let head_lo = head.slice(0..half * n_embd * 2);
26199            let HeadSplit { logits_e, .. } = &mut *ws;
26200            // Writes rows [0, half) of logits_e; rank1's raw push fills [half, n_vocab).
26201            e.matvec_bf16_view_into(&head_lo, hn, logits_e, n_embd, half)?;
26202            e.stream().wait(&ws.ev_done)?;
26203            Ok(Some(()))
26204        }
26205    }
26206
26207    /// Device twin of `head_split_matvec` for the chain: fills the persistent full-logits
26208    /// row exactly like the host variant (identical halves, identical concat) and runs the
26209    /// device argmax into `token_d` — NO host readback. Returns false when the split is
26210    /// ineligible (caller falls back to the plain matmul head).
26211    pub(crate) fn head_split_argmax_device(
26212        &self,
26213        e: &Engine,
26214        hn: &CudaSlice<f32>,
26215        token_d: &mut CudaSlice<u32>,
26216    ) -> Result<bool, Box<dyn std::error::Error>> {
26217        if self.head_split_fill_device(e, hn)?.is_none() {
26218            return Ok(false);
26219        }
26220        let n_vocab = self.cfg.n_vocab as usize;
26221        let guard = HEAD_SPLIT_WS
26222            .lock()
26223            .map_err(|_| "head split lock is poisoned")?;
26224        let ws = guard.as_ref().expect("filled above");
26225        let _main = e.gpu.enter_main()?;
26226        e.argmax_token_device_into(&ws.logits_e, token_d, n_vocab)?;
26227        Ok(true)
26228    }
26229
26230    /// SAMPLED twin of `head_split_argmax_device`. The split head already materializes the
26231    /// full concatenated row in `ws.logits_e`, so sampling does NOT have to give up HEAD_SPLIT
26232    /// — it draws from that row on device (filter thresholds, Gumbel perturbation, argmax)
26233    /// exactly as the serve tick does. Worth ~0.2 ms/token: the post-W8 census had the
26234    /// unsplit q8 head at ~364 us against ~82 us per half.
26235    pub(crate) fn head_split_sample_device(
26236        &self,
26237        e: &Engine,
26238        hn: &CudaSlice<f32>,
26239        token_d: &mut CudaSlice<u32>,
26240        samp: &crate::decode_batch::DevSamp,
26241        ctr: u32,
26242    ) -> Result<bool, Box<dyn std::error::Error>> {
26243        if self.head_split_fill_device(e, hn)?.is_none() {
26244            return Ok(false);
26245        }
26246        let n_vocab = self.cfg.n_vocab as usize;
26247        let guard = HEAD_SPLIT_WS
26248            .lock()
26249            .map_err(|_| "head split lock is poisoned")?;
26250        let mut guard = guard;
26251        let ws = guard.as_mut().expect("filled above");
26252        let _main = e.gpu.enter_main()?;
26253        if ws.samp.is_none() {
26254            ws.samp = Some(SampScratch {
26255                pb: e.zeros(n_vocab)?,
26256                th: e.zeros(1)?,
26257                z: e.zeros(1)?,
26258                mx: e.zeros(1)?,
26259                rows: e.htod_i32(&[0i32])?,
26260            });
26261        }
26262        let filtered = samp.top_k > 0 || samp.top_p < 1.0 || samp.min_p > 0.0;
26263        let HeadSplit {
26264            logits_e,
26265            samp: scratch,
26266            ..
26267        } = &mut *ws;
26268        let sc = scratch.as_mut().expect("armed above");
26269        if filtered {
26270            e.filter_stats(
26271                logits_e, n_vocab, &sc.rows, &mut sc.th, &mut sc.z, &mut sc.mx, n_vocab, 1,
26272                samp.temp, samp.top_k, samp.top_p, samp.min_p,
26273            )?;
26274            let SampScratch { pb, th, mx, .. } = sc;
26275            e.gumbel_perturb_filtered_col(
26276                logits_e, 0, pb, n_vocab, samp.seed, ctr, samp.temp, mx, th, 0,
26277            )?;
26278        } else {
26279            e.gumbel_perturb_col(logits_e, 0, &mut sc.pb, n_vocab, samp.seed, ctr, samp.temp)?;
26280        }
26281        e.argmax_token_device_col(&sc.pb, 0, n_vocab, token_d, 0)?;
26282        Ok(true)
26283    }
26284
26285    /// One-per-chunk host readback of the persistent full-logits row (the LAST chain
26286    /// token's row).
26287    pub(crate) fn head_split_logits_dtoh(
26288        &self,
26289        e: &Engine,
26290    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
26291        let guard = HEAD_SPLIT_WS
26292            .lock()
26293            .map_err(|_| "head split lock is poisoned")?;
26294        let ws = guard.as_ref().ok_or("head split logits not armed")?;
26295        let _main = e.gpu.enter_main()?;
26296        e.dtoh(&ws.logits_e)
26297    }
26298
26299    /// Chunk-loop replay (F-lite, MEMRA_STEP_TP_GRAPH_LOOP): run up to `k_target` greedy
26300    /// tokens as back-to-back graph launches chained through the in-graph tail argmax —
26301    /// ONE host sync, ONE history readback, and ONE bulk KV transaction per layer per
26302    /// chunk. Returns None when ineligible (caller falls back to the per-token path).
26303    /// The chunk consumes `token` (already emitted by the caller) as launch 0's input and
26304    /// returns the ids the chain argmax'd (hist[0..k]) plus the LAST launch's logits row —
26305    /// hist[k-1] is exactly argmax(logits), so the caller emits hist[..k-1] and lets its
26306    /// own loop re-derive hist[k-1] from the returned row.
26307    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
26308    pub fn step35_token_graph_chunk(
26309        &self,
26310        e: &Engine,
26311        token: u32,
26312        k_target: usize,
26313        cache: &mut Cache,
26314    ) -> Result<Option<(Vec<u32>, Vec<f32>)>, Box<dyn std::error::Error>> {
26315        if !self.uses_sliding_gated_moe_program()
26316            || !crate::tp::step_tp_graph_enabled()?
26317            || !crate::tp::step_tp_dcw_enabled()?
26318            || !crate::tp::step_tp_qkv_fused_enabled()?
26319            || !crate::tp::step_tp_dev_router_enabled()?
26320            || !crate::tp::step_nvfp4_dev_routes_enabled()?
26321        {
26322            return Ok(None);
26323        }
26324        // GRAPH-LAUNCH HEADROOM GUARD: same guard, same eager twin as
26325        // `step35_token_graph_step` (the chunk is that step replayed k times).
26326        if !crate::spec::graph_launch_headroom_ok(e) {
26327            static NOTED: std::sync::Once = std::sync::Once::new();
26328            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
26329            return Ok(None);
26330        }
26331        let n_layers = self.layers.len();
26332        let pos = cache.pos;
26333        let staged_next = pos + 1;
26334        if staged_next < 96 {
26335            return Ok(None);
26336        }
26337        // Bucket for the FIRST token; the chunk must not cross the bucket boundary (the
26338        // exec's n_splits ladder must match eager per depth).
26339        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
26340        if !fa_vec {
26341            return Ok(None);
26342        }
26343        let sp = crate::fa_split_keys(staged_next, 8);
26344        let bucket_max = (n_splits * sp).max(staged_next);
26345        let to_boundary = bucket_max.saturating_sub(staged_next) + 1;
26346        let mut k = k_target.min(to_boundary).min(16);
26347        if k < 2 {
26348            return Ok(None);
26349        }
26350        // Every layer must contiguous-append all k rows (no rebase inside the chunk).
26351        for il in 0..n_layers {
26352            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
26353                return Ok(None);
26354            };
26355            while k >= 2 && tp_kv.peek_append_ring(k)?.1 {
26356                k -= 1;
26357            }
26358            if k < 2 {
26359                return Ok(None);
26360            }
26361        }
26362
26363        let mut state_guard = self
26364            .step35_token_graph
26365            .lock()
26366            .map_err(|_| "step35 token graph lock is poisoned")?;
26367        let Some(state) = state_guard.as_mut() else {
26368            return Ok(None); // per-token path arms the state + stages first
26369        };
26370        if state.graphs.is_empty() {
26371            return Ok(None);
26372        }
26373        {
26374            let (b, g) = state.graphs.first_mut().expect("checked above");
26375            if *b != bucket_max {
26376                g.retarget_bucket(bucket_max)?;
26377                *b = bucket_max;
26378            }
26379        }
26380        let graph = state.graphs.first().map(|(_, g)| g).expect("checked above");
26381
26382        // Rank-stream fence (eager stragglers; see the per-token path).
26383        {
26384            let fa0 = match &self.layers[0].mixer {
26385                Mixer::Full(fa) => fa,
26386                _ => return Err("step35 token graph expects full-attention layers".into()),
26387            };
26388            let tp0 = fa0
26389                .step_tp_qkv
26390                .as_ref()
26391                .ok_or("step35 token graph lost its TP state")?;
26392            for rank in 0..tp0.runtime.devices().len() {
26393                let engine = tp0
26394                    .runtime
26395                    .rank_engine(rank)
26396                    .ok_or("step35 token graph lost a rank engine")?;
26397                let _main = engine.gpu.enter_main()?;
26398                engine.stream().synchronize()?;
26399            }
26400        }
26401
26402        // Seed the chain and fire k launches back-to-back: launch i embeds the token the
26403        // PREVIOUS launch's tail argmax wrote (launch 0 embeds the host-seeded `token`).
26404        {
26405            let _main = e.gpu.enter_main()?;
26406            e.set_u32_one(&mut state.token_d, token)?;
26407            e.set_i32_one(&mut state.pos_d, pos as i32)?;
26408            e.set_i32_one(&mut state.hist_idx, 0)?;
26409        }
26410        for _ in 0..k {
26411            graph.launch(e)?;
26412        }
26413        // Bulk host bookkeeping overlaps the replays: one k-row txn per layer.
26414        for il in 0..n_layers {
26415            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
26416            let transaction = tp_kv.begin_transaction()?;
26417            let fa = match &self.layers[il].mixer {
26418                Mixer::Full(fa) => fa,
26419                _ => return Err("step35 token graph expects full-attention layers".into()),
26420            };
26421            let tp = fa
26422                .step_tp_qkv
26423                .as_ref()
26424                .ok_or("step35 token graph lost its TP state")?;
26425            let empty: [CudaSlice<f32>; 0] = [];
26426            tp.runtime.append_tp_kv_transaction_inner(
26427                tp_kv,
26428                transaction,
26429                &empty,
26430                &empty,
26431                k,
26432                true,
26433            )?;
26434            tp.runtime
26435                .commit_tp_kv_transaction_external(tp_kv, transaction, k)?;
26436            if let Some(local) = cache.kv[il].as_mut() {
26437                local.len = pos + k;
26438                let _main = e.gpu.enter_main()?;
26439                e.set_i32_one(&mut local.len_d, (pos + k) as i32)?;
26440            }
26441        }
26442        cache.pos = pos + k;
26443        let (hist, logits) = {
26444            let _main = e.gpu.enter_main()?;
26445            e.stream().synchronize()?;
26446            (e.dtoh_u32(&state.token_hist)?, e.dtoh(&state.logits_stage)?)
26447        };
26448        Ok(Some((hist[..k].to_vec(), logits)))
26449    }
26450}
26451
26452impl HybridModel {
26453    /// Capture the whole-token parent for one fa bucket. Capture executes nothing; the
26454    /// section closures issue the SAME calls the eager dcw path runs (bit-proven), with the
26455    /// per-token operands living at the fixed stage addresses. Groups: the two rank sections
26456    /// of each phase fork in parallel and merge into the following root section.
26457    #[allow(clippy::too_many_arguments)]
26458    fn step35_token_graph_build(
26459        &self,
26460        e: &Engine,
26461        cache: &mut Cache,
26462        state: &mut Step35TokenGraphState,
26463        bucket_max: usize,
26464    ) -> Result<(), Box<dyn std::error::Error>> {
26465        use cudarc::driver::DevicePtr;
26466        let n_embd = self.cfg.n_embd as usize;
26467        let eps = self.cfg.rms_eps;
26468        let n_layers = self.layers.len();
26469        let started = std::time::Instant::now();
26470        if !crate::router_kernel_on() {
26471            return Err(
26472                "step35 token graph requires the router kernel (MEMRA_ROUTER_KERNEL=0)".into(),
26473            );
26474        }
26475        if !Engine::bf16_mmv_on() || !n_embd.is_multiple_of(8) {
26476            return Err("step35 token graph requires MEMRA_BF16_MMV bf16-resident matvecs".into());
26477        }
26478
26479        // Device embed table (the spec/graph lanes' lazily-uploaded copy).
26480        let embd_gpu = self
26481            .embd_gpu_try(e)
26482            .ok_or("step35 token graph could not upload the device embed table")?;
26483        let embd_qtype = match self.embd.ggml_type {
26484            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
26485            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
26486            other => return Err(format!("token graph embed dtype {other:?} unhandled").into()),
26487        };
26488        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
26489
26490        // Fixed-stage pointers the sections reference.
26491        let (p_mixed, p_kshadow, p_vshadow) = {
26492            let _main = e.gpu.enter_main()?;
26493            let stream = e.stream();
26494            let (a, _g) = state.mixed_stage.device_ptr(&stream);
26495            let (b, _g) = state.k_shadow_stage.device_ptr(&stream);
26496            let (c, _g) = state.v_shadow_stage.device_ptr(&stream);
26497            (a, b, c)
26498        };
26499
26500        crate::tp::token_graph_build_begin()?;
26501        let mut group_id: u32 = 0;
26502        for il in 0..n_layers {
26503            let layer = &self.layers[il];
26504            let fa = match &layer.mixer {
26505                Mixer::Full(fa) => fa,
26506                _ => return Err("step35 token graph expects full-attention layers".into()),
26507            };
26508            let tp = fa
26509                .step_tp_qkv
26510                .as_ref()
26511                .ok_or("step35 token graph lost its TP state")?;
26512            let attention = tp
26513                .attention
26514                .as_ref()
26515                .ok_or("step35 token graph lost its attention aux")?;
26516            let geometry = self.step35_geom(il);
26517            let window = geometry.window.map(|w| w as usize);
26518            let head_dim = geometry.head_dim_k as usize;
26519            let heads = geometry.n_head as usize;
26520            let kv_heads = geometry.n_head_kv as usize;
26521            let ranks = tp.runtime.devices().len();
26522            let local_heads = heads / ranks;
26523            let local_kv_heads = kv_heads / ranks;
26524            let layer_bucket = window.map(|w| bucket_max.min(w)).unwrap_or(bucket_max);
26525            let use_gate_shards =
26526                attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some();
26527            if !use_gate_shards {
26528                return Err("step35 token graph requires the fused gate shards".into());
26529            }
26530
26531            let ws_index = tp
26532                .runtime
26533                .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
26534            let ws_mutex = tp.runtime.decode_v2_workspace();
26535            let mut ws_guard = ws_mutex
26536                .lock()
26537                .map_err(|_| "step TP decode v2 workspace lock is poisoned")?;
26538            let ws = ws_guard
26539                .get_mut(ws_index)
26540                .ok_or("step TP decode v2 workspace missing after ensure")?;
26541            tp.runtime
26542                .decode_v2_arm_token_mirrors(ws, p_mixed, (p_kshadow, p_vshadow))?;
26543            let mut rope_freqs = Vec::with_capacity(ranks);
26544            for rank in 0..ranks {
26545                let engine = tp
26546                    .runtime
26547                    .rank_engine(rank)
26548                    .ok_or("step35 token graph lost a rank engine")?;
26549                rope_freqs.push(if geometry.rope_factors {
26550                    self.step35_aux
26551                        .as_ref()
26552                        .and_then(|aux| aux.rope_freqs(engine))
26553                } else {
26554                    None
26555                });
26556            }
26557            let gate_shards_arg = if let Some(shards) = attention.gate_shards.as_deref() {
26558                Some(crate::tp::StepTpGateShards::F32(shards))
26559            } else {
26560                attention
26561                    .gate_shards_bf16
26562                    .as_deref()
26563                    .map(crate::tp::StepTpGateShards::Bf16)
26564            };
26565
26566            // ---- E1: embed (layer 0) / attn norm into h_stage + pos copy ----
26567            let decode_input = attention
26568                .decode_input
26569                .as_ref()
26570                .ok_or("step35 token graph requires the replicated decode input")?;
26571            let mut decode_input = decode_input
26572                .lock()
26573                .map_err(|_| "replicated decode input lock is poisoned")?;
26574            // Stage arming happens through the eager stage flow once; require it here.
26575            if ws.h_stage.is_none() {
26576                return Err(
26577                    "step35 token graph requires the stage flow armed (run eager dcw first)".into(),
26578                );
26579            }
26580            {
26581                let state_x = &mut state.x;
26582                let token_d = &state.token_d;
26583                let pos_d = &state.pos_d;
26584                crate::tp::graph_section(e, None, || {
26585                    let _main = e.gpu.enter_main()?;
26586                    if il == 0 {
26587                        e.embed_gather_device_into(
26588                            embd_gpu,
26589                            token_d,
26590                            state_x,
26591                            n_embd,
26592                            embd_qtype,
26593                            embd_row_bytes,
26594                        )?;
26595                    }
26596                    {
26597                        let h_stage = ws.h_stage.as_mut().expect("stage armed checked above");
26598                        e.rms_norm(
26599                            state_x,
26600                            layer.attn_norm.float_data(),
26601                            h_stage,
26602                            n_embd,
26603                            1,
26604                            eps,
26605                        )?;
26606                    }
26607                    {
26608                        let pos_stage = ws.pos_stage.as_mut().expect("stage armed above");
26609                        let mut dst = pos_stage.slice_mut(0..1);
26610                        e.stream().memcpy_dtod(&pos_d.slice(0..1), &mut dst)?;
26611                    }
26612                    Ok(())
26613                })?;
26614            }
26615
26616            // ---- R0/R1 (parallel): projections + dcw attention interior ----
26617            group_id += 1;
26618            for rank in 0..ranks {
26619                let engine = tp
26620                    .runtime
26621                    .rank_engine(rank)
26622                    .ok_or("step35 token graph lost a rank engine")?;
26623                {
26624                    // fa partial pool must reach the RUN CEILING before capture — an
26625                    // in-capture grow is a mem node (child graphs reject those), and the
26626                    // retarget path (increment C) widens the baked memsets up to the ceiling
26627                    // without moving the pool pointers. Two ensures cover both sp rungs.
26628                    let ceiling = window
26629                        .map(|w| cache.max_ctx.min(w))
26630                        .unwrap_or(cache.max_ctx);
26631                    let _main = engine.gpu.enter_main()?;
26632                    engine.fa_dcw_pool_ensure(
26633                        head_dim,
26634                        local_heads,
26635                        local_kv_heads,
26636                        ceiling.min(2048),
26637                    )?;
26638                    engine.fa_dcw_pool_ensure(head_dim, local_heads, local_kv_heads, ceiling)?;
26639                    engine.fa_dcw_pool_ensure(
26640                        head_dim,
26641                        local_heads,
26642                        local_kv_heads,
26643                        layer_bucket,
26644                    )?;
26645                }
26646                let runtime = &tp.runtime;
26647                let q_norm = &attention.q_norm;
26648                let k_norm = &attention.k_norm;
26649                let gate_ref = gate_shards_arg.as_ref();
26650                crate::tp::graph_section(engine, Some(group_id), || {
26651                    runtime.decode_v2_input_qkv_rank(
26652                        ws,
26653                        &state.pos_d,
26654                        &mut decode_input,
26655                        &tp.q,
26656                        &tp.k,
26657                        &tp.v,
26658                        q_norm,
26659                        k_norm,
26660                        head_dim,
26661                        geometry.n_rot as usize,
26662                        geometry.rope_base,
26663                        &rope_freqs,
26664                        eps,
26665                        gate_ref,
26666                        true,
26667                        true,
26668                        false,
26669                        rank,
26670                        None,
26671                    )?;
26672                    // Merged dcw interior at the BUCKET geometry (one-partition law makes the
26673                    // replayed values track the live counters).
26674                    let distributed = cache.tp_kv[il]
26675                        .as_mut()
26676                        .ok_or("step35 token graph lost a TP cache")?;
26677                    let (kv_dim_k, kv_dim_v) = (distributed.kv_dim_k(), distributed.kv_dim_v());
26678                    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
26679                    let capacity = distributed.physical_capacity();
26680                    {
26681                        let rank_cache = distributed
26682                            .rank_mut(rank)
26683                            .ok_or("step35 token graph lost a rank cache")?;
26684                        let (k_plane, v_plane, len_d, base_d) =
26685                            rank_cache.planes_and_counters_mut();
26686                        engine.append_kv_quantized_dcw(
26687                            &ws.k[rank],
26688                            &ws.v_raw[rank],
26689                            k_plane,
26690                            v_plane,
26691                            len_d,
26692                            base_d,
26693                            kv_dim_k,
26694                            kv_dim_v,
26695                            ktb,
26696                            vtb,
26697                        )?;
26698                    }
26699                    {
26700                        let rank_cache = distributed
26701                            .rank_mut(rank)
26702                            .ok_or("step35 token graph lost a rank cache")?;
26703                        engine.inc_i32(rank_cache.len_d_mut())?;
26704                    }
26705                    let rank_cache = distributed
26706                        .rank(rank)
26707                        .ok_or("step35 token graph lost a rank cache")?;
26708                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * ktb);
26709                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * vtb);
26710                    // Graph build keeps the UNFUSED combine + gate pair: the bucket
26711                    // retarget addresses combine's nsp at arg slot 6, and the fused
26712                    // combine_gate kernel shifts it to 7 — bit-identical either way, so
26713                    // only the eager arm takes FUSION #2d.
26714                    engine.fa_decode_dcw(
26715                        &ws.q[rank],
26716                        &k_ring,
26717                        &v_ring,
26718                        &mut ws.attn_out[rank],
26719                        head_dim,
26720                        local_heads,
26721                        local_kv_heads,
26722                        rank_cache.len_d(),
26723                        rank_cache.base_d(),
26724                        window.unwrap_or(0),
26725                        layer_bucket,
26726                        geometry.attention_scale(),
26727                        ktb,
26728                        vtb,
26729                        None,
26730                    )?;
26731                    engine.attn_head_gate(
26732                        &ws.attn_out[rank],
26733                        &ws.gate[rank],
26734                        &mut ws.gated[rank],
26735                        None,
26736                        head_dim,
26737                        local_heads,
26738                        1,
26739                    )?;
26740                    runtime.decode_v2_finish_rank_partial(ws, &tp.o, true, rank)?;
26741                    Ok(())
26742                })?;
26743            }
26744
26745            // ---- ROOT: combine + shadows + e-mirrors ----
26746            {
26747                let root = tp
26748                    .runtime
26749                    .rank_engine(0)
26750                    .ok_or("step35 token graph lost the root engine")?;
26751                let runtime = &tp.runtime;
26752                crate::tp::graph_section(root, None, || runtime.decode_v2_finish_root_fused(ws))?;
26753            }
26754            drop(ws_guard);
26755            drop(decode_input);
26756
26757            let probe_layer: Option<usize> = std::env::var("MEMRA_TG_PROBE_LAYER")
26758                .ok()
26759                .and_then(|v| v.parse().ok());
26760            if probe_layer == Some(il) {
26761                let Step35TokenGraphState {
26762                    mixed_stage,
26763                    probe_mixed,
26764                    ..
26765                } = &mut *state;
26766                crate::tp::graph_section(e, None, || {
26767                    let _main = e.gpu.enter_main()?;
26768                    let mut dst = probe_mixed.slice_mut(0..n_embd);
26769                    e.stream()
26770                        .memcpy_dtod(&mixed_stage.slice(0..n_embd), &mut dst)?;
26771                    Ok(())
26772                })?;
26773            }
26774
26775            // ---- FFN half ----
26776            match &layer.ffn {
26777                crate::hybrid::Ffn::Dense {
26778                    ffn_gate,
26779                    ffn_up,
26780                    ffn_down,
26781                } => {
26782                    let n_ff = ffn_gate.out_features();
26783                    let lim = self.cfg.clamp_shexp_at(il as u32);
26784                    // Alloc-free inline of ffn_swiglu_decode's bf16 tail: dual gate/up matvec is
26785                    // bit-identical per row to the two matmul-dispatched matvec_bf16 launches.
26786                    // A clamped dense layer would take eager's q8_1 branch instead -- refuse.
26787                    if lim.is_some() {
26788                        return Err("step35 token graph dense FFN with clamp unsupported".into());
26789                    }
26790                    let (wg_d, wu_d, wd_d) = match (ffn_gate, ffn_up, ffn_down) {
26791                        (
26792                            crate::model::GpuTensor::FloatBf16 { data: wg, .. },
26793                            crate::model::GpuTensor::FloatBf16 { data: wu, .. },
26794                            crate::model::GpuTensor::FloatBf16 { data: wd, .. },
26795                        ) => (wg, wu, wd),
26796                        _ => {
26797                            return Err(
26798                                "step35 token graph dense FFN requires bf16-resident weights"
26799                                    .into(),
26800                            );
26801                        }
26802                    };
26803                    crate::tp::graph_section(e, None, || {
26804                        let _main = e.gpu.enter_main()?;
26805                        let Step35TokenGraphState {
26806                            x,
26807                            x1,
26808                            mixed_stage,
26809                            dense_z,
26810                            dense_gate,
26811                            dense_up,
26812                            dense_act,
26813                            sh_stage,
26814                            ..
26815                        } = &mut *state;
26816                        e.add_rms_norm(
26817                            x,
26818                            mixed_stage,
26819                            layer.post_attn_norm.float_data(),
26820                            x1,
26821                            dense_z,
26822                            n_embd,
26823                            1,
26824                            eps,
26825                        )?;
26826                        // TWO SINGLE matvecs, not the dual: eager dense rides two
26827                        // matmul-dispatched matvec_bf16 launches; the dual twin measured a
26828                        // ~2e-9 residual difference here (token-graph bisection, 2026-08-21).
26829                        e.matvec_bf16_into(wg_d, dense_z, dense_gate, n_embd, n_ff)?;
26830                        e.matvec_bf16_into(wu_d, dense_z, dense_up, n_embd, n_ff)?;
26831                        Self::ffn_act_lim(
26832                            e, &self.cfg, dense_gate, dense_up, 1.0, 1.0, lim, dense_act, n_ff,
26833                        )?;
26834                        e.matvec_bf16_into(wd_d, dense_act, sh_stage, n_ff, n_embd)?;
26835                        e.add(x1, sh_stage, x, n_embd)?;
26836                        Ok(())
26837                    })?;
26838                }
26839                crate::hybrid::Ffn::Moe(m) => {
26840                    let moe = self
26841                        .cfg
26842                        .moe
26843                        .as_ref()
26844                        .ok_or("step35 token graph needs moe cfg")?;
26845                    let n_expert = moe.expert_count as usize;
26846                    let n_used = moe.expert_used_count as usize;
26847                    let sigmoid = self
26848                        .cfg
26849                        .sigmoid_router()
26850                        .ok_or("step35 token graph needs the sigmoid router")?;
26851                    let step_tp = m
26852                        .step_tp
26853                        .as_ref()
26854                        .ok_or("step35 token graph needs TP experts")?;
26855                    let bank = match &step_tp.experts {
26856                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => bank,
26857                        _ => return Err("step35 token graph needs the NVFP4 bank".into()),
26858                    };
26859                    let routes_ws_mutex = bank.device_workspace_handle();
26860                    let mut routes_guard = routes_ws_mutex
26861                        .lock()
26862                        .map_err(|_| "routes workspace lock is poisoned")?;
26863                    let routes_ws = routes_guard
26864                        .as_mut()
26865                        .ok_or("step35 token graph requires the routes workspace warmed")?;
26866                    routes_ws.arm_stages(e, bank.input_width, n_used)?;
26867                    step_tp.runtime.routes_arm_raw(bank, routes_ws)?;
26868                    let p_z = {
26869                        let root = step_tp
26870                            .runtime
26871                            .rank_engine(0)
26872                            .ok_or("routes root engine missing")?;
26873                        let _main = root.gpu.enter_main()?;
26874                        let stream = root.stream();
26875                        let in_stage = routes_ws
26876                            .in_stage_handle()
26877                            .ok_or("routes in stage not armed")?;
26878                        let (a, _g) = in_stage.device_ptr(&stream);
26879                        a
26880                    };
26881                    let local_out = bank.expert_width / ranks;
26882
26883                    // ---- E2: post-attn norm into the routes in-stage + router + staging ----
26884                    crate::tp::graph_section(e, None, || {
26885                        let _main = e.gpu.enter_main()?;
26886                        {
26887                            let in_stage = routes_ws
26888                                .in_stage_mut()
26889                                .ok_or("routes in stage not armed")?;
26890                            let Step35TokenGraphState {
26891                                x, x1, mixed_stage, ..
26892                            } = &mut *state;
26893                            e.add_rms_norm(
26894                                x,
26895                                mixed_stage,
26896                                layer.post_attn_norm.float_data(),
26897                                x1,
26898                                in_stage,
26899                                n_embd,
26900                                1,
26901                                eps,
26902                            )?;
26903                        }
26904                        {
26905                            let z_ref = routes_ws
26906                                .in_stage_handle()
26907                                .ok_or("routes in stage not armed")?;
26908                            e.router_gemv_into(
26909                                m.gate_inp.float_data(),
26910                                z_ref,
26911                                &mut state.router_logits,
26912                                n_embd,
26913                                n_expert,
26914                                1,
26915                            )?;
26916                        }
26917                        let (sel_e, w_e) = routes_ws
26918                            .dev_route_e_mut()
26919                            .ok_or("routes staging not armed")?;
26920                        e.moe_router_sigmoid_topk_into(
26921                            &state.router_logits,
26922                            1,
26923                            n_expert,
26924                            n_used,
26925                            m.active_count(),
26926                            &m.exp_probs_b_dev,
26927                            &m.active_experts_dev,
26928                            sigmoid.0,
26929                            sigmoid.1,
26930                            sel_e,
26931                            w_e,
26932                        )?;
26933                        Ok(())
26934                    })?;
26935
26936                    // ---- R0r/R1r (parallel): routes sweeps ----
26937                    group_id += 1;
26938                    for rank in 0..ranks {
26939                        let engine = step_tp
26940                            .runtime
26941                            .rank_engine(rank)
26942                            .ok_or("routes rank engine missing")?;
26943                        let runtime = &step_tp.runtime;
26944                        crate::tp::graph_section(engine, Some(group_id), || {
26945                            runtime.routes_rank_section(
26946                                bank,
26947                                routes_ws,
26948                                p_z,
26949                                local_out,
26950                                n_used,
26951                                step_tp.activation_limit,
26952                                rank,
26953                            )
26954                        })?;
26955                    }
26956
26957                    // ---- ROOTr: combine into the out stage ----
26958                    {
26959                        let root = step_tp
26960                            .runtime
26961                            .rank_engine(0)
26962                            .ok_or("routes root engine missing")?;
26963                        let runtime = &step_tp.runtime;
26964                        crate::tp::graph_section(root, None, || {
26965                            runtime.routes_root_section(bank, routes_ws)
26966                        })?;
26967                    }
26968
26969                    // ---- E3: shexp + add_shared onto the out stage + residual ----
26970                    // Alloc-free inline of moe_ffn_grouped_add_shared's bf16_dual arm (the arm
26971                    // eager takes under MEMRA_BF16_MMV, guarded at fn entry).
26972                    let lim_sh = self.cfg.clamp_shexp_at(il as u32);
26973                    let (wg_sh, wu_sh, wd_sh) = match (&m.gate_shexp, &m.up_shexp, &m.down_shexp) {
26974                        (
26975                            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
26976                            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
26977                            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
26978                        ) => (wg, wu, wd),
26979                        _ => {
26980                            return Err(
26981                                "step35 token graph shexp requires bf16-resident weights".into()
26982                            );
26983                        }
26984                    };
26985                    let n_ff_sh = m
26986                        .gate_shexp
26987                        .as_ref()
26988                        .expect("matched Some above")
26989                        .out_features();
26990                    // No sigmoid gate on this model's shexp -> gate_sig stays at its 1.0
26991                    // init, reproducing eager's ones vector without a launch.
26992                    let gate_inp_shexp = m.gate_inp_shexp.as_ref();
26993                    crate::tp::graph_section(e, None, || {
26994                        let _main = e.gpu.enter_main()?;
26995                        let (z_ref, out_stage) = routes_ws
26996                            .in_and_out_stages_mut()
26997                            .ok_or("routes stages not armed")?;
26998                        let Step35TokenGraphState {
26999                            x,
27000                            x1,
27001                            sh_stage,
27002                            shexp_gate,
27003                            shexp_up,
27004                            shexp_act,
27005                            gate_sig,
27006                            ..
27007                        } = &mut *state;
27008                        e.matvec_bf16_dual_into(
27009                            wg_sh, wu_sh, z_ref, shexp_gate, shexp_up, n_embd, n_ff_sh,
27010                        )?;
27011                        Self::ffn_act_lim(
27012                            e, &self.cfg, shexp_gate, shexp_up, 1.0, 1.0, lim_sh, shexp_act,
27013                            n_ff_sh,
27014                        )?;
27015                        e.matvec_bf16_into(wd_sh, shexp_act, sh_stage, n_ff_sh, n_embd)?;
27016                        if let Some(gate_w) = gate_inp_shexp {
27017                            e.sigmoid_dot_rows_into(
27018                                z_ref,
27019                                gate_w.float_data(),
27020                                gate_sig,
27021                                n_embd,
27022                                1,
27023                            )?;
27024                        }
27025                        e.add_scaled_rows(sh_stage, gate_sig, out_stage, n_embd, 1)?;
27026                        e.add(x1, out_stage, x, n_embd)?;
27027                        Ok(())
27028                    })?;
27029                }
27030            }
27031            if probe_layer == Some(il) {
27032                let Step35TokenGraphState { x, probe_x, .. } = &mut *state;
27033                crate::tp::graph_section(e, None, || {
27034                    let _main = e.gpu.enter_main()?;
27035                    let mut dst = probe_x.slice_mut(0..n_embd);
27036                    e.stream().memcpy_dtod(&x.slice(0..n_embd), &mut dst)?;
27037                    Ok(())
27038                })?;
27039            }
27040        }
27041
27042        // ---- Tail: output norm + head into the logits stage ----
27043        let head = match &self.output {
27044            crate::model::GpuTensor::FloatBf16 { data, .. } => data,
27045            _ => return Err("step35 token graph head requires the bf16-resident output".into()),
27046        };
27047        crate::tp::graph_section(e, None, || {
27048            let _main = e.gpu.enter_main()?;
27049            let Step35TokenGraphState {
27050                x,
27051                hn,
27052                logits_stage,
27053                token_d,
27054                pos_d,
27055                token_hist,
27056                hist_idx,
27057                ..
27058            } = &mut *state;
27059            e.rms_norm(x, self.output_norm.float_data(), hn, n_embd, 1, eps)?;
27060            e.matvec_bf16_into(head, hn, logits_stage, n_embd, self.cfg.n_vocab as usize)?;
27061            // Chunk-loop tail: greedy argmax feeds token_d (host-identical tie-break,
27062            // argmax_gate-validated), the id lands in the history ring, and pos advances on
27063            // device — consecutive launches chain with NO host sync. Single-token mode
27064            // overwrites token_d/pos_d from the host before each launch, so these nodes are
27065            // harmless there.
27066            e.argmax_token_device_into(logits_stage, token_d, self.cfg.n_vocab as usize)?;
27067            e.u32_hist_append(token_d, token_hist, hist_idx)?;
27068            e.inc_i32(pos_d)?;
27069            Ok(())
27070        })?;
27071
27072        let graph = crate::tp::token_graph_build_finish()?;
27073        state.graphs.push((bucket_max, graph));
27074        eprintln!(
27075            "[step35-token-graph] built bucket={bucket_max} layers={n_layers} \
27076             build_ms={:.0} performance_claim=false",
27077            started.elapsed().as_secs_f64() * 1e3
27078        );
27079        Ok(())
27080    }
27081}