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
5use crate::Engine;
6use crate::cache::Cache;
7use cudarc::driver::CudaSlice;
8use memra_gguf::config::{ModelConfig, SwigluClamp};
9
10/// Resident trunk transients for the eager prime (piecewise-graph foundation; see
11/// HybridModel::prime_slabs). Every live buffer prefix is fully overwritten before use per prime;
12/// capacity beyond the current token count must never cross a shape-sensitive boundary.
13pub struct PrimeSlabs {
14    pub t_cap: usize,
15    pub h: CudaSlice<f32>,
16    pub x1: CudaSlice<f32>,
17    pub z: CudaSlice<f32>,
18    pub act: CudaSlice<f32>,
19    pub xa: CudaSlice<f32>,
20    pub xb: CudaSlice<f32>,
21    pub h16: CudaSlice<u8>,
22    pub z16: CudaSlice<u8>,
23    /// piecewise boundary slabs (increment 2): GEMM outputs land here so the
24    /// downstream captured segments see fixed addresses.
25    pub gate: CudaSlice<f32>, // t * n_ff_max
26    pub up: CudaSlice<f32>,      // t * n_ff_max
27    pub ffn_out: CudaSlice<f32>, // t * n_embd
28    /// piecewise increment 3: per-layer S-glue segment graphs (down-add + next
29    /// attn-norm, ALL-slab IO, zero in-graph allocations -> keeperless capture is
30    /// clean). Baked at this t_cap; replay only when t == t_cap. seg_glue[il] fires
31    /// between layer il and il+1 (ping-pong parity is deterministic per il).
32    pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
33    /// increment 5 (core-split edition): the mixer out-GEMM writes _into_ `mixed`
34    /// directly (no staging copy — the increment-4 copy route was refuted), making
35    /// S-mid [add + post-norm] all-slab and capturable.
36    pub mixed: CudaSlice<f32>,
37    pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
38    pub seg_t: usize,
39}
40
41// Split prime ranges cannot enter the full-range segment-graph arm, and every slab access
42// is serialized by its device mutex after binding that device's CUDA context on the thread.
43unsafe impl Send for PrimeSlabs {}
44
45/// Shared-expert gate+up at t==1: NVFP4 fused2 (the ornith15/qwen35moe NVFP4 mints keep
46/// gate/up_shexp uniformly NVFP4, so the Q8-only fused2 never fired there and the pair fell
47/// to two mr2 singles + two re-quantizes of the same z — 2 of the 8 unfused launches/layer
48/// the orndecode B=1 census ranked at 17.1%), else the Q8_0 fused2 (the Q8 35B mint), else
49/// two singles. ONE helper for all three shexp dispatch sites — the MEMRA_GDN_MMA
50/// three-read-sites defect is the precedent for not inlining this thrice. Fusion law
51/// everywhere: per (tensor,row) the fused seg body is verbatim, so fused == singles
52/// bit-identically, and the shared (zq, zd) is the same quantize each single recomputes.
53fn shexp_gate_up_t1(
54    e: &Engine,
55    gate_shexp: &crate::model::GpuTensor,
56    up_shexp: &crate::model::GpuTensor,
57    z: &CudaSlice<f32>,
58    zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
59) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
60    let is_nvfp4 = |w: &crate::model::GpuTensor| matches!(w, crate::model::GpuTensor::Quant { qtype, .. } if *qtype == crate::QT_NVFP4);
61    if is_nvfp4(gate_shexp) && is_nvfp4(up_shexp) {
62        // Reuse the caller's t==1 z-quantize when one exists (the zq8 seam the dev arm
63        // already consumes) — the helper's own quantize is the identical kernel on the
64        // identical input, so this drops one launch per MoE layer without moving a byte.
65        let pair = match zq8 {
66            Some((zq, zd)) => e.matmul_nvfp4_fused2(gate_shexp, up_shexp, zq, zd, 1)?,
67            None => {
68                let (zq, zd) = e.quantize_q8_1(z, 1, gate_shexp.in_features())?;
69                e.matmul_nvfp4_fused2(gate_shexp, up_shexp, &zq, &zd, 1)?
70            }
71        };
72        if let Some(pair) = pair {
73            return Ok(pair);
74        }
75    }
76    match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
77        Some(pair) => Ok(pair),
78        None => Ok((e.matmul(gate_shexp, z, 1)?, e.matmul(up_shexp, z, 1)?)),
79    }
80}
81
82fn active_matrix_values(
83    available: usize,
84    rows: usize,
85    columns: usize,
86    label: &str,
87) -> Result<usize, String> {
88    let required = rows
89        .checked_mul(columns)
90        .ok_or_else(|| format!("{label} shape overflows: {rows}x{columns}"))?;
91    if available < required {
92        return Err(format!(
93            "{label} has {available} values, fewer than the active {rows}x{columns} ({required})"
94        ));
95    }
96    Ok(required)
97}
98
99fn step_grouped_decode_shape(prefill: bool, tokens: usize) -> bool {
100    !prefill && tokens == 1
101}
102
103fn parse_step_ep_grouped_prefill(value: Option<&str>) -> Result<bool, String> {
104    match value {
105        None | Some("") | Some("0") => Ok(false),
106        Some("1") => Ok(true),
107        Some(value) => Err(format!(
108            "MEMRA_STEP_EP_GROUPED_PREFILL={value:?} is invalid; expected 0 or 1"
109        )),
110    }
111}
112
113fn step_ep_grouped_prefill_enabled() -> Result<bool, String> {
114    parse_step_ep_grouped_prefill(
115        std::env::var("MEMRA_STEP_EP_GROUPED_PREFILL")
116            .ok()
117            .as_deref(),
118    )
119}
120
121fn step_grouped_prefill_shape(enabled: bool, prefill: bool, tokens: usize) -> bool {
122    enabled && prefill && (PRIME_MIN_T..=crate::cache::PRIME_CHUNK_MAX_TOKENS).contains(&tokens)
123}
124
125fn parse_step_tp_prefill(value: Option<&str>) -> Result<bool, String> {
126    match value {
127        None | Some("") | Some("0") => Ok(false),
128        Some("1") => Ok(true),
129        Some(value) => Err(format!(
130            "MEMRA_STEP_TP_PREFILL={value:?} is invalid; expected 0 or 1"
131        )),
132    }
133}
134
135fn step_tp_prefill_enabled() -> Result<bool, String> {
136    parse_step_tp_prefill(std::env::var("MEMRA_STEP_TP_PREFILL").ok().as_deref())
137}
138
139fn validate_step_prime_batch_modes(tp_prefill: bool, grouped_prefill: bool) -> Result<(), String> {
140    if grouped_prefill && !tp_prefill {
141        return Err("MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into());
142    }
143    if tp_prefill {
144        return Err(
145            "Step TP4 cross-request prime batching did not clear the live-server performance \
146             gate; use per-session grouped prefill"
147                .into(),
148        );
149    }
150    Ok(())
151}
152
153fn step_tp_prefill_shape(
154    enabled: bool,
155    tokens: usize,
156    ranks: usize,
157    native_p2p: bool,
158    has_rank_local_attention: bool,
159    fp8_kv: bool,
160) -> bool {
161    // TP2 admitted 2026-08-25 behind the same off-by-default door. The prefill body is
162    // rank-count-generic (every geometry check divides by `ranks`); only this shape ever
163    // named 4. TP4's high-context NO-GO was a TRANSPORT verdict — token-row column
164    // gathers plus remote O blocks issue ~61,440 peer copies per 4K attention layer on
165    // that placement — which is not evidence about a 2-card native-P2P placement that
166    // reduces O rank-locally. TP2 is UNQUALIFIED until its own prefill argmax + TTFT
167    // receipts land; the door stays off by default.
168    enabled
169        && tokens >= PRIME_MIN_T
170        && matches!(ranks, 2 | 4)
171        && native_p2p
172        && has_rank_local_attention
173        && !fp8_kv
174}
175
176fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
177    std::iter::repeat_with(|| None).take(n).collect()
178}
179
180fn prime_cache_stage_for_layer(fence: &[usize], layer: usize) -> usize {
181    debug_assert!(fence.len() >= 3);
182    match fence[1..fence.len() - 1].binary_search(&layer) {
183        Ok(index) => index + 1,
184        Err(index) => index,
185    }
186}
187
188fn move_prime_cache_layers<T>(
189    parent: &mut [Option<T>],
190    stages: &mut [Vec<Option<T>>],
191    fence: &[usize],
192) {
193    assert_eq!(stages.len() + 1, fence.len());
194    assert!(stages.iter().all(|stage| stage.len() == parent.len()));
195    for (layer, value) in parent.iter_mut().enumerate() {
196        let stage = prime_cache_stage_for_layer(fence, layer);
197        debug_assert!(stages[stage][layer].is_none());
198        stages[stage][layer] = value.take();
199    }
200}
201
202#[cfg(test)]
203fn restore_prime_cache_layers<T>(
204    parent: &mut [Option<T>],
205    stages: &mut [Vec<Option<T>>],
206    fence: &[usize],
207) {
208    assert_eq!(stages.len() + 1, fence.len());
209    assert!(stages.iter().all(|stage| stage.len() == parent.len()));
210    for (layer, value) in parent.iter_mut().enumerate() {
211        let stage = prime_cache_stage_for_layer(fence, layer);
212        debug_assert!(value.is_none());
213        *value = stages[stage][layer].take();
214    }
215}
216
217/// Temporarily move a PP cache's layer state into independently-owned stage shells. The stage
218/// walkers then receive disjoint `&mut Cache` values and can run on separate host threads without
219/// aliasing. GPU buffers are moved, not copied; Drop restores every layer and publishes the last
220/// position completed by every stage.
221struct PrimeCacheStages<'a> {
222    parent: &'a mut Cache,
223    fence: Vec<usize>,
224    stages: Vec<std::sync::Mutex<Cache>>,
225    committed: bool,
226}
227
228impl<'a> PrimeCacheStages<'a> {
229    fn new(parent: &'a mut Cache, fence: &[usize]) -> Self {
230        let n = parent.kv.len();
231        assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
232        assert_eq!(parent.tp_kv.len(), n, "cache layer vectors disagree");
233        assert_eq!(parent.latent.len(), n, "cache layer vectors disagree");
234        let n_stages = fence.len().checked_sub(1).expect("PP cache fence is empty");
235        assert!((2..=4).contains(&n_stages), "PP cache needs 2..=4 stages");
236        assert_eq!(fence[0], 0, "PP cache fence must start at layer zero");
237        assert!(
238            fence.windows(2).all(|pair| pair[0] < pair[1]),
239            "PP cache fence must be strictly increasing"
240        );
241        assert!(fence[n_stages] <= n, "PP cache fence exceeds {n} layers");
242
243        let mut latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
244        let mut g5_recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
245        let mut g5_latent: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
246        let mut kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
247        let mut tp_kv: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
248        let mut recur: Vec<_> = (0..n_stages).map(|_| empty_cache_layers(n)).collect();
249        move_prime_cache_layers(&mut parent.kv, &mut kv, fence);
250        move_prime_cache_layers(&mut parent.tp_kv, &mut tp_kv, fence);
251        move_prime_cache_layers(&mut parent.recur, &mut recur, fence);
252
253        move_prime_cache_layers(&mut parent.latent, &mut latent, fence);
254        move_prime_cache_layers(&mut parent.glm5_tp_recur, &mut g5_recur, fence);
255        move_prime_cache_layers(&mut parent.glm5_tp_latent_peer, &mut g5_latent, fence);
256        let pos = parent.pos;
257        let max_ctx = parent.max_ctx;
258        // Indexed rather than zipped: six per-stage layer vectors (main's kv/tp_kv/recur plus
259        // this lane's latent/glm5_tp_recur/glm5_tp_latent_peer) do not read as a zip chain, and
260        // a nested-tuple pattern is exactly where a field silently lands on the wrong stage.
261        let stages = (0..n_stages)
262            .map(|stage| {
263                std::sync::Mutex::new(Cache {
264                    kv: std::mem::take(&mut kv[stage]),
265                    tp_kv: std::mem::take(&mut tp_kv[stage]),
266                    recur: std::mem::take(&mut recur[stage]),
267                    latent: std::mem::take(&mut latent[stage]),
268                    glm5_tp_recur: std::mem::take(&mut g5_recur[stage]),
269                    glm5_tp_latent_peer: std::mem::take(&mut g5_latent[stage]),
270                    pos,
271                    max_ctx,
272                    tainted: false,
273                    last_logits_dev: None,
274                    dflash_taps: None,
275                    hc_taps: None,
276                })
277            })
278            .collect();
279        Self {
280            parent,
281            fence: fence.to_vec(),
282            stages,
283            committed: false,
284        }
285    }
286
287    fn pp2_parts(&mut self) -> (&mut Cache, &mut Cache) {
288        assert_eq!(self.stages.len(), 2);
289        let (stage0, stage1) = self.stages.split_at_mut(1);
290        (
291            stage0[0]
292                .get_mut()
293                .unwrap_or_else(|poisoned| poisoned.into_inner()),
294            stage1[0]
295                .get_mut()
296                .unwrap_or_else(|poisoned| poisoned.into_inner()),
297        )
298    }
299
300    fn stages(&self) -> &[std::sync::Mutex<Cache>] {
301        &self.stages
302    }
303
304    fn commit(&mut self) {
305        self.committed = true;
306    }
307}
308
309impl Drop for PrimeCacheStages<'_> {
310    fn drop(&mut self) {
311        let n = self.parent.kv.len();
312        for i in 0..n {
313            let stage = prime_cache_stage_for_layer(&self.fence, i);
314            let source = self.stages[stage]
315                .get_mut()
316                .unwrap_or_else(|poisoned| poisoned.into_inner());
317            debug_assert!(self.parent.kv[i].is_none());
318            debug_assert!(self.parent.tp_kv[i].is_none());
319            debug_assert!(self.parent.recur[i].is_none());
320            debug_assert!(self.parent.latent[i].is_none());
321            self.parent.kv[i] = source.kv[i].take();
322            self.parent.tp_kv[i] = source.tp_kv[i].take();
323            self.parent.recur[i] = source.recur[i].take();
324            self.parent.latent[i] = source.latent[i].take();
325            self.parent.glm5_tp_recur[i] = source.glm5_tp_recur[i].take();
326            self.parent.glm5_tp_latent_peer[i] = source.glm5_tp_latent_peer[i].take();
327        }
328        self.parent.pos = self
329            .stages
330            .iter_mut()
331            .map(|stage| {
332                stage
333                    .get_mut()
334                    .unwrap_or_else(|poisoned| poisoned.into_inner())
335                    .pos
336            })
337            .min()
338            .unwrap_or(self.parent.pos);
339        if !self.committed {
340            self.parent.mark_tainted();
341        }
342    }
343}
344
345/// Fail-stop transaction marker for concat-prime paths. These paths mutate several independent
346/// caches before their final epilogue can fail; an error must make every member permanently
347/// ineligible for retry/reuse rather than replaying a queue over partially advanced state.
348struct CacheTaintGuard {
349    caches: Vec<*mut Cache>,
350    committed: bool,
351}
352
353impl CacheTaintGuard {
354    fn arm(caches: &mut [&mut Cache]) -> Self {
355        Self {
356            caches: caches
357                .iter_mut()
358                .map(|cache| *cache as *mut Cache)
359                .collect(),
360            committed: false,
361        }
362    }
363
364    fn commit(&mut self) {
365        self.committed = true;
366    }
367}
368
369impl Drop for CacheTaintGuard {
370    fn drop(&mut self) {
371        if self.committed {
372            return;
373        }
374        for cache in &self.caches {
375            // SAFETY: `arm` receives the function's unique cache references. The guard never
376            // escapes that call or dereferences them until unwind/return after active borrows end.
377            unsafe { (&mut **cache).mark_tainted() };
378        }
379    }
380}
381
382#[derive(Debug, Clone, Copy, PartialEq, Eq)]
383struct PrimePpWaveSlot {
384    wave: usize,
385    slot: usize,
386}
387
388#[derive(Debug)]
389enum PrimePpSignal {
390    Slot(PrimePpWaveSlot),
391    Error(String),
392}
393
394#[derive(Default)]
395struct PrimePpWaveCredits {
396    next_wave: usize,
397    pending: std::collections::VecDeque<PrimePpWaveSlot>,
398}
399
400impl PrimePpWaveCredits {
401    fn release_required(&self) -> Option<PrimePpWaveSlot> {
402        (self.pending.len() == 2).then(|| self.pending[0])
403    }
404
405    fn record_release(&mut self, released: PrimePpWaveSlot) -> Result<(), String> {
406        let expected =
407            self.pending.front().copied().ok_or_else(|| {
408                "prime PP received a slot release with no pending wave".to_string()
409            })?;
410        if released != expected {
411            return Err(format!(
412                "prime PP slot release {:?} does not match oldest pending {:?}",
413                released, expected
414            ));
415        }
416        self.pending.pop_front();
417        Ok(())
418    }
419
420    fn record_send(&mut self, sent: PrimePpWaveSlot) -> Result<(), String> {
421        if sent.wave != self.next_wave {
422            return Err(format!(
423                "prime PP sent wave {} while wave {} was next",
424                sent.wave, self.next_wave
425            ));
426        }
427        if sent.slot >= 2 {
428            return Err(format!(
429                "prime PP boundary returned invalid slot {}",
430                sent.slot
431            ));
432        }
433        if self.pending.iter().any(|pending| pending.slot == sent.slot) {
434            return Err(format!(
435                "prime PP reused slot {} before its exact-wave release",
436                sent.slot
437            ));
438        }
439        self.pending.push_back(sent);
440        self.next_wave += 1;
441        Ok(())
442    }
443}
444
445fn recv_prime_pp_signal(
446    receiver: &std::sync::mpsc::Receiver<PrimePpSignal>,
447    expected: PrimePpWaveSlot,
448    exact_slot: bool,
449    label: &str,
450) -> Result<PrimePpWaveSlot, String> {
451    match receiver.recv() {
452        Ok(PrimePpSignal::Error(error)) => Err(error),
453        Ok(PrimePpSignal::Slot(received))
454            if received.wave == expected.wave
455                && (!exact_slot || received.slot == expected.slot) =>
456        {
457            if received.slot >= 2 {
458                Err(format!(
459                    "{label}: wave {} carried invalid slot {}",
460                    received.wave, received.slot
461                ))
462            } else {
463                Ok(received)
464            }
465        }
466        Ok(PrimePpSignal::Slot(received)) => Err(format!(
467            "{label}: expected wave/slot {:?}, received {:?}",
468            expected, received
469        )),
470        Err(_) => Err(format!(
471            "{label}: channel closed while waiting for wave {}",
472            expected.wave
473        )),
474    }
475}
476
477fn send_prime_pp_signal(
478    sender: &std::sync::mpsc::Sender<PrimePpSignal>,
479    signal: PrimePpSignal,
480    label: &str,
481) -> Result<(), String> {
482    sender
483        .send(signal)
484        .map_err(|_| format!("{label}: channel closed"))
485}
486
487struct PrimePpWave<'a> {
488    start: usize,
489    end: usize,
490    tokens: &'a [u32],
491}
492
493struct PrimePpStageChannels {
494    incoming: Option<std::sync::mpsc::Receiver<PrimePpSignal>>,
495    release_upstream: Option<std::sync::mpsc::Sender<PrimePpSignal>>,
496    outgoing: std::sync::mpsc::Sender<PrimePpSignal>,
497    released_downstream: std::sync::mpsc::Receiver<PrimePpSignal>,
498}
499
500impl PrimePpStageChannels {
501    fn notify_failure(&self, error: &str) {
502        if let Some(upstream) = &self.release_upstream {
503            let _ = upstream.send(PrimePpSignal::Error(error.to_string()));
504        }
505        let _ = self.outgoing.send(PrimePpSignal::Error(error.to_string()));
506    }
507}
508
509/// The DSA k-pool indexer's resident state, borrowed for one `mla_attn_core` call.
510///
511/// TWO PLANES, DIFFERENT LIFETIMES. `state` is the packed `[k_norm | gate]` row per cached token
512/// (`LatentKvLayer::index_rows`); it is append-only and grows with the cache. `pool_keys` is the
513/// collapsed key per COMPLETE pool of `pool` such rows, and it is the residency win: a pool's key
514/// is final the moment its last row lands, so pools `[0, *ready)` are never recomputed and each
515/// call builds only the pools its own tokens completed. `ready` is written back through the
516/// borrow, so the caller must persist it alongside the buffers.
517///
518/// `pool_keys` is `Option` because its size needs the indexer's `pool`, which the state plan does
519/// not carry — `mla_kpool_indices` allocates it on first use and leaves it resident thereafter.
520/// A caller that hands over a fresh `None` every call (the stateless arm) gets the old
521/// rebuild-everything behaviour, which is exactly right when the state itself is per-call.
522pub struct IndexerPlanes<'a> {
523    pub state: &'a mut CudaSlice<f32>,
524    pub pool_keys: &'a mut Option<CudaSlice<f32>>,
525    pub ready: &'a mut usize,
526    /// PHYSICAL rows of `state` when it is a TAIL RING; 0 when the plane is flat (one row per
527    /// cached token, absolute addressing). `mla_kpool_indices` rounds this DOWN to a multiple of
528    /// the indexer's `pool` — the state plan does not carry `pool`, so the allocator cannot — and
529    /// proves the liveness bound against the rounded value before it appends.
530    pub state_ring_rows: usize,
531    /// Token capacity of the session, which sizes `pool_keys`. It is NOT derivable from
532    /// `state.len()` once `state` is a ring: the ring holds one call's tail, the pool-key plane
533    /// holds the whole context collapsed `pool`-to-one.
534    pub capacity_tokens: usize,
535}
536
537/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
538pub(crate) struct AttnPre {
539    pub q: cudarc::driver::CudaSlice<f32>,
540    pub k: cudarc::driver::CudaSlice<f32>,
541    pub v: cudarc::driver::CudaSlice<f32>,
542    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
543}
544
545/// task #18: one sequence's GDN prep outputs (the scan inputs).
546pub(crate) struct GdnPrep {
547    pub hk: usize,
548    pub q_l2: cudarc::driver::CudaSlice<f32>,
549    pub k_l2: cudarc::driver::CudaSlice<f32>,
550    pub v_g: cudarc::driver::CudaSlice<f32>,
551    pub beta: cudarc::driver::CudaSlice<f32>,
552    pub g_log: cudarc::driver::CudaSlice<f32>,
553    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
554    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
555}
556
557/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
558pub(crate) struct VerifyStreamScratch {
559    pub pos_d: CudaSlice<i32>,
560    pub row_ctrs: Vec<CudaSlice<i32>>,
561}
562use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
563
564struct MoeInputTraceWriter {
565    dir: std::path::PathBuf,
566    index: std::fs::File,
567    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
568}
569
570static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
571    std::sync::OnceLock::new();
572
573/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
574/// per-expert launch chain). See `moe_gdec_token`.
575fn gdec_enabled() -> bool {
576    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
577    *E.get_or_init(|| {
578        std::env::var("MEMRA_MOE_GDEC")
579            .map(|v| v != "0")
580            .unwrap_or(true)
581    })
582}
583
584/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
585/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
586/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
587/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
588/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
589/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
590/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
591/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
592/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
593/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
594fn moe_slab_enabled() -> bool {
595    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
596}
597
598/// `MEMRA_MOE_FUSED_EPI` — the glm5_next fused MoE epilogue (sigmoid-routed, PRE-clamped SwiGLU,
599/// per-expert macro fold) collapsed into one launch pair per token-layer.
600///
601/// DEFAULT OFF, deliberately (docs/FLAGS.md carries the row and the reasons). The arm is proven
602/// EXACT against `memra_reference` by `tests/glm5_moe_epilogue_gpu.rs`, but it has no throughput
603/// receipt: the rig is correctness-only by law and the 190.7 GB artifact has never been on it, so
604/// the launch-count claim is arithmetic from source and nothing has been measured on serving
605/// hardware. Unmeasured behavior does not default ON.
606///
607/// Read PER CALL, not latched in a `OnceLock`: the acceptance gate flips both arms inside one
608/// test process (the interleave unit is a model load, not a boot), and a latched flag would make
609/// the second arm silently a copy of the first.
610fn moe_fused_epi_enabled() -> bool {
611    std::env::var("MEMRA_MOE_FUSED_EPI")
612        .map(|v| v != "0")
613        .unwrap_or(false)
614}
615
616/// `MEMRA_HC_DECODE_WS` — the persistent hc-glue decode workspace (lane/glm5-decode-diet
617/// lever 2): the T=1 hc decode walk lands its glue transients (mixes, gates, comb, collapse
618/// y, both norm scratches, the per-site post output) in one per-engine `HyperDecodeWs`
619/// instead of ~12 fresh `cuMemAllocAsync`+free pairs per layer per token (the launch-diet
620/// census's 2,358-calls/token class). Same kernels, same call order, same operand bytes —
621/// byte identity ON/OFF gated by `tests/hc_decode_ws_gpu.rs`.
622///
623/// DEFAULT OFF, deliberately (docs/FLAGS.md row): the alloc-call reduction is proven on the
624/// rig by counter receipt, but the ms/token value is arithmetic against the box's measured
625/// launch/alloc constants — nothing has been measured on serving hardware yet. Unmeasured
626/// behavior does not default ON.
627///
628/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent).
629fn hyper_decode_ws_on() -> bool {
630    std::env::var("MEMRA_HC_DECODE_WS").as_deref() == Ok("1")
631}
632
633/// Engagement counter for the workspace walk — the receipt the gate and any box A/B arm
634/// must show (engagement lines are receipts, never inferred).
635pub static HC_DECODE_WS_DISPATCHES: std::sync::atomic::AtomicU64 =
636    std::sync::atomic::AtomicU64::new(0);
637
638/// `MEMRA_MLA_TC_PREFILL` — the glm5_next tensor-core MLA prefill chain
639/// (lane/glm5-mla-tc-prefill, 2026-08-30): at prefill widths the three per-position f32
640/// kernels the launch-diet census named (`memra_mla_attn_gathered_kernel` 139 ms +
641/// `memra_mla_absorb_q_kernel` 44.5 ms + `memra_mla_decompress_v_kernel` 43.6 ms per
642/// layer-chunk, 75.8% of a 98%-GPU-busy cold prime) are replaced by two strided-batched
643/// bf16 tensor-core GEMMs (absorb / decompress, one launch each) and one gathered
644/// flash-attention MMA kernel (`fa_mla_gathered_bf16`). Selection, the latent cache, the
645/// q/kv projections, and decode are UNTOUCHED.
646///
647/// DEFAULT ON (owner acceptance 2026-08-30, "why not? i dont see why not", on the two-box
648/// A/B receipts): interleaved x5 fresh boots per arm on BOTH the Server-Edition and
649/// Workstation-Edition 4-card boxes, zero violations, zero argmax flips across 20 boots,
650/// TTFD -62%..-69% (7.45->2.83 s @4.6k / 6.58->2.51 s), prefill 619-724 -> 1629-2255 tok/s,
651/// vendor-default sampled twin -66/-67%, decode untouched, engagement receipted in every ON
652/// boot with no cuBLASLt declines (docs/FLAGS.md row carries the pointers). The numeric
653/// config remains band-gated (bf16 operands, f32 accumulate — the fa_prefill/MEMRA_PP_BF16
654/// class, `tests/mla_tc_prefill_gpu.rs`, never bit). `MEMRA_MLA_TC_PREFILL=0` is the
655/// rollback seam.
656///
657/// Read PER CALL, not latched (the `MEMRA_MOE_FUSED_EPI` rollback-seam precedent): the gate
658/// flips both arms inside one test process, and a latched flag would make the second arm
659/// silently a copy of the first.
660fn mla_tc_prefill_enabled() -> bool {
661    std::env::var("MEMRA_MLA_TC_PREFILL")
662        .map(|v| v != "0")
663        .unwrap_or(true)
664}
665
666/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
667/// default flip. `=0` selects the established path, while any other explicit value enables the
668/// grouped research arm for the current call.
669fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
670    std::env::var("MEMRA_MOE_GROUPED")
671        .map(|value| value != "0")
672        .unwrap_or(false)
673}
674
675/// `MEMRA_MOE_GROUPED_PREFILL`: the glm5_next expert-grouped MoE PREFILL arm, token-sort by
676/// expert (host CSR, the `moe_align_block_size` shape), then ONE grouped tensor-core GEMM per
677/// projection per layer-chunk over the resident NVFP4 bank, with the sigmoid `noaux_tc` routing,
678/// the PRE-clamped SwiGLU epilogue and the per-expert `weight_scale_2` macro fold the fused
679/// epilogue lane qualified for this family.
680///
681/// DEFAULT ON since 2026-08-29 (owner acceptance; `=0` is the rollback seam). The flip carries
682/// its receipts, per the flag-default law: reference-band + routing-exactness gate green
683/// (`tests/glm5_moe_grouped_prefill_gpu.rs`; grouped GEMM is measured non-bit-stable, so byte
684/// identity is not the honest bar; routing sel/w stay bit-identical by construction, the same
685/// `moe_router_logits` + `moe_route_sigmoid_cfg` invocation as the sequential arm), plus the
686/// interleaved x5 box A/B on the serving card class: TTFD 54.2 -> 7.5 s / 65.5 -> 8.9 /
687/// 75.9 -> 10.3 at 4.6/5.5/6.5k-token real prompts (85 -> 616-639 tok/s prefill, decode
688/// unchanged, sampled vendor-default twin green, engagement 42/42). The one greedy first-token
689/// flip (B5550) sits at a position the 8-draw vendor-default census measured as SOFT in both
690/// arms (the OFF arm itself draws the ON arm's token there) and was accepted by the OWNER on
691/// 2026-08-29, the MEMRA_BF16_MMV acceptance class. Receipts:
692/// `research/glm53-flash-bringup-20260827/moe-grouped-prefill-receipts/` (`box-ab-20260829/`).
693///
694/// Read PER CALL, not latched: the acceptance gate flips both arms inside one test process.
695fn moe_grouped_prefill_enabled() -> bool {
696    std::env::var("MEMRA_MOE_GROUPED_PREFILL")
697        .map(|v| v != "0")
698        .unwrap_or(true)
699}
700
701/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
702/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
703fn moe_prefetch_enabled() -> bool {
704    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
705    *E.get_or_init(|| {
706        std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
707            || crate::spill_pread::worker_enabled()
708    })
709}
710
711/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
712/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
713/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
714/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
715fn moe_page_prefetch_window() -> usize {
716    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
717    *W.get_or_init(|| {
718        page_prefetch_window_from_values(
719            std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
720            std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
721                .ok()
722                .as_deref(),
723        )
724    })
725}
726
727fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
728    if !enabled {
729        return 0;
730    }
731    raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
732}
733
734/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
735/// full window; each later position adds one expert at the far edge. Thus widening the window does
736/// not repeatedly issue `MADV_WILLNEED` for the same range.
737fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
738    if window == 0 || position >= len {
739        return len..len;
740    }
741    let (start, count) = if position == 0 {
742        (1, window)
743    } else {
744        (position.saturating_add(window), 1)
745    };
746    let start = start.min(len);
747    start..start.saturating_add(count).min(len)
748}
749
750/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
751/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
752fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
753    let position = current.map_or(0, |position| position.saturating_add(1));
754    (position < order_len).then_some(position)
755}
756
757/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
758/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
759/// window. Position zero primes the current expert too: its three independent reads can run in
760/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
761fn worker_prefetch_window() -> usize {
762    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
763    *WINDOW.get_or_init(|| {
764        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
765        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
766            .ok()
767            .and_then(|value| value.parse::<usize>().ok())
768            .unwrap_or(automatic.max(1))
769    })
770}
771
772/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
773/// this includes the current expert when the window is seeded so all three current projections
774/// enter the CPU pool together.
775fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
776    if window == 0 || position >= len {
777        return len..len;
778    }
779    let (start, count) = if position == 0 {
780        (0, window)
781    } else {
782        (position.saturating_add(window).saturating_sub(1), 1)
783    };
784    let start = start.min(len);
785    start..start.saturating_add(count).min(len)
786}
787
788/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
789/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
790/// expert weight pointers come from the per-layer device table. Requires the fused router (the
791/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
792fn moe_dev_enabled() -> bool {
793    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
794    *E.get_or_init(|| {
795        std::env::var("MEMRA_MOE_DEV")
796            .map(|v| v != "0")
797            .unwrap_or(true)
798            && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
799    })
800}
801
802/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
803/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
804/// Where the verify-rows MoE pair's routed selection lives for one layer-call
805/// (lane/glm5-moe-loc door D). `Host` is the shipped arm: the router's pinned readback gave the
806/// host `sel`/`w`, and the host builds the pointer/scale tables. `Dev` is door D's arm: the
807/// router's own device `sel_idx`/`sel_w` are still live, so the tables are built where they are
808/// and the readback (2 DtoH + a full `cuStreamSynchronize` per MoE layer-call) never happens.
809/// ONE launch path consumes both — only the table build differs, term-for-term identically.
810enum VrowsSel<'a> {
811    Host(&'a [u32], &'a [f32]),
812    Dev(&'a CudaSlice<i32>, &'a CudaSlice<f32>),
813}
814
815fn sigmoid_router_enabled() -> bool {
816    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
817    *E.get_or_init(|| {
818        std::env::var("MEMRA_SIG_ROUTER")
819            .map(|v| v != "0")
820            .unwrap_or(true)
821    })
822}
823
824/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
825/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
826/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
827/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
828fn moe_q8_enabled() -> bool {
829    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
830    *E.get_or_init(|| {
831        std::env::var("MEMRA_MOE_Q8")
832            .map(|v| v != "0")
833            .unwrap_or(true)
834    })
835}
836
837/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
838/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
839fn expert_dp4a_supported(qt: i32) -> bool {
840    qt == crate::QT_Q4_0
841        || qt == crate::QT_IQ3_S
842        || qt == crate::QT_IQ4_XS
843        || qt == crate::QT_Q3_K
844        || qt == crate::QT_Q4_K
845        || qt == crate::QT_Q6_K
846}
847
848fn q8_expert_supported(qt: i32) -> bool {
849    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
850    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
851    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
852    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
853    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
854    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
855    let kq = *KQ.get_or_init(|| {
856        std::env::var("MEMRA_MOE_Q8_KQ")
857            .map(|v| v != "0")
858            .unwrap_or(true)
859    });
860    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
861    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
862    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
863    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
864    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
865    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
866    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
867        .map(|v| v != "0")
868        .unwrap_or(true);
869    qt == crate::QT_IQ3_S
870        || qt == crate::QT_IQ4_XS
871        || (nvfp4_q8 && qt == crate::QT_NVFP4)
872        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
873}
874
875/// ModelOpt `W4A16_NVFP4` is weight-only: feeding its expert weights a q8_1 activation changes
876/// the declared numeric program to W4A8. Keep the established q8 path for every other artifact,
877/// but force Hy3 W4A16 experts through the BF16-activation `qmatvec_view` oracle.
878fn q8_expert_supported_for_model(cfg: &ModelConfig, qt: i32) -> bool {
879    let weight_only_nvfp4 = cfg.hy3.as_ref().is_some_and(|hy3| hy3.weight_only_nvfp4);
880    q8_expert_supported(qt) && !(weight_only_nvfp4 && qt == crate::QT_NVFP4)
881}
882
883fn moe_q8_enabled_for_model(cfg: &ModelConfig, m: &MoeWeights) -> bool {
884    m.has_uniform_expert_layout()
885        && moe_q8_enabled()
886        && q8_expert_supported_for_model(cfg, m.gate_exps.qtype)
887        && q8_expert_supported_for_model(cfg, m.up_exps.qtype)
888        && q8_expert_supported_for_model(cfg, m.down_exps.qtype)
889}
890
891#[cfg(test)]
892mod w4a16_dispatch_tests {
893    use super::q8_expert_supported_for_model;
894    use memra_gguf::config::{HfConfig, ModelConfig};
895
896    #[test]
897    fn hy3_w4a16_never_admits_q8_activations() {
898        let hf = HfConfig::parse(
899            r#"{"model_type":"hy_v3","num_hidden_layers":2,"hidden_size":8,
900            "num_attention_heads":2,"intermediate_size":16,"vocab_size":32,
901            "max_position_embeddings":32,
902            "quantization_config":{"quant_method":"modelopt","quant_algo":"W4A16_NVFP4"}}"#,
903        );
904        let cfg = ModelConfig::from_hf(&hf);
905        assert!(!q8_expert_supported_for_model(&cfg, crate::QT_NVFP4));
906        assert!(q8_expert_supported_for_model(&cfg, crate::QT_IQ4_XS));
907    }
908}
909
910/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
911/// k-quant tensors must fall to the _em dot path instead.
912fn q8_expert_dec_supported(qt: i32) -> bool {
913    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
914}
915
916/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
917/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
918/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
919/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
920/// q35 layers, which is why that cell measured FLAT.
921fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
922    match qt {
923        crate::QT_Q4_0 => in_f.is_multiple_of(32),
924        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
925            in_f.is_multiple_of(256)
926        }
927        // NVFP4 (block 64) added lane/moebatch-q35moe 2026-08-21: the ornith15 expert bank is
928        // uniform NVFP4, which passed the pairs q8 gate but missed BOTH batched doors
929        // (use_mma's dec set and this table), so 14.7k-token prefill rode the per-pair _em
930        // fallback — 88.6% of the prime wall (prime-anatomy receipt).
931        crate::QT_NVFP4 => in_f.is_multiple_of(64),
932        _ => false,
933    }
934}
935
936/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
937/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
938fn moe_prewarm_enabled() -> bool {
939    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
940    *E.get_or_init(|| {
941        std::env::var("MEMRA_MOE_PREWARM")
942            .map(|v| v != "0")
943            .unwrap_or(true)
944    })
945}
946
947/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
948/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
949/// can then vote for and exercise those experts on GPU before the cache is frozen.
950fn cpu_expert_profile_admit_enabled() -> bool {
951    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
952    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
953}
954
955/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
956/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
957/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
958pub const PRIME_MIN_T: usize = 16;
959
960/// MEMRA_STEP_GEMM_PRIME_SUFFIX: does a CONTINUATION prime (`cache.pos > 0` — a rewound
961/// session's suffix, or a prompt remainder split across scheduler ticks) ride the batched
962/// GEMM prime, like a fresh prompt does?
963///
964/// DEFAULT ON since 2026-08-29, by decision, under the flip bar the OFF-era FLAGS row
965/// wrote down (never byte identity — a prime-decomposition m-dependence that EVERY
966/// measured prime path shares, walk included, bars that gate for all of them):
967///  1. vendor-default sampled rows: the blind, rubric-pre-registered 8-turn quality A/B
968///     (research/step37-sampled-quality-20260828, 72/72 valid rows, engagement receipts
969///     per row) — WARM-GEMM sits inside COLD's own self-spread at t4 and t8 (t8 carried
970///     at n=16; the round-1 walk-over-gemm signal collapsed at p~0.91).
971///  2. the 8-turn cache-on twin: warm TTFT 0.58 s (door) vs 7.15 s (walk) on the real
972///     warm serving shape, zero faults.
973///  3. the batched prime's own standard: acceptance 0.80-0.86 across all arms with the
974///     door arm highest at t8; interleaved arms; zero ILLEGAL/#87/panics in 19 boots.
975///
976/// Precondition shipped first: the SWA-ring checkpoint restore fix (c9a617ca99) — real
977/// session reuse crosses the grow path before any door question matters.
978/// Why it is worth it, measured: the walk continuation costs 5.5978 ms/suffix-token
979/// against this path's 0.99 ms/token (five-point sweep, R^2 0.9976), a 7.97x suffix
980/// slope collapse.
981///
982/// The `seq_end` fix beneath is NOT gated on this door — it is unconditional, because
983/// the chunk-local `seq_end` it replaced is wrong for a fresh prompt of 4096+k tokens
984/// (k in [PRIME_MIN_T, 512)) with no continuation anywhere in sight.
985///
986/// `=0` is the kill switch (continuations back on the walk, fresh primes keep the fast
987/// path); `=1` forces; `MEMRA_STEP_GEMM_PRIME=0` remains the whole-path seam. Read per
988/// call, not cached — probes flip it in process.
989fn step_gemm_prime_suffix_on() -> bool {
990    std::env::var("MEMRA_STEP_GEMM_PRIME_SUFFIX").as_deref() != Ok("0")
991}
992
993/// Widest tick the MoE DEV per-token program serves (lane/orndecode-20260822). PRIME_MIN_T
994/// doubled as the dev-arm's upper bound on the assumption that t==16 only ever meant real
995/// prefill; the exact-16 decode tier broke that assumption — at B=16 the MoE stage crossed
996/// onto the t>=MMA_T grouped/kq GEMM program (m_e ~1.6 rows/expert: 52.6% of the tick at
997/// ~104 us/launch) or the `_em` per-pair fallback (67.7 us), both catastrophically slower
998/// than the dev q8 kernels that serve B<=8 (8.8 us gate_up covering a token's whole expert
999/// set). Decode widths 2..=16 now ride dev; the grouped/pairs prefill programs start at 17.
1000/// gate2/gate3 byte batteries at B=12/16 are the qualification (bit-checked vs isolated).
1001const MOE_DEV_MAX_T: usize = 16;
1002const PRIME_PIPE_MICROBATCHES: usize = 8;
1003const PRIME_PIPE_MIN_CHUNK: usize = 128;
1004const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
1005const PRIME_PIPE_LINEAR_WORK: usize = 8;
1006
1007fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
1008    crate::pp::prime_pp_on()
1009        && !crate::pp::pp2_streams_off()
1010        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
1011}
1012
1013fn prime_ppn_wave_auto_geometry(n_layers: usize) -> bool {
1014    crate::pp::prime_pp_on()
1015        && !crate::pp::pp2_streams_off()
1016        && crate::pp::pp_wave_on() == Ok(true)
1017        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| matches!(cuts.len(), 4 | 5))
1018}
1019
1020fn prime_pipeline_auto_geometry(n_layers: usize) -> bool {
1021    prime_pp2_auto_geometry(n_layers) || prime_ppn_wave_auto_geometry(n_layers)
1022}
1023
1024/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative.
1025/// Pipelined PP primes use the measured PP-2 geometry: up to eight microchunks, never below
1026/// 128 tokens, while the legacy 4096-token cap remains the long-context bound. PP-3/4 inherit
1027/// only the geometry when their separate MEMRA_PP_WAVE door is explicitly open.
1028pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
1029    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
1030        let parsed = value
1031            .parse::<usize>()
1032            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
1033        return if crate::cache::swa_ring_on() {
1034            if parsed == 0 {
1035                crate::cache::PRIME_CHUNK_MAX_TOKENS
1036            } else {
1037                parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
1038            }
1039        } else {
1040            parsed
1041        };
1042    }
1043    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
1044    if prime_pipeline_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
1045        chunk.min(
1046            t.div_ceil(PRIME_PIPE_MICROBATCHES)
1047                .max(PRIME_PIPE_MIN_CHUNK),
1048        )
1049    } else {
1050        chunk
1051    }
1052}
1053
1054fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
1055    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
1056}
1057
1058fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
1059    if chunk == 0 || t <= chunk {
1060        return vec![(0, t)];
1061    }
1062    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
1063    let mut start = 0usize;
1064    while start < t {
1065        let mut end = (start + chunk).min(t);
1066        if t - end > 0 && t - end < PRIME_MIN_T {
1067            if ring_on {
1068                let shifted = t - PRIME_MIN_T;
1069                end = if shifted > start { shifted } else { t };
1070            } else {
1071                end = t;
1072            }
1073        }
1074        ranges.push((start, end));
1075        start = end;
1076    }
1077    ranges
1078}
1079
1080fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
1081    let prefix = prefix as u128;
1082    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
1083}
1084
1085fn dynamic_prime_chunk_ranges(
1086    t: usize,
1087    fixed_chunk: usize,
1088    fixed: &[(usize, usize)],
1089) -> Vec<(usize, usize)> {
1090    let n = fixed.len();
1091    if n < 3 {
1092        return fixed.to_vec();
1093    }
1094
1095    let max_first = t - (n - 1) * PRIME_MIN_T;
1096    let first = fixed_chunk
1097        .div_ceil(2)
1098        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
1099        .min(max_first);
1100    let mut ranges = Vec::with_capacity(n);
1101    ranges.push((0, first));
1102
1103    let first_work = prime_chunk_work(first, t);
1104    let work_span = prime_chunk_work(t, t) - first_work;
1105    let denominator = (n - 1) as u128;
1106    let mut previous = first;
1107    for boundary in 1..n - 1 {
1108        let target = first_work * denominator + work_span * (boundary as u128);
1109        let remaining = n - 1 - boundary;
1110        let mut low = previous + PRIME_MIN_T;
1111        let mut high = t - remaining * PRIME_MIN_T;
1112        while low < high {
1113            let mid = low + (high - low) / 2;
1114            if prime_chunk_work(mid, t) * denominator >= target {
1115                high = mid;
1116            } else {
1117                low = mid + 1;
1118            }
1119        }
1120        ranges.push((previous, low));
1121        previous = low;
1122    }
1123    ranges.push((previous, t));
1124    ranges
1125}
1126
1127/// Internal prime ranges. A pipelined PP prime defaults to a short-fill, equal-modeled-time
1128/// schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
1129/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
1130///
1131/// `gdn_grid`: the model runs the chunked GDN WY scan (`HybridModel::gdn_prime_grid_on`) —
1132/// AUTO-scheduled internal boundaries are then snapped down to the WY-chunk grid
1133/// (`align_prime_ranges_to_gdn`; the spec-longctx grid law, extended from serve splits to
1134/// the PP prime microchunks). Explicit MEMRA_PRIME_CHUNK keeps its operator-authoritative
1135/// (fixed, unaligned) semantics — the FLAGS caveat documents that identity contract.
1136pub fn prime_chunk_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1137    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
1138    let chunk = prime_chunk_tokens(t, n_layers);
1139    let fixed = fixed_prime_chunk_ranges(t, chunk);
1140    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
1141        Ok(value) => value == "dynamic",
1142        Err(_) => true,
1143    };
1144    if explicit_chunk {
1145        return fixed;
1146    }
1147    let ranges = if !dynamic || !prime_pipeline_auto_geometry(n_layers) {
1148        fixed
1149    } else {
1150        dynamic_prime_chunk_ranges(t, chunk, &fixed)
1151    };
1152    // MEMRA_PRIME_GRID_ALIGN=0 is the shared rollback seam of the grid law (same env the
1153    // worker's serve-boundary alignment honors, read per call so gates can flip it
1154    // in-process): the legacy off-grid auto schedule — the toothed cell's broken arm.
1155    if gdn_grid && std::env::var("MEMRA_PRIME_GRID_ALIGN").as_deref() != Ok("0") {
1156        align_prime_ranges_to_gdn(&ranges, t, Engine::gdn_chunk_size())
1157    } else {
1158        ranges
1159    }
1160}
1161
1162/// The mHC prime schedule: how `prime_cache_hyper` SPLITS one prompt into calls.
1163///
1164/// THIS IS THE RULE THE CAPACITY GATE ASSERTS ON, and it is separate from
1165/// [`prime_chunk_ranges`] so the hyper walk's split can be read, gated and changed without
1166/// touching the serial trunk's. It DELEGATES to the serial schedule rather than deriving a
1167/// second one: the transient pressure it answers is the same pressure, and two schedules would
1168/// be two things to keep aligned with `MEMRA_PRIME_CHUNK`.
1169///
1170/// WHY THE SPLIT IS SEMANTICALLY INERT, term by term. Read this precisely: it says the split
1171/// computes the SAME PROGRAM, not that it computes the same BITS. The bit claim is FALSE on this
1172/// trunk and measured so — see the near-tie note at the end.
1173///
1174///   * **The mHC residual is strictly PER TOKEN.** `crate::hyper`'s contract is `mixes[t,:]`, an
1175///     RMS rescale over that token's own `streams*hidden` slab, a Sinkhorn per token per site,
1176///     a per-token collapse and a per-token post. The stream state is expanded at the start of a
1177///     call and collapsed at its end; it carries NOTHING between tokens. Splitting the token
1178///     axis cannot move a value.
1179///   * **KDA prefill is a SEQUENTIAL scan** (`kda.rs`: `memra_kda_scan_s128` runs prefill and
1180///     decode alike, the chunked UT transform is not the shipped path). A sequential recurrence
1181///     has no fold grid, so the GDN WY grid law has NO KDA analogue to violate and
1182///     `align_prime_ranges_to_gdn` has nothing to align. The conv ring carries across calls
1183///     already — it is the seam every decode step uses. **DEBT, named:** if the chunked KDA twin
1184///     ever becomes the prefill path, it acquires a fold grid and this schedule's internal
1185///     boundaries must be snapped to it exactly as the GDN ones are, or chunked prime stops
1186///     being bit-identical. The `gdn_grid` argument is the seam that change lands on.
1187///   * **The latent KV plane is f32** (`LatentKvLayer::rows`), so a later call reads earlier
1188///     calls' rows in the SAME numeric class it would have computed them in. There is no
1189///     analogue of the serial trunk's f32-vs-quantized-KV class edge — the thing that made
1190///     `MEMRA_PRIME_CHUNK` steer arithmetic until the 2026-08-05 grain-free fix.
1191///   * **The DSA pool keys are already incremental** (`index_pools_ready`): a pool key is a pure
1192///     function of its own `pool` state rows and the constant `kpool_ape`, final the instant the
1193///     pool's last row lands, so no boundary can move one. Selection is per query over resident
1194///     keys, with visibility keyed on the query's ABSOLUTE cache row.
1195///
1196/// NOT BIT-STABLE ACROSS CHUNK SIZES, and the cause is NOT the split. Measured on the rig
1197/// (`research/glm53-flash-bringup-20260827/1m-context-20260828/02`): the arms diverge at ROW 0,
1198/// which no cross-token state can reach, and `Engine::linear` — the cuBLASLt f32 `mixes` GEMM in
1199/// `hyper::pre` — is itself not m-invariant (m=32 vs m=200 moves 9601/12288 output bits at worst
1200/// 3.815e-6, the same worst the chunked prime reports; m=128 and m=199 vs m=200 are identical).
1201/// cuBLASLt reselects its algorithm by shape and the reduction order goes with it, which
1202/// `hyper.rs`'s header already concedes ("a serving trunk, not a byte-parity oracle"). So this
1203/// is a documented near-tie class the split EXPOSES, not one it creates. It is written into the
1204/// `MEMRA_PRIME_CHUNK` FLAGS row, and `glm5_chunked_prime_gpu` holds the split to a calibrated
1205/// band anchored on `memra_reference::execute` rather than on the monolithic sibling.
1206///
1207/// One thing the split provably cannot break here: POSITIONS. glm5_next is NoPE end to end
1208/// (`qk_rope_head_dim = 0`, `mla_use_nope`), and KDA is positionless, so `pos_d` reaches no
1209/// kernel on this path — a mutation that made it call-local instead of session-absolute moved
1210/// nothing at all.
1211///
1212/// A prompt at or under one chunk takes the monolithic body unchanged, and `MEMRA_PRIME_CHUNK=0`
1213/// restores the monolithic walk at any length — the rollback seam, and the oracle arm the
1214/// correctness gate compares against.
1215pub fn hyper_prime_ranges(t: usize, n_layers: usize, gdn_grid: bool) -> Vec<(usize, usize)> {
1216    prime_chunk_ranges(t, n_layers, gdn_grid)
1217}
1218
1219/// The largest number of token rows any ONE mHC prime call carries under
1220/// [`hyper_prime_ranges`]. Every per-call transient in the walk — the `t*streams*hidden` stream
1221/// state, the MLA query planes, and the DSA indexer's `t * n_pools` score plane — is
1222/// proportional to this, so it is the single number a capacity assertion needs.
1223pub fn hyper_prime_call_rows(t: usize, n_layers: usize, gdn_grid: bool) -> usize {
1224    hyper_prime_ranges(t, n_layers, gdn_grid)
1225        .iter()
1226        .map(|&(start, end)| end - start)
1227        .max()
1228        .unwrap_or(0)
1229}
1230
1231/// Per-request prefill WORKSPACE coefficients for a HyperConnections trunk, published to
1232/// admission (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper model: their
1233/// admission arithmetic is byte-identical to the pre-lane behavior.
1234///
1235/// These are the FORMULA behind the 262k 2-card cell's measured ~0.8 MiB/token/card prefill
1236/// wall (`research/glm53-flash-bringup-20260827/262k-2card-20260830/LANE.md`), not the slope
1237/// itself: each term is the size of a named allocation in the walk, summed per token of ONE
1238/// prime call. On GLM-5.3-Flash geometry (H=4096, S=4, F=2048, U=8, heads=64, qk=256, v=256,
1239/// topk=2048, P=4) `chunk_token_bytes` evaluates to ~0.86 MiB — the receipt's slope with the
1240/// conservative side up. The attribution table naming every term lives in
1241/// `research/glm53-flash-bringup-20260827/gpf-workspace-20260830/LANE.md` §1.
1242#[derive(Debug, Clone, Copy)]
1243pub struct HyperPrimeWorkspaceShape {
1244    /// Bytes of per-call prefill transients PER TOKEN OF ONE PRIME CALL: the double-buffered
1245    /// `[t, streams, hidden]` stream state + ppN boundary slots, the pre/norm transients, the
1246    /// grouped-MoE staging (CSR activations + three f32 partial planes + f16 mirrors + scatter
1247    /// planes), the MLA query/attention planes, the k-pool idx plane, and the prime-tail
1248    /// hidden/norm pair. Multiplied by [`hyper_prime_call_rows`] this bounds the workspace of
1249    /// the CHUNKED prime; on the monolithic rollback (`MEMRA_PRIME_CHUNK=0`) the call rows are
1250    /// the whole prompt and the same product stays honest.
1251    pub chunk_token_bytes: usize,
1252    /// Bytes per PROMPT token that live for the WHOLE prime on the last stage: the returned
1253    /// pre-output_norm `hiddens` stack (`n_embd` f32), consumed by the MTP-spec `prompt_h`
1254    /// and the embed capture.
1255    pub prompt_bytes_per_token: usize,
1256    /// DSA k-pool group size `P`, or 0 when the model runs no k-pool indexer. The selection
1257    /// score plane of ONE call is `call_rows * (ctx / P)` f32 — the one prefill transient that
1258    /// stays COUPLED TO CONTEXT DEPTH after chunking (it is the allocation the 3-card 1M prime
1259    /// died on at 97.2 GiB).
1260    pub kpool_score_pool: usize,
1261    /// Trunk layer count, for re-deriving [`hyper_prime_call_rows`] at admission time with the
1262    /// same env-sensitive schedule the prime itself will walk.
1263    pub n_layers: usize,
1264    /// The model's own GDN grid-alignment input to the schedule.
1265    pub gdn_grid: bool,
1266}
1267
1268impl HyperPrimeWorkspaceShape {
1269    /// The admission charge for one request: workspace of the LARGEST prime call this request
1270    /// can produce, plus the ctx-coupled score plane at that call width, plus the prompt-long
1271    /// hiddens stack.
1272    ///
1273    /// Keyed on PROMPT rows, deliberately not on `ctx_cap`: every term here is a function of
1274    /// what the PRIME walks, and a `max_tokens`-omitted request carries a `ctx_cap` of the
1275    /// whole server window — charging the window would refuse every vendor-default short
1276    /// prompt on a deep-window box for workspace it never allocates. A continuation request's
1277    /// `prompt` is the full rendered conversation (the suffix optimization is internal
1278    /// reuse), so the score plane's `t_kv` is covered too.
1279    pub fn admission_bytes(&self, prompt_rows: usize) -> usize {
1280        let rows = hyper_prime_call_rows(prompt_rows, self.n_layers, self.gdn_grid);
1281        let chunk = self.chunk_token_bytes.saturating_mul(rows);
1282        let score = prompt_rows
1283            .checked_div(self.kpool_score_pool)
1284            .map(|pools| rows.saturating_mul(pools).saturating_mul(size_of::<f32>()))
1285            .unwrap_or(0);
1286        chunk
1287            .saturating_add(score)
1288            .saturating_add(self.prompt_bytes_per_token.saturating_mul(prompt_rows))
1289    }
1290}
1291
1292/// Snap AUTO prime-range internal boundaries DOWN to the GDN WY-chunk grid (lane/
1293/// hermes-perf-fixes, 2026-08-23 — the missing helper the PP-auto-ranges finding names).
1294///
1295/// THE LAW THIS EXTENDS (measured, research/multiturn-cache-20260821/
1296/// LONGCTX-EXACTNESS-20260821.md; the serve-split half already ships as the worker's
1297/// `grid_align_boundary`): under the chunked WY scan a prompt primed as two calls split at
1298/// L is bit-identical to the monolithic prime iff `L % gdn_chunk_size() == 0` — an off-grid
1299/// call start shifts the fold grid and materializes recurrent state at a point the
1300/// monolithic program never computes. The prime loop walks these ranges as separate
1301/// `prime_layers` calls, so INTERNAL microchunk boundaries are the same seam: the PP-2
1302/// auto geometry (`t.div_ceil(8).max(128)` fills, and every dynamic short-fill boundary)
1303/// lands off the 32-token grid for most prompt lengths, which is exactly the
1304/// chunk-value bit-identity the GDN lane falsified (FLAGS PRIME_CHUNK/SCHED caveat).
1305///
1306/// Boundaries only move DOWN (earlier is always semantically safe — same argument as the
1307/// worker's alignment); a boundary that collapses onto its predecessor is dropped (ranges
1308/// merge). The final range always ends at `t`. Aligning down only GROWS the tail
1309/// remainder, so the fixed-schedule tail-merge rule is never re-violated. Cost bound: at
1310/// most `c-1` tokens shift per boundary.
1311pub fn align_prime_ranges_to_gdn(
1312    ranges: &[(usize, usize)],
1313    t: usize,
1314    c: usize,
1315) -> Vec<(usize, usize)> {
1316    if c == 0 || ranges.len() < 2 {
1317        return ranges.to_vec();
1318    }
1319    let mut out: Vec<(usize, usize)> = Vec::with_capacity(ranges.len());
1320    let mut start = 0usize;
1321    for (i, &(_, end)) in ranges.iter().enumerate() {
1322        let e = if i + 1 == ranges.len() {
1323            t
1324        } else {
1325            end / c * c
1326        };
1327        if e > start {
1328            out.push((start, e));
1329            start = e;
1330        } // else: boundary collapsed onto its predecessor — merge into the next range
1331    }
1332    debug_assert_eq!(out.last().map(|&(_, e)| e), Some(t));
1333    out
1334}
1335
1336struct HeadSplit {
1337    pin: u64,
1338    w1: CudaSlice<u8>,
1339    hn1: CudaSlice<f32>,
1340    y1: CudaSlice<f32>,
1341    logits_e: CudaSlice<f32>,
1342    ev_hn: cudarc::driver::CudaEvent,
1343    ev_done: cudarc::driver::CudaEvent,
1344    raw_hn1: u64,
1345    raw_y1: u64,
1346    raw_logits_hi: u64,
1347    /// SAMPLED-TAIL scratch (perturbed row + the filter's threshold/z/max slots + the row
1348    /// index). Allocating these per token cost more than the split head saved: the first
1349    /// sampled-split measurement came in at 78.25 tok/s against 78.96 for the unsplit head,
1350    /// which is five allocations per token, not arithmetic.
1351    samp: Option<SampScratch>,
1352}
1353
1354struct SampScratch {
1355    pb: CudaSlice<f32>,
1356    th: CudaSlice<f32>,
1357    z: CudaSlice<f32>,
1358    mx: CudaSlice<f32>,
1359    rows: CudaSlice<i32>,
1360}
1361/// HEAD-SPLIT workspace (host + device twins share it).
1362static HEAD_SPLIT_WS: std::sync::Mutex<Option<HeadSplit>> = std::sync::Mutex::new(None);
1363
1364/// DEV1-LOCAL ROUTER replicas (MEMRA_DEV1_ROUTER): per-layer (gate_inp_f32, exp_probs_b,
1365/// active_experts) on rank1 + a shared logits scratch. Deterministic kernels on identical
1366/// input bits — rank1's local selection is bit-equal to the root's.
1367#[allow(clippy::type_complexity)]
1368static DEV1_ROUTER_REPS: std::sync::Mutex<
1369    Option<(
1370        std::collections::HashMap<u16, (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<u8>)>,
1371        Option<CudaSlice<f32>>,
1372    )>,
1373> = std::sync::Mutex::new(None);
1374
1375/// SHEXP-ON-DEV1 workspace (MEMRA_SHEXP_DEV1): replica weights + scratch on rank1, the
1376/// down row lands ROOT-resident over P2P (single store pass), and apply adds it on e
1377/// behind ev_done. (pins, wg1, wu1, wd1, act1, sh_root, ev_z, ev_done).
1378#[allow(clippy::type_complexity)]
1379static SHEXP_D1_REPS: std::sync::Mutex<
1380    Option<std::collections::HashMap<u16, (CudaSlice<u8>, CudaSlice<u8>, CudaSlice<u8>)>>,
1381> = std::sync::Mutex::new(None);
1382#[allow(clippy::type_complexity)]
1383static SHEXP_D1_WS: std::sync::Mutex<
1384    Option<(
1385        (usize, usize),
1386        CudaSlice<f32>,
1387        CudaSlice<f32>,
1388        CudaSlice<f32>,
1389        cudarc::driver::CudaEvent,
1390        cudarc::driver::CudaEvent,
1391    )>,
1392> = std::sync::Mutex::new(None);
1393
1394/// SHEXP OVERLAP workspace (issue writes, apply reads): (device, n_embd, n_ff_sh, act, sh).
1395#[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1396static SHEXP_OV_WS: std::sync::Mutex<
1397    Option<(usize, usize, usize, CudaSlice<f32>, CudaSlice<f32>)>,
1398> = std::sync::Mutex::new(None);
1399
1400impl HybridModel {
1401    /// Does this model's prime schedule live under the GDN WY-chunk grid law? True when the
1402    /// trunk has GDN (linear-attention) layers AND the chunked scan is on — the regime where
1403    /// an off-grid prime-call boundary shifts the WY fold grid (see
1404    /// `align_prime_ranges_to_gdn`). Attention-only models and the sequential scan
1405    /// (`MEMRA_GDN_CHUNKED=0`) are split-invariant, so the grid is a no-op contract there.
1406    pub fn gdn_prime_grid_on(&self) -> bool {
1407        Engine::gdn_chunked_enabled()
1408            && self
1409                .layers
1410                .iter()
1411                .any(|l| matches!(l.mixer, crate::hybrid::Mixer::Linear(_)))
1412    }
1413
1414    /// Can the step TP runtime run the DEVICE-RESIDENT activation path from this serving
1415    /// engine? Native P2P (peer copies replace the host staging) AND a shared root context
1416    /// (the device buffers must be addressable on both sides — the TP registry builds its
1417    /// own Engine per rank, so this is a real seam, not a formality).
1418    fn full_attn_tp_device_resident(e: &Engine, tp: &crate::hybrid::StepTpQkv) -> bool {
1419        tp.runtime.native_p2p() && tp.runtime.root_shares_ctx(e)
1420    }
1421
1422    pub(crate) fn full_attn_tp_qkv(
1423        &self,
1424        e: &Engine,
1425        fa: &FullAttnLayer,
1426        h: &CudaSlice<f32>,
1427        t: usize,
1428    ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
1429        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1430            return Ok(None);
1431        };
1432        let values = active_matrix_values(
1433            h.len(),
1434            t,
1435            self.cfg.n_embd as usize,
1436            "Step TP QKV activation",
1437        )?;
1438        // DEVICE-RESIDENT NATIVE PATH (lane/hermes-perf-fixes, 2026-08-23 — the host-bounce
1439        // finding): the native-P2P arm used to dtoh the FULL hidden state per layer, run
1440        // from a host copy, gather q/k/v to host vectors, and htod all three back — a host
1441        // round-trip on every execute that the peer transport exists to remove. The
1442        // device twins are byte-identical by construction (the same bytes travel dtod
1443        // instead of dtoh+htod; kernels, peer copies, and gather order are shared code).
1444        // The host arm below remains the transport for !native_p2p (host staging IS that
1445        // transport) and for a root context this engine cannot address.
1446        if Self::full_attn_tp_device_resident(e, tp) {
1447            // Producer fence: h was written on THIS engine's stream; the TP ranks read it
1448            // on theirs (same context, different streams).
1449            e.stream().synchronize()?;
1450            let q = tp
1451                .runtime
1452                .bf16_column_parallel_resident_native_device(&tp.q, h, t)?;
1453            let k = tp
1454                .runtime
1455                .bf16_column_parallel_resident_native_device(&tp.k, h, t)?;
1456            let v = tp
1457                .runtime
1458                .bf16_column_parallel_resident_native_device(&tp.v, h, t)?;
1459            Self::full_attn_tp_log_once(tp, "qkv", "device-resident");
1460            return Ok(Some(vec![q, k, v]));
1461        }
1462        let host = e.dtoh_view(&h.slice(0..values))?;
1463        let q = if tp.runtime.native_p2p() {
1464            tp.runtime
1465                .bf16_column_parallel_resident_native(&tp.q, &host, t)?
1466        } else {
1467            tp.runtime
1468                .bf16_column_parallel_resident(&tp.q, &host, t)?
1469                .gathered
1470        };
1471        let k = if tp.runtime.native_p2p() {
1472            tp.runtime
1473                .bf16_column_parallel_resident_native(&tp.k, &host, t)?
1474        } else {
1475            tp.runtime
1476                .bf16_column_parallel_resident(&tp.k, &host, t)?
1477                .gathered
1478        };
1479        let v = if tp.runtime.native_p2p() {
1480            tp.runtime
1481                .bf16_column_parallel_resident_native(&tp.v, &host, t)?
1482        } else {
1483            tp.runtime
1484                .bf16_column_parallel_resident(&tp.v, &host, t)?
1485                .gathered
1486        };
1487        Self::full_attn_tp_log_once(tp, "qkv", "host-canonical");
1488        Ok(Some(vec![e.htod(&q)?, e.htod(&k)?, e.htod(&v)?]))
1489    }
1490
1491    /// One transport banner per (projection, transport) — the old per-call eprintln fired
1492    /// on EVERY layer of EVERY step, itself a decode-rate cost on the path this lane is
1493    /// unbouncing (the sibling grouped-EP path already learned this).
1494    fn full_attn_tp_log_once(tp: &crate::hybrid::StepTpQkv, proj: &str, activation: &'static str) {
1495        use std::sync::atomic::{AtomicBool, Ordering};
1496        static LOGGED: [AtomicBool; 4] = [
1497            AtomicBool::new(false),
1498            AtomicBool::new(false),
1499            AtomicBool::new(false),
1500            AtomicBool::new(false),
1501        ];
1502        let idx = 2 * usize::from(proj == "o") + usize::from(activation == "device-resident");
1503        if LOGGED[idx].swap(true, Ordering::Relaxed) {
1504            return;
1505        }
1506        eprintln!(
1507            "[step-tp-{proj}] execute layer={} devices={:?} projections={proj} \
1508             tensor_parallel=true attention_local=true kv_local=true transport={} \
1509             native_p2p={} bulk_p2p={} activation={activation} \
1510             output={} performance_claim=false (logged once per transport)",
1511            tp.layer,
1512            tp.devices,
1513            tp.runtime.transport_label(),
1514            tp.runtime.native_p2p(),
1515            tp.runtime.bulk_p2p(),
1516            if activation == "device-resident" {
1517                "root-resident"
1518            } else {
1519                "root-readback"
1520            },
1521        );
1522    }
1523
1524    pub(crate) fn full_attn_tp_o(
1525        &self,
1526        e: &Engine,
1527        fa: &FullAttnLayer,
1528        activation: &CudaSlice<f32>,
1529        tokens: usize,
1530    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1531        let Some(tp) = fa.step_tp_qkv.as_ref() else {
1532            return Ok(None);
1533        };
1534        // DEVICE-RESIDENT NATIVE PATH — the O-projection half of the same finding: no DtoH
1535        // of the attention output, no host O staging, root-resident reduction consumed in
1536        // place (byte-identical shared core: `step_bf16_row_native_reduce_from_root`).
1537        if Self::full_attn_tp_device_resident(e, tp) {
1538            e.stream().synchronize()?; // producer fence, as the QKV half
1539            let output = tp
1540                .runtime
1541                .step_bf16_row_parallel_resident_native_device(&tp.o, activation, tokens)?;
1542            Self::full_attn_tp_log_once(tp, "o", "device-resident");
1543            return Ok(Some(output));
1544        }
1545        let host = e.dtoh(activation)?;
1546        let output = if tp.runtime.native_p2p() {
1547            tp.runtime
1548                .step_bf16_row_parallel_resident_native(&tp.o, &host, tokens)?
1549        } else {
1550            tp.runtime
1551                .step_bf16_row_parallel_resident(&tp.o, &host, tokens)?
1552        };
1553        Self::full_attn_tp_log_once(tp, "o", "host-canonical");
1554        Ok(Some(e.htod(&output)?))
1555    }
1556
1557    fn full_attn_o(
1558        &self,
1559        e: &Engine,
1560        fa: &FullAttnLayer,
1561        activation: &CudaSlice<f32>,
1562        tokens: usize,
1563    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1564        match self.full_attn_tp_o(e, fa, activation, tokens)? {
1565            Some(output) => Ok(output),
1566            None => e.matmul(&fa.wo, activation, tokens),
1567        }
1568    }
1569
1570    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
1571    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
1572    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
1573    /// (it forces a dtoh + host hash per layer).
1574    fn prime_trace_path() -> Option<&'static str> {
1575        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
1576        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
1577            .as_deref()
1578    }
1579
1580    /// PRIME ANATOMY (diagnostic): `MEMRA_PRIME_ANATOMY=1` synchronizes the stream around
1581    /// each prime_layers stage and accumulates wall time per stage class, printed after
1582    /// every prime_layers call (cumulative across chunks/reps). The per-stage syncs
1583    /// serialize launch/execute overlap, so the summed total exceeds the naked prime wall —
1584    /// attribution ratios only, never a measured default run. Non-seg serial arm only.
1585    fn prime_anatomy_on() -> bool {
1586        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1587        *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
1588    }
1589
1590    fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
1591        static S: [std::sync::atomic::AtomicU64; 5] = [
1592            std::sync::atomic::AtomicU64::new(0), // 0 mixer full-attn
1593            std::sync::atomic::AtomicU64::new(0), // 1 mixer linear-attn (GDN)
1594            std::sync::atomic::AtomicU64::new(0), // 2 ffn MoE (router + experts + shexp)
1595            std::sync::atomic::AtomicU64::new(0), // 3 ffn dense
1596            std::sync::atomic::AtomicU64::new(0), // 4 norms/adds/glue
1597        ];
1598        &S
1599    }
1600
1601    /// Fail closed on a path that has not been taught the mHC residual program.
1602    ///
1603    /// A serial residual on an hc model is not a degraded answer, it is a DIFFERENT function
1604    /// computed at full speed and full confidence — the exact failure `crate::hyper` exists to
1605    /// prevent. Every trunk entry point that has not been converted calls this first, so the
1606    /// unconverted set is a list of named refusals rather than a list of silent wrong answers.
1607    pub(crate) fn refuse_hyper(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
1608        if let Some(topology) = self.hyper.as_ref() {
1609            return Err(format!(
1610                "{path} runs a serial residual, but this model's ModelPlan declares \
1611                 ResidualTopology::HyperConnections{{ streams: {}, collapse: {:?} }}. Refusing: \
1612                 that path would compute a different model. Converted paths: forward, \
1613                 forward_last, prime_cache, decode_step, and the batched serving chain \
1614                 decode_step_batch / _sampled / _lean / _masked.",
1615                topology.streams, topology.collapse
1616            )
1617            .into());
1618        }
1619        Ok(())
1620    }
1621
1622    /// The FFN branch of one hc site, from an already-normed `[t, hidden]` input.
1623    ///
1624    /// Split out because under hyper-connections the FFN's input is `rms_norm(hc_pre(x))`, not
1625    /// `rms_norm(x + attn)` — the fused add+norm+quantize forms the serial paths use have no
1626    /// residual to fold, so this is the unfused dispatch by construction.
1627    fn hyper_ffn_branch(
1628        &self,
1629        e: &Engine,
1630        layer: &crate::hybrid::HybridLayer,
1631        z: &CudaSlice<f32>,
1632        t: usize,
1633        il: usize,
1634        prefill: bool,
1635    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1636        match &layer.ffn {
1637            crate::hybrid::Ffn::Dense {
1638                ffn_gate,
1639                ffn_up,
1640                ffn_down,
1641            } => {
1642                let n_ff = ffn_gate.out_features();
1643                let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], z, t)?;
1644                let up = g2.pop().unwrap();
1645                let gate = g2.pop().unwrap();
1646                let mut act = e.uninit(t * n_ff)?;
1647                // A dense FFN reads the SHEXP clamp array — see forward()'s note.
1648                Self::ffn_act_lim(
1649                    e,
1650                    &self.cfg,
1651                    &gate,
1652                    &up,
1653                    1.0,
1654                    1.0,
1655                    self.cfg.clamp_shexp_at(il as u32),
1656                    &mut act,
1657                    t * n_ff,
1658                )?;
1659                e.matmul(ffn_down, &act, t)
1660            }
1661            crate::hybrid::Ffn::Moe(m) => {
1662                if prefill {
1663                    self.moe_ffn_il_prefill(e, m, z, t, il as u16)
1664                } else {
1665                    self.moe_ffn_il_zq8(e, m, z, None, t, il as u16)
1666                }
1667            }
1668        }
1669    }
1670
1671    /// Stateless prefill under the mHC residual (`crate::hyper`), the hc twin of `forward` /
1672    /// `forward_last`. The mixers, the FFNs and the norms are the SAME calls the serial paths
1673    /// make; only the residual program around them changes, which is the whole point — a mixer
1674    /// never sees the stream dimension.
1675    fn forward_hyper(
1676        &self,
1677        e: &Engine,
1678        tokens: &[u32],
1679        last_only: bool,
1680    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
1681        let topology = *self
1682            .hyper
1683            .as_ref()
1684            .ok_or("forward_hyper on a model with no HyperConnections topology")?;
1685        // M2 ppN door for the mHC trunk. The generic arm's door lives in `decode_step_h`;
1686        // this walk is reached BEFORE it (decode.rs routes `hyper.is_some()` first), so the
1687        // hc walks own their own door. Loud refusal, never silent fallback: an unqualified
1688        // pipeline rewrite errors here rather than running a single-engine walk over weights
1689        // the loader has already sharded across devices.
1690        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1691            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1692                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1693            }
1694            return self.forward_hyper_ppn(e, tokens, last_only, &topology, &fence);
1695        }
1696        let n_embd = self.cfg.n_embd as usize;
1697        let t = tokens.len();
1698        let eps = self.cfg.rms_eps;
1699        let pos: Vec<i32> = (0..t as i32).collect();
1700        let pos_d = e.htod_i32(&pos)?;
1701
1702        let embedded = self.embed(e, tokens)?;
1703        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
1704        let trace = memra_reference::hidden_trace::enabled();
1705        if trace {
1706            memra_reference::hidden_trace::emit_tokens(tokens);
1707            let streams = x.len() / (t * n_embd);
1708            memra_reference::hidden_trace::emit_last_row(
1709                "expand",
1710                -1,
1711                t,
1712                streams * n_embd,
1713                &e.dtoh(&x)?,
1714            );
1715        }
1716
1717        x = self.hyper_range_forward(e, &topology, x, 0, self.layers.len(), &pos_d, t, trace)?;
1718
1719        // SHARED EXIT with the ppN twin: one trunk exit, so the split and unsplit arms cannot
1720        // drift apart in the head. That is what makes `glm5-hyper-ppn-gate`'s bit-identity bar
1721        // a structural property rather than a coincidence of two maintained copies.
1722        self.hyper_head_logits(e, &topology, &x, t, n_embd, eps, last_only)
1723    }
1724
1725    /// Stateful prefill under the mHC residual: `prime_cache_overlaid`'s contract (leave a
1726    /// decode-ready cache behind, return last-row logits + the pre-output_norm hidden seed and
1727    /// stack) over the hc layer program.
1728    ///
1729    /// Deliberately UNCHUNKED and UNCAPTURED. The serial prime's chunking, prime slabs, S-mid
1730    /// graph capture and core-split arms are all keyed to the serial residual's transient set;
1731    /// re-deriving them for a stream state is a tuning lane, not a correctness one, and this
1732    /// path is the one the reference gate pins. Long prompts therefore hold `T*streams*hidden`
1733    /// f32 of stream state — 4x the serial trunk's — and that ceiling is the named cost of the
1734    /// simple form.
1735    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1736    fn prime_cache_hyper(
1737        &self,
1738        e: &Engine,
1739        tokens: &[u32],
1740        cache: &mut Cache,
1741        queued_after: usize,
1742        overlay: Option<&crate::vision::EmbedOverlay>,
1743    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1744        let topology = *self
1745            .hyper
1746            .as_ref()
1747            .ok_or("prime_cache_hyper on a model with no HyperConnections topology")?;
1748        // M2 ppN door for the mHC prime (see `forward_hyper`'s note). glm5_next has no
1749        // batched prime arm: the split twin is a straight layer-range split with the
1750        // [t, streams, hidden] state on the wire.
1751        //
1752        // CHUNKED, like the single-engine walk below (lane/glm53-1m-demo, 2026-08-29 — the
1753        // follow-up the previous note named). The monolithic ppN prime carried the WHOLE
1754        // prompt as one call, which capped it three independent ways on the 4x96 GB box:
1755        // per-call transients proportional to t OOM'd from ~32k tokens, and every launch
1756        // that places t in grid.y (kda_conv_silu, kda_gate, rms_norm over rows, the router)
1757        // hits the CUDA 65,535 grid.y ceiling from t=65,536 (measured: instant
1758        // CUDA_ERROR_INVALID_VALUE at a 128,566-token prime, receipts in
1759        // research/glm53-flash-bringup-20260827/1m-demo-20260829/). The chunk loop reuses
1760        // the SAME schedule as the single-engine walk (`hyper_prime_ranges`), so per-chunk
1761        // t is bounded and — because the per-chunk staged walk is bit-identical to the
1762        // per-chunk unsplit walk (glm5_hyper_ppn_gate arm 2) — the chunked ppN prime
1763        // composes to bit-identity with the chunked single-engine prime over the same
1764        // schedule. `queued_after + (t - end)` keeps the REQUEST-level `seq_end` invariant
1765        // across chunks (each call recomputes pos0+start + (end-start) + rest = pos0 + t +
1766        // queued_after). A prompt at or under one chunk takes the monolithic ppN body
1767        // unchanged, and `MEMRA_PRIME_CHUNK=0` restores the monolithic walk — the same
1768        // rollback seam the single-engine chunk walk documents.
1769        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1770            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1771                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1772            }
1773            // The mixed-embedding overlay rides the ppN twin too (lane/glm5-vision-default-on,
1774            // 2026-08-30): the splice is an EMBEDDING-INTAKE transform, and the ppN walk embeds
1775            // on stage 0 only — every later stage receives the already-expanded stream state.
1776            // Under the chunked ppN prime each chunk takes the overlay WINDOWED to its own
1777            // call-relative range (`EmbedOverlay::window`, the same rebase seam the serve
1778            // prefill tick uses), so splice placement is chunk-schedule-invariant. Gated by
1779            // glm5-hyper-ppn-gate's overlay arm (bit-identity vs the substituted-token truth,
1780            // red arm = shifted spans).
1781            let n_embd = self.cfg.n_embd as usize;
1782            let t = tokens.len();
1783            if cache.pos + t > cache.max_ctx {
1784                return Err("prime_cache: prompt exceeds cache max_ctx".into());
1785            }
1786            let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1787            if ranges.len() == 1 {
1788                return self.prime_cache_hyper_ppn(
1789                    e,
1790                    tokens,
1791                    cache,
1792                    queued_after,
1793                    &topology,
1794                    &fence,
1795                    overlay,
1796                );
1797            }
1798            let mut hiddens = e.uninit(t * n_embd)?;
1799            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1800            for &(start, end) in &ranges {
1801                let ov = overlay.and_then(|o| o.window(start, end - start));
1802                let (l, hs, x) = self.prime_cache_hyper_ppn(
1803                    e,
1804                    &tokens[start..end],
1805                    cache,
1806                    queued_after + (t - end),
1807                    &topology,
1808                    &fence,
1809                    ov.as_ref(),
1810                )?;
1811                e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1812                last = Some((l, hs));
1813            }
1814            let (logits, h_seed) =
1815                last.expect("hyper_prime_ranges never returns an empty schedule");
1816            return Ok((logits, h_seed, hiddens));
1817        }
1818        let n_embd = self.cfg.n_embd as usize;
1819        let t = tokens.len();
1820        if cache.pos + t > cache.max_ctx {
1821            return Err("prime_cache: prompt exceeds cache max_ctx".into());
1822        }
1823        // The REQUEST's absolute end position, computed ONCE before the walk: every chunk sees
1824        // the same value whatever the chunk size, which is the tick-seg law the serial loop
1825        // above carries verbatim (`+ queued_after` closes the serve-split axis).
1826        let seq_end = cache.pos + t + queued_after;
1827        let ranges = hyper_prime_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
1828        if ranges.len() == 1 {
1829            return self.prime_chunk_hyper(e, tokens, cache, seq_end, 0, overlay);
1830        }
1831        let mut hiddens = e.uninit(t * n_embd)?;
1832        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1833        for &(start, end) in &ranges {
1834            let (l, hs, x) =
1835                self.prime_chunk_hyper(e, &tokens[start..end], cache, seq_end, start, overlay)?;
1836            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
1837            last = Some((l, hs));
1838        }
1839        let (logits, h_seed) = last.expect("hyper_prime_ranges never returns an empty schedule");
1840        Ok((logits, h_seed, hiddens))
1841    }
1842
1843    /// Publish this model's prefill-workspace coefficients to admission
1844    /// (lane/glm5-gpf-workspace, 2026-08-30). `None` for every non-hyper trunk — the server's
1845    /// admission arithmetic is then byte-identical to the pre-lane behavior for that family.
1846    ///
1847    /// Term-by-term, each anchored on a named allocation of the hyper prime walk (glm5 numbers
1848    /// in parentheses; the full attribution table is the lane doc's §1):
1849    ///   * stream state, double-buffered at `hyper::post`, PLUS the ppN boundary tx/rx pair of
1850    ///     the same `[t, streams, hidden]` payload: `4 * S * H * 4` (256 KiB/token);
1851    ///   * pre/norm transients (`hyper::pre` y + `rms_norm` h/z + ffn_out): `4 * H * 4`
1852    ///     (64 KiB/token);
1853    ///   * prime-tail hidden/norm pair (`collapse` + output-norm stack): `2 * H * 4`
1854    ///     (32 KiB/token);
1855    ///   * grouped-MoE prefill staging (`moe_ffn_grouped_prefill_sigmoid`): the f16 CSR
1856    ///     activations `U*2H`, three f32 partial planes `3*U*4F` (gate/up/act), the f16 down
1857    ///     mirror `U*2F`, the CSR-order down output + pair-order permute `2*U*4H`, and the
1858    ///     scatter target `4H` — `U*(10H + 14F) + 4H` (560 KiB/token);
1859    ///   * MLA query/attention planes: `heads * (qk_head_dim + v_head_dim) * 4`
1860    ///     (128 KiB/token) and the k-pool idx plane `(topk/P + 1) * 4` (~2 KiB/token).
1861    ///
1862    /// Validation against truth: at GLM-5.3-Flash geometry the sum is ~0.92 MiB per call
1863    /// token, against the 262k cell's MEASURED retained slope of ~0.8 MiB/token/card
1864    /// (vramwatch.csv: +6.3 GiB across the 8,072-token prime) — the formula sits above the
1865    /// measurement, never below it. The ctx-coupled score plane and the prompt-long hiddens
1866    /// stack are separate coefficients on the shape; see [`HyperPrimeWorkspaceShape`].
1867    pub fn hyper_prime_workspace_shape(&self) -> Option<HyperPrimeWorkspaceShape> {
1868        let topology = self.hyper.as_ref()?;
1869        let h = self.cfg.n_embd as usize;
1870        let s = topology.streams;
1871        let f32b = std::mem::size_of::<f32>();
1872        // Stream state (x2) + ppN boundary slots (x2), pre/norm transients, prime tail.
1873        let mut chunk_token_bytes = 4 * s * h * f32b + 4 * h * f32b + 2 * h * f32b;
1874        if let Some(moe) = self.cfg.moe.as_ref() {
1875            let u = moe.expert_used_count as usize;
1876            let f = moe.expert_ff_length as usize;
1877            chunk_token_bytes += u * (10 * h + 14 * f) + 4 * h;
1878        }
1879        let mut kpool_score_pool = 0;
1880        if let Some(glm5) = self.cfg.glm5.as_ref() {
1881            let heads = self.cfg.n_head as usize;
1882            chunk_token_bytes +=
1883                heads * (glm5.qk_head_dim as usize + glm5.v_head_dim as usize) * f32b;
1884            if glm5.index_kpool > 0 {
1885                chunk_token_bytes += (glm5.index_topk as usize / glm5.index_kpool as usize + 1)
1886                    * std::mem::size_of::<i32>();
1887                kpool_score_pool = glm5.index_kpool as usize;
1888            }
1889        }
1890        Some(HyperPrimeWorkspaceShape {
1891            chunk_token_bytes,
1892            prompt_bytes_per_token: h * f32b,
1893            kpool_score_pool,
1894            n_layers: self.layers.len(),
1895            gdn_grid: self.gdn_prime_grid_on(),
1896        })
1897    }
1898
1899    /// One T=1 decode step under the mHC residual: `decode_step_h`'s contract over the hc layer
1900    /// program. The stream state is INTRA-STEP — expanded from the embedded row, collapsed for
1901    /// the logits — so no cache format changes and the mixers keep their own state exactly as
1902    /// they do on the serial path.
1903    pub(crate) fn decode_step_hyper(
1904        &self,
1905        e: &Engine,
1906        token: u32,
1907        cache: &mut Cache,
1908    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1909        let topology = *self
1910            .hyper
1911            .as_ref()
1912            .ok_or("decode_step_hyper on a model with no HyperConnections topology")?;
1913        // M2 ppN door for the mHC decode step (see `forward_hyper`'s note). This is the door
1914        // the GLM-5.3-Flash residency arc turns on: with it shut, every routed expert has to
1915        // fit beside card 0's trunk, which 171.2 GB of experts cannot do.
1916        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1917            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
1918                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
1919            }
1920            return self.decode_step_hyper_ppn(e, token, cache, &topology, &fence);
1921        }
1922        let n_embd = self.cfg.n_embd as usize;
1923        let eps = self.cfg.rms_eps;
1924        let pos = cache.pos;
1925        let pos_d = e.htod_i32(&[pos as i32])?;
1926
1927        let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
1928        let mut x = crate::hyper::expand(e, &topology, &embedded, 1, n_embd)?;
1929
1930        x = self.hyper_range_decode(e, &topology, x, 0, self.layers.len(), &pos_d, pos, cache)?;
1931
1932        // SHARED EXIT with the ppN twin (see `forward_hyper`'s note).
1933        self.hyper_decode_tail(e, &topology, &x, n_embd, eps, cache)
1934    }
1935
1936    /// One hc layer RANGE `[lo, hi)` of the STATELESS prefill walk, driven by engine `e`.
1937    ///
1938    /// Extracted so the unsplit walk and every pipeline stage run the SAME code over their own
1939    /// range: the ppN arm's bit-identity claim is then structural, not a coincidence of two
1940    /// hand-kept-in-sync copies. `x` enters and leaves as the `[t, streams, hidden]` stream
1941    /// state, which is exactly the payload a stage boundary transports.
1942    #[allow(clippy::too_many_arguments)]
1943    fn hyper_range_forward(
1944        &self,
1945        e: &Engine,
1946        topology: &crate::hyper::HyperTopology,
1947        mut x: CudaSlice<f32>,
1948        lo: usize,
1949        hi: usize,
1950        pos_d: &CudaSlice<i32>,
1951        t: usize,
1952        trace: bool,
1953    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1954        let n_embd = self.cfg.n_embd as usize;
1955        let eps = self.cfg.rms_eps;
1956        for il in lo..hi {
1957            let layer = &self.layers[il];
1958            let hyper = layer.hyper.as_ref().ok_or_else(|| {
1959                format!("layer {il} carries no hyper-connection weights under an hc plan")
1960            })?;
1961
1962            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
1963            let mut h = e.uninit(t * n_embd)?;
1964            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1965            let mixed = match &layer.mixer {
1966                Mixer::Full(fa) => self.full_attn(e, fa, &h, pos_d, t, il)?,
1967                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
1968                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, pos_d, t, il)?,
1969                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
1970            };
1971            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
1972            if trace {
1973                let index = il as i64;
1974                memra_reference::hidden_trace::emit_last_row(
1975                    "mixer",
1976                    index,
1977                    t,
1978                    n_embd,
1979                    &e.dtoh(&mixed)?,
1980                );
1981                let streams = x.len() / (t * n_embd);
1982                memra_reference::hidden_trace::emit_last_row(
1983                    "attn",
1984                    index,
1985                    t,
1986                    streams * n_embd,
1987                    &e.dtoh(&x)?,
1988                );
1989            }
1990
1991            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
1992            let mut z = e.uninit(t * n_embd)?;
1993            e.rms_norm(
1994                &y,
1995                layer.post_attn_norm.float_data(),
1996                &mut z,
1997                n_embd,
1998                t,
1999                eps,
2000            )?;
2001            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2002            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2003            if trace {
2004                let index = il as i64;
2005                memra_reference::hidden_trace::emit_last_row(
2006                    "ffn",
2007                    index,
2008                    t,
2009                    n_embd,
2010                    &e.dtoh(&ffn_out)?,
2011                );
2012                let streams = x.len() / (t * n_embd);
2013                memra_reference::hidden_trace::emit_last_row(
2014                    "layer",
2015                    index,
2016                    t,
2017                    streams * n_embd,
2018                    &e.dtoh(&x)?,
2019                );
2020            }
2021        }
2022        Ok(x)
2023    }
2024
2025    /// One hc layer RANGE `[lo, hi)` of the STATEFUL prime walk (see `hyper_range_forward`).
2026    /// Every mixer writes its own layer's cache state through `e`, so under the ppN door a
2027    /// stage's KDA conv ring / delta-rule state and its MLA latent rows + kpool indexer plane
2028    /// are written by the SAME engine `pp::new_cache` allocated them on.
2029    #[allow(clippy::too_many_arguments)]
2030    fn hyper_range_prime(
2031        &self,
2032        e: &Engine,
2033        topology: &crate::hyper::HyperTopology,
2034        mut x: CudaSlice<f32>,
2035        lo: usize,
2036        hi: usize,
2037        pos_d: &CudaSlice<i32>,
2038        t: usize,
2039        cache: &mut Cache,
2040        seq_end: usize,
2041    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2042        let n_embd = self.cfg.n_embd as usize;
2043        let eps = self.cfg.rms_eps;
2044        for il in lo..hi {
2045            let layer = &self.layers[il];
2046            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2047                format!("layer {il} carries no hyper-connection weights under an hc plan")
2048            })?;
2049
2050            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, t, n_embd)?;
2051            let mut h = e.uninit(t * n_embd)?;
2052            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2053            let mixed = match &layer.mixer {
2054                Mixer::Full(fa) => {
2055                    self.full_attn_prime(e, fa, &h, None, pos_d, t, cache, il, seq_end)?
2056                }
2057                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2058                Mixer::Mla(mla) if mla.tp.is_some() => {
2059                    self.mla_tp_attn_cached(e, mla, &h, pos_d, t, il, cache, false)?
2060                }
2061                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, t, il, cache)?,
2062                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2063                    e,
2064                    la,
2065                    &h,
2066                    t,
2067                    eps,
2068                    cache,
2069                    il,
2070                    crate::kda::ConvArm::Prefill,
2071                )?,
2072                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2073            };
2074            x = crate::hyper::post(e, topology, &mixed, &x, &mix, t, n_embd)?;
2075
2076            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, t, n_embd)?;
2077            let mut z = e.uninit(t * n_embd)?;
2078            e.rms_norm(
2079                &y,
2080                layer.post_attn_norm.float_data(),
2081                &mut z,
2082                n_embd,
2083                t,
2084                eps,
2085            )?;
2086            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2087            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, t, n_embd)?;
2088            // glm5 DFlash2 feature tap (lane/glm5-dflash-draft-src): the CONTRACTED
2089            // completed layer output, host-staged — one Option check when unarmed.
2090            self.glm5_hc_tap(e, cache, topology, il, &x, t)?;
2091        }
2092        Ok(x)
2093    }
2094
2095    /// One hc layer RANGE `[lo, hi)` of the T=1 decode step (see `hyper_range_forward`).
2096    #[allow(clippy::too_many_arguments)]
2097    fn hyper_range_decode(
2098        &self,
2099        e: &Engine,
2100        topology: &crate::hyper::HyperTopology,
2101        mut x: CudaSlice<f32>,
2102        lo: usize,
2103        hi: usize,
2104        pos_d: &CudaSlice<i32>,
2105        pos: usize,
2106        cache: &mut Cache,
2107    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2108        // MEMRA_HC_DECODE_WS=1 (default OFF, read per call — rollback seam): the persistent-
2109        // workspace twin of this walk. Same kernels, same call order, same operand bytes;
2110        // only the hc-glue allocations (mixes/gates/comb/y/h/z/post-out, ~12 alloc+free pairs
2111        // per layer per token of the census's 2,358) disappear. Byte identity ON/OFF is gated
2112        // by hc_decode_ws_gpu.rs; refusal shapes fall through to the allocating walk below.
2113        if hyper_decode_ws_on() {
2114            return self.hyper_range_decode_ws(e, topology, x, lo, hi, pos_d, pos, cache);
2115        }
2116        let n_embd = self.cfg.n_embd as usize;
2117        let eps = self.cfg.rms_eps;
2118        for il in lo..hi {
2119            let layer = &self.layers[il];
2120            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2121                format!("layer {il} carries no hyper-connection weights under an hc plan")
2122            })?;
2123
2124            let (y, mix) = crate::hyper::pre(e, topology, &hyper.attn, &x, 1, n_embd)?;
2125            let mut h = e.uninit(n_embd)?;
2126            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
2127            let mixed = match &layer.mixer {
2128                Mixer::Full(fa) => self.full_attn_decode(e, fa, &h, pos_d, pos, cache, il)?,
2129                Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il)?,
2130                Mixer::Mla(mla) if mla.tp.is_some() => {
2131                    self.mla_tp_attn_cached(e, mla, &h, pos_d, 1, il, cache, false)?
2132                }
2133                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache)?,
2134                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2135                    e,
2136                    la,
2137                    &h,
2138                    1,
2139                    eps,
2140                    cache,
2141                    il,
2142                    crate::kda::ConvArm::Decode,
2143                )?,
2144                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?,
2145            };
2146            x = crate::hyper::post(e, topology, &mixed, &x, &mix, 1, n_embd)?;
2147
2148            let (y, mix) = crate::hyper::pre(e, topology, &hyper.mlp, &x, 1, n_embd)?;
2149            let mut z = e.uninit(n_embd)?;
2150            e.rms_norm(
2151                &y,
2152                layer.post_attn_norm.float_data(),
2153                &mut z,
2154                n_embd,
2155                1,
2156                eps,
2157            )?;
2158            let ffn_out = self.hyper_ffn_branch(e, layer, &z, 1, il, false)?;
2159            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, 1, n_embd)?;
2160        }
2161        Ok(x)
2162    }
2163
2164    /// The persistent-workspace twin of `hyper_range_decode` (MEMRA_HC_DECODE_WS, lever 2 of
2165    /// the decode diet). One `HyperDecodeWs` per engine (so each ppN stage owns its own,
2166    /// allocated on its own device); the walk TAKES it from the engine pool, rotates the
2167    /// stream state against `ws.xb` (an ownership swap, not a copy), and puts it back. The
2168    /// mixers and the FFN/MoE branches are the SAME calls with the SAME inputs — their
2169    /// internal allocations are untouched by this lever.
2170    #[allow(clippy::too_many_arguments)]
2171    fn hyper_range_decode_ws(
2172        &self,
2173        e: &Engine,
2174        topology: &crate::hyper::HyperTopology,
2175        x: CudaSlice<f32>,
2176        lo: usize,
2177        hi: usize,
2178        pos_d: &CudaSlice<i32>,
2179        pos: usize,
2180        cache: &mut Cache,
2181    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2182        let n_embd = self.cfg.n_embd as usize;
2183        let mut ws = match e.hyper_ws_take() {
2184            Some(ws) if ws.matches(topology, n_embd) => ws,
2185            _ => crate::hyper::HyperDecodeWs::new(e, topology, n_embd)?,
2186        };
2187        if HC_DECODE_WS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
2188            eprintln!(
2189                "[hc-decode-ws] engaged streams={} hidden={n_embd} (persistent hc-glue \
2190                 workspace, per-engine pool; MEMRA_HC_DECODE_WS=1)",
2191                topology.streams
2192            );
2193        }
2194        let out =
2195            self.hyper_range_decode_ws_body(e, topology, x, lo, hi, pos_d, pos, cache, &mut ws);
2196        e.hyper_ws_put(ws);
2197        out
2198    }
2199
2200    /// The walk itself — `hyper_range_decode`'s loop with the hc glue landing in `ws`.
2201    /// KEPT CALL-FOR-CALL IN STEP with the allocating walk above: same kernels, same order
2202    /// (pre -> rms_norm -> mixer -> post -> pre -> rms_norm -> ffn -> post), so the
2203    /// byte-identity gate is a structural claim, not a coincidence.
2204    #[allow(clippy::too_many_arguments)]
2205    fn hyper_range_decode_ws_body(
2206        &self,
2207        e: &Engine,
2208        topology: &crate::hyper::HyperTopology,
2209        mut x: CudaSlice<f32>,
2210        lo: usize,
2211        hi: usize,
2212        pos_d: &CudaSlice<i32>,
2213        pos: usize,
2214        cache: &mut Cache,
2215        ws: &mut crate::hyper::HyperDecodeWs,
2216    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2217        let n_embd = self.cfg.n_embd as usize;
2218        let eps = self.cfg.rms_eps;
2219        for il in lo..hi {
2220            let layer = &self.layers[il];
2221            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2222                format!("layer {il} carries no hyper-connection weights under an hc plan")
2223            })?;
2224
2225            crate::hyper::pre_t1_ws(e, topology, &hyper.attn, &x, ws, n_embd)?;
2226            e.rms_norm(
2227                &ws.y,
2228                layer.attn_norm.float_data(),
2229                &mut ws.h,
2230                n_embd,
2231                1,
2232                eps,
2233            )?;
2234            let mixed = match &layer.mixer {
2235                Mixer::Full(fa) => self.full_attn_decode(e, fa, &ws.h, pos_d, pos, cache, il)?,
2236                Mixer::Linear(la) => self.linear_attn_decode(e, la, &ws.h, cache, il)?,
2237                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &ws.h, pos_d, 1, il, cache)?,
2238                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &ws.h, eps, cache, il)?,
2239            };
2240            crate::hyper::post_t1_ws(e, topology, &mixed, &x, ws, n_embd)?;
2241            std::mem::swap(&mut x, &mut ws.xb);
2242
2243            crate::hyper::pre_t1_ws(e, topology, &hyper.mlp, &x, ws, n_embd)?;
2244            e.rms_norm(
2245                &ws.y,
2246                layer.post_attn_norm.float_data(),
2247                &mut ws.z,
2248                n_embd,
2249                1,
2250                eps,
2251            )?;
2252            let ffn_out = self.hyper_ffn_branch(e, layer, &ws.z, 1, il, false)?;
2253            crate::hyper::post_t1_ws(e, topology, &ffn_out, &x, ws, n_embd)?;
2254            std::mem::swap(&mut x, &mut ws.xb);
2255        }
2256        Ok(x)
2257    }
2258
2259    /// One hc layer RANGE `[lo, hi)` of the BATCHED T=1 decode step: B independent sessions
2260    /// share one walk over the `[B, streams, n_embd]` stream state. The batched twin of
2261    /// `hyper_range_decode`, and the trunk of `decode_step_batch_hyper` (decode_batch.rs).
2262    ///
2263    /// SHAPE — batched where the arithmetic is row-independent, per-session where the state
2264    /// is, decode-exact where a reduction is width-dependent:
2265    ///
2266    ///   * The hc glue (`expand`/`pre_finish` kernels/`post`) is block-per-token by
2267    ///     construction (grid over t), so t=B batches it with per-row bytes unchanged.
2268    ///   * The hc mixing GEMM is the ONE width-dependent reduction in the glue
2269    ///     (cuBLASLt's n-dependent split — the lt_ndep probe), so this walk calls
2270    ///     `hyper::pre_exact`, which runs each row through the m=1 program the serial step
2271    ///     runs. rms_norm at m=B is a per-row program.
2272    ///   * The MIXERS (KDA conv ring + delta rule, MLA latent rows + kpool indexer plane,
2273    ///     and the Full/Linear classes for completeness) hold per-session recurrent or
2274    ///     latent state, so each session's row is routed to ITS OWN cache through the SAME
2275    ///     t=1 call its solo step makes — the per-seq loop is the v1 exactness doctrine
2276    ///     from this module's sibling (`decode_batch.rs` header), and the row copies in and
2277    ///     out are arithmetic-free materializations.
2278    ///   * The FFN batches at t=B: the MoE body's router is the fixed per-row program at
2279    ///     t < PRIME_MIN_T, expert dispatch is per-token, and the shexp trio rides the
2280    ///     per-column decode-exact arm at decode widths; the dense branch runs per-row so
2281    ///     each row executes the serial `hyper_ffn_branch` program verbatim.
2282    ///
2283    /// EXACTNESS BAR: row b of a B-row step must be BIT-IDENTICAL to session b decoding
2284    /// alone through `decode_step_hyper` — full-logit compare, per step. Gate:
2285    /// `glm5-hyper-batch-gate` (fixture-driven, red-armed with a swapped-row and a
2286    /// wrong-cache-slot mutation; receipts in
2287    /// `research/glm53-flash-bringup-20260827/batched-decode-gate/`).
2288    ///
2289    /// `pos_rows[bi]` is session bi's single-position device buffer, uploaded by the caller
2290    /// through THIS range's engine (the per-stage pos_d law under a pp split). `caches[bi]`
2291    /// advances exactly as its solo step would; `cache.pos` itself is bumped by the caller's
2292    /// epilogue, never here.
2293    #[allow(clippy::too_many_arguments)]
2294    pub(crate) fn hyper_batch_range_decode(
2295        &self,
2296        e: &Engine,
2297        topology: &crate::hyper::HyperTopology,
2298        mut x: CudaSlice<f32>,
2299        lo: usize,
2300        hi: usize,
2301        pos_rows: &[CudaSlice<i32>],
2302        caches: &mut [&mut Cache],
2303    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2304        let b_n = caches.len();
2305        assert_eq!(
2306            pos_rows.len(),
2307            b_n,
2308            "hyper_batch_range_decode: pos_rows built for a different batch width"
2309        );
2310        let n_embd = self.cfg.n_embd as usize;
2311        let eps = self.cfg.rms_eps;
2312        for il in lo..hi {
2313            let layer = &self.layers[il];
2314            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2315                format!("layer {il} carries no hyper-connection weights under an hc plan")
2316            })?;
2317
2318            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.attn, &x, b_n, n_embd)?;
2319            let mut h = e.uninit(b_n * n_embd)?;
2320            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, b_n, eps)?;
2321            // ---- mixers: per-session, each row through its OWN cache (the t=1 serial
2322            // program — cross-session contamination here is the failure mode the gate's
2323            // swapped-row mutation exists to catch) ----
2324            let mut mixed = e.uninit(b_n * n_embd)?;
2325            for bi in 0..b_n {
2326                let mut h_row = e.uninit(n_embd)?;
2327                e.dtod_copy_view(&h.slice(bi * n_embd..(bi + 1) * n_embd), &mut h_row)?;
2328                let cache: &mut Cache = &mut *caches[bi];
2329                let pos = cache.pos;
2330                let out_row = match &layer.mixer {
2331                    Mixer::Full(fa) => {
2332                        self.full_attn_decode(e, fa, &h_row, &pos_rows[bi], pos, cache, il)?
2333                    }
2334                    Mixer::Linear(la) => self.linear_attn_decode(e, la, &h_row, cache, il)?,
2335                    Mixer::Mla(mla) => {
2336                        self.mla_attn_cached(e, mla, &h_row, &pos_rows[bi], 1, il, cache)?
2337                    }
2338                    Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h_row, eps, cache, il)?,
2339                };
2340                e.copy_into(&mut mixed, bi * n_embd, &out_row, n_embd)?;
2341            }
2342            x = crate::hyper::post(e, topology, &mixed, &x, &mix, b_n, n_embd)?;
2343
2344            let (y, mix) = crate::hyper::pre_exact(e, topology, &hyper.mlp, &x, b_n, n_embd)?;
2345            let mut z = e.uninit(b_n * n_embd)?;
2346            e.rms_norm(
2347                &y,
2348                layer.post_attn_norm.float_data(),
2349                &mut z,
2350                n_embd,
2351                b_n,
2352                eps,
2353            )?;
2354            let ffn_out = self.hyper_ffn_branch_batch(e, layer, &z, b_n, il, false)?;
2355            x = crate::hyper::post(e, topology, &ffn_out, &x, &mix, b_n, n_embd)?;
2356        }
2357        Ok(x)
2358    }
2359
2360    /// The FFN branch of the batched hc decode walk (see `hyper_batch_range_decode`).
2361    ///
2362    /// MoE batches at t=B — that is the weight win this walk exists for (one stream of the
2363    /// routed experts serves B rows), and every stage of `moe_ffn_il_zq8` is per-row exact
2364    /// at decode widths (router: fixed per-row program at t < PRIME_MIN_T; experts:
2365    /// per-(token,expert) programs; shexp: the per-column decode-exact arm). The DENSE
2366    /// branch runs PER ROW through the serial `hyper_ffn_branch` instead: its
2367    /// `matmul_group` dispatch carries no per-row bit-identity contract across widths for
2368    /// every weight class this walk must serve, and a first-k-dense plan carries one such
2369    /// layer — per-row costs nothing and each row executes the solo step's program verbatim.
2370    /// `vrows` (lane/glm5-vrest): the VERIFY walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`)
2371    /// passes `true`, which lets the MoE body take the pairs-shaped batched routed-expert
2372    /// program across the t rows (`moe_vrows_pairs_q8` — bit-identical per row, fail-closed
2373    /// to the sequential loop for every unqualified shape). The batched DECODE walk
2374    /// (`decode_step_batch_hyper`) passes `false` — its priced dispatch class stays
2375    /// byte-stable; porting it is a named follow-up with its own re-price. The DENSE branch
2376    /// is per-row in both arms (its `matmul_group` dispatch carries no cross-width per-row
2377    /// bit-identity contract for every weight class this walk must serve; ~3 layers, named
2378    /// out of scope in the vrest attribution).
2379    pub(crate) fn hyper_ffn_branch_batch(
2380        &self,
2381        e: &Engine,
2382        layer: &crate::hybrid::HybridLayer,
2383        z: &CudaSlice<f32>,
2384        b_n: usize,
2385        il: usize,
2386        vrows: bool,
2387    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2388        let n_embd = self.cfg.n_embd as usize;
2389        match &layer.ffn {
2390            crate::hybrid::Ffn::Dense { .. } => {
2391                let mut out = e.uninit(b_n * n_embd)?;
2392                for bi in 0..b_n {
2393                    let mut z_row = e.uninit(n_embd)?;
2394                    e.dtod_copy_view(&z.slice(bi * n_embd..(bi + 1) * n_embd), &mut z_row)?;
2395                    let row = self.hyper_ffn_branch(e, layer, &z_row, 1, il, false)?;
2396                    e.copy_into(&mut out, bi * n_embd, &row, n_embd)?;
2397                }
2398                Ok(out)
2399            }
2400            crate::hybrid::Ffn::Moe(m) => {
2401                if vrows {
2402                    self.moe_ffn_il_zq8_vrows(e, m, z, b_n, il as u16)
2403                } else {
2404                    self.moe_ffn_il_zq8(e, m, z, None, b_n, il as u16)
2405                }
2406            }
2407        }
2408    }
2409
2410    // =============================== M2 ppN, mHC arm ===============================
2411    //
2412    // The three walks below are the hc twins of `decode_step_h_ppn`. They exist because the
2413    // GLM-5.3-Flash residency arithmetic does not close on one card: 171.2 GB of routed
2414    // experts against 2x96 GB means the second card is the only route to full residency, and
2415    // the pp door is how weights get there. Until these landed, all three hc walks refused
2416    // the door outright ("the sharded stage handoff is unwired for this residual topology"),
2417    // which is a loud refusal and was the right behaviour — a single-engine walk over
2418    // stage-sharded weights dereferences another device's pointers.
2419    //
2420    // WHAT IS DIFFERENT FROM THE GENERIC ARM, and it is exactly one thing: the payload on the
2421    // wire. The serial trunk hands `[n_embd]` (decode) or `[t, n_embd]` (prime) across a
2422    // boundary; the mHC trunk carries `streams` residual streams between layers, so the
2423    // boundary payload is `[streams, n_embd]` / `[t, streams, n_embd]`. `pp.rs`'s BoundarySlot
2424    // buffers are lazily sized from the caller's `n` and grow to the high-water mark, so no
2425    // slot-sizing change was needed for that — the wider payload just makes them wider.
2426    // `hyper::expand` runs on stage 0 (it takes no weights) and `hyper::collapse` +
2427    // output_norm + lm head on the last stage, which is where the loader already put the head
2428    // (`pp::layer_engine(e, n_trunk, n_trunk - 1)` in `hybrid.rs`) and, under
2429    // `HcCollapse::GatedHead`, the head trio.
2430    //
2431    // Per-layer state placement needed NO new contract: `pp::new_cache` already picks the
2432    // owning stage's `KvDev` per layer for all three of glm5_next's state classes
2433    // (`Recurrent` = KDA conv ring + delta-rule state, `LatentKvCache` = MLA rows + the kpool
2434    // indexer plane, `KvCache` = full attention), and the kpool `index_pool_keys` plane is
2435    // lazily allocated through the engine the mixer is called with, which under these walks is
2436    // the stage's engine. `glm5-hyper-ppn-gate` asserts the fence actually separates those
2437    // classes across stages, so that is a tested property rather than an argued one.
2438    //
2439    // NOT WIRED, and refused rather than approximated: the deferred-readback (pipelined) arm.
2440    // `decode_step_h_ppn_deferred` calls `refuse_hyper`, and this lane did not change that.
2441    //
2442    // Gate: `glm5-hyper-ppn-gate` (bit-identical logits vs the unsplit hc walk, decode and
2443    // prime, at every N/knob combination), receipts in
2444    // `research/glm53-flash-bringup-20260827/ppn-hyper-gate/`.
2445
2446    /// ppN twin of `forward_hyper`: the stateless prefill as N stage subgraphs.
2447    fn forward_hyper_ppn(
2448        &self,
2449        e: &Engine,
2450        tokens: &[u32],
2451        last_only: bool,
2452        topology: &crate::hyper::HyperTopology,
2453        fence: &[usize],
2454    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2455        let n_embd = self.cfg.n_embd as usize;
2456        let eps = self.cfg.rms_eps;
2457        let t = tokens.len();
2458        let width = topology.streams * n_embd;
2459        let trace = memra_reference::hidden_trace::enabled();
2460        if trace {
2461            memra_reference::hidden_trace::emit_tokens(tokens);
2462        }
2463        let pos: Vec<i32> = (0..t as i32).collect();
2464
2465        if crate::pp::pp2_streams_off() {
2466            // Same-stream rollback seam: one engine, one stream, an explicit copy pair per
2467            // boundary. Structurally identical to the split walk, which is the point — it is
2468            // the arm that says "the split is the split, not the streams".
2469            let pos_d = e.htod_i32(&pos)?;
2470            let embedded = self.embed(e, tokens)?;
2471            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2472            x = self.hyper_range_forward(e, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2473            for s in 1..fence.len() - 1 {
2474                let boundary_tx = e.clone_dtod(&x)?;
2475                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2476                x = self.hyper_range_forward(
2477                    e,
2478                    topology,
2479                    boundary_rx,
2480                    fence[s],
2481                    fence[s + 1],
2482                    &pos_d,
2483                    t,
2484                    trace,
2485                )?;
2486            }
2487            return self.hyper_head_logits(e, topology, &x, t, n_embd, eps, last_only);
2488        }
2489
2490        let rt = crate::pp::PpNRt::get(e)?;
2491        let n_st = fence.len() - 1;
2492        assert_eq!(
2493            rt.n_stages(),
2494            n_st,
2495            "PpNRt stage count {} != fence stages {n_st}",
2496            rt.n_stages()
2497        );
2498        // #87 REVERSE PUBLICATION: order every stage stream behind the caller's stream before
2499        // the first stage allocation (anatomy: `PpNRt::fence_stages_behind`).
2500        rt.fence_stages_behind(&e.stream())?;
2501
2502        let mut slot = {
2503            let _st0 = rt.enter(0);
2504            let e0 = rt.engine(0, e);
2505            // PER-STAGE pos_d (M2 pipelining law): each stage uploads its OWN copy on ITS
2506            // stream, so the buffer is allocated, consumed and freed on one stream.
2507            let pos_d = e0.htod_i32(&pos)?;
2508            let embedded = self.embed(e0, tokens)?;
2509            let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2510            let x =
2511                self.hyper_range_forward(e0, topology, x, fence[0], fence[1], &pos_d, t, trace)?;
2512            rt.tx(0, &x, t * width)?
2513        };
2514        for s in 1..n_st - 1 {
2515            let _st = rt.enter(s);
2516            let es = rt.engine(s, e);
2517            let pos_d = es.htod_i32(&pos)?;
2518            let x = rt.rx(s - 1, slot, t * width)?;
2519            let x = self.hyper_range_forward(
2520                es,
2521                topology,
2522                x,
2523                fence[s],
2524                fence[s + 1],
2525                &pos_d,
2526                t,
2527                trace,
2528            )?;
2529            slot = rt.tx(s, &x, t * width)?;
2530        }
2531        let _stl = rt.enter(n_st - 1);
2532        let el = rt.engine(n_st - 1, e);
2533        let pos_d = el.htod_i32(&pos)?;
2534        let x = rt.rx(n_st - 2, slot, t * width)?;
2535        let x = self.hyper_range_forward(
2536            el,
2537            topology,
2538            x,
2539            fence[n_st - 1],
2540            fence[n_st],
2541            &pos_d,
2542            t,
2543            trace,
2544        )?;
2545        self.hyper_head_logits(el, topology, &x, t, n_embd, eps, last_only)
2546    }
2547
2548    /// Trunk exit shared by `forward_hyper` and its ppN twin: collapse the stream state, apply
2549    /// `output_norm`, and project. Runs on the LAST stage's engine under the pp door, which is
2550    /// where `hybrid.rs` uploaded `output_norm`, the lm head and (under `HcCollapse::GatedHead`)
2551    /// the head trio.
2552    #[allow(clippy::too_many_arguments)]
2553    fn hyper_head_logits(
2554        &self,
2555        e: &Engine,
2556        topology: &crate::hyper::HyperTopology,
2557        x: &CudaSlice<f32>,
2558        t: usize,
2559        n_embd: usize,
2560        eps: f32,
2561        last_only: bool,
2562    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2563        let collapsed =
2564            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2565        if memra_reference::hidden_trace::enabled() {
2566            memra_reference::hidden_trace::emit_last_row(
2567                "collapse",
2568                -1,
2569                t,
2570                n_embd,
2571                &e.dtoh(&collapsed)?,
2572            );
2573        }
2574        let mut hn = e.uninit(t * n_embd)?;
2575        e.rms_norm(
2576            &collapsed,
2577            self.output_norm.float_data(),
2578            &mut hn,
2579            n_embd,
2580            t,
2581            eps,
2582        )?;
2583        let logits = if last_only {
2584            let last = e.view(&hn, t * n_embd);
2585            let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2586            let mut hlast = e.uninit(n_embd)?;
2587            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2588            e.matmul(&self.output, &hlast, 1)?
2589        } else {
2590            e.matmul(&self.output, &hn, t)?
2591        };
2592        e.dtoh(&logits)
2593    }
2594
2595    /// ppN twin of `prime_cache_hyper`: the monolithic stateful prime as N stage subgraphs.
2596    /// The returned device buffers (`h_seed`, `hiddens`) are owned by the LAST stage's engine,
2597    /// the same contract `decode_step_h_ppn` publishes for its `h_seed`.
2598    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2599    #[allow(clippy::too_many_arguments)]
2600    fn prime_cache_hyper_ppn(
2601        &self,
2602        e: &Engine,
2603        tokens: &[u32],
2604        cache: &mut Cache,
2605        queued_after: usize,
2606        topology: &crate::hyper::HyperTopology,
2607        fence: &[usize],
2608        overlay: Option<&crate::vision::EmbedOverlay>,
2609    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2610        let n_embd = self.cfg.n_embd as usize;
2611        let eps = self.cfg.rms_eps;
2612        let t = tokens.len();
2613        let width = topology.streams * n_embd;
2614        if cache.pos + t > cache.max_ctx {
2615            return Err("prime_cache: prompt exceeds cache max_ctx".into());
2616        }
2617        let seq_end = cache.pos + t + queued_after;
2618        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2619        // glm5 DFlash2 feature tap: set ONCE per call, before any stage walk — every stage
2620        // range of this chunk writes the same rows at the chunk's absolute offset.
2621        if let Some(sink) = cache.hc_taps.as_mut() {
2622            sink.base = cache.pos;
2623        }
2624
2625        if crate::pp::pp2_streams_off() {
2626            let pos_d = e.htod_i32(&pos)?;
2627            let mut embedded = self.embed(e, tokens)?;
2628            if let Some(ov) = overlay {
2629                // Mixed-embedding splice at embedding intake (the same point the streams-on
2630                // arm and prime_chunk_hyper use). The caller's overlay is already windowed
2631                // to THIS call (prefill_tick rebases spans call-relative), so chunk_off = 0.
2632                ov.splice_into(e, &mut embedded, 0, t, n_embd)?;
2633            }
2634            let mut x = crate::hyper::expand(e, topology, &embedded, t, n_embd)?;
2635            x = self.hyper_range_prime(
2636                e, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2637            )?;
2638            for s in 1..fence.len() - 1 {
2639                let boundary_tx = e.clone_dtod(&x)?;
2640                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2641                x = self.hyper_range_prime(
2642                    e,
2643                    topology,
2644                    boundary_rx,
2645                    fence[s],
2646                    fence[s + 1],
2647                    &pos_d,
2648                    t,
2649                    cache,
2650                    seq_end,
2651                )?;
2652            }
2653            return self.hyper_prime_tail(e, topology, &x, t, n_embd, eps, cache);
2654        }
2655
2656        {
2657            let rt = crate::pp::PpNRt::get(e)?;
2658            let n_st = fence.len() - 1;
2659            assert_eq!(
2660                rt.n_stages(),
2661                n_st,
2662                "PpNRt stage count {} != fence stages {n_st}",
2663                rt.n_stages()
2664            );
2665            rt.fence_stages_behind(&e.stream())?;
2666            // OVERLAY DEVICE LAW: the overlay rows are built by the caller on the PRIMARY
2667            // engine (build_vision_overlay runs the tower there). Stage 0 keeps the primary
2668            // engine whenever it lives on the primary device (pp::PpNRt::build); a placement
2669            // that moves stage 0 to another device would make the splice a cross-device copy
2670            // with no receipts — refuse loudly instead of silently peer-reading.
2671            if overlay.is_some() && !std::ptr::eq(rt.engine(0, e), e) {
2672                return Err(
2673                    "vision embedding overlay requires stage 0 on the primary device \
2674                     (overlay rows are primary-resident); place stage 0 on the primary \
2675                     or run MEMRA_PP_STREAMS=0"
2676                        .into(),
2677                );
2678            }
2679            let mut slot = {
2680                let _st0 = rt.enter(0);
2681                let e0 = rt.engine(0, e);
2682                let pos_d = e0.htod_i32(&pos)?;
2683                let mut embedded = self.embed(e0, tokens)?;
2684                if let Some(ov) = overlay {
2685                    // Mixed-embedding splice at stage-0 embedding intake, BEFORE stream
2686                    // expansion — the reference's execute_multimodal splice point. Stages
2687                    // s > 0 only ever see the [t, streams, hidden] boundary payload, so no
2688                    // other stage carries overlay arithmetic. chunk_off = 0: the caller's
2689                    // overlay is already windowed to this call.
2690                    ov.splice_into(e0, &mut embedded, 0, t, n_embd)?;
2691                }
2692                let x = crate::hyper::expand(e0, topology, &embedded, t, n_embd)?;
2693                let x = self.hyper_range_prime(
2694                    e0, topology, x, fence[0], fence[1], &pos_d, t, cache, seq_end,
2695                )?;
2696                rt.tx(0, &x, t * width)?
2697            };
2698            for s in 1..n_st - 1 {
2699                let _st = rt.enter(s);
2700                let es = rt.engine(s, e);
2701                let pos_d = es.htod_i32(&pos)?;
2702                let x = rt.rx(s - 1, slot, t * width)?;
2703                let x = self.hyper_range_prime(
2704                    es,
2705                    topology,
2706                    x,
2707                    fence[s],
2708                    fence[s + 1],
2709                    &pos_d,
2710                    t,
2711                    cache,
2712                    seq_end,
2713                )?;
2714                slot = rt.tx(s, &x, t * width)?;
2715            }
2716            let _stl = rt.enter(n_st - 1);
2717            let el = rt.engine(n_st - 1, e);
2718            let pos_d = el.htod_i32(&pos)?;
2719            let x = rt.rx(n_st - 2, slot, t * width)?;
2720            let x = self.hyper_range_prime(
2721                el,
2722                topology,
2723                x,
2724                fence[n_st - 1],
2725                fence[n_st],
2726                &pos_d,
2727                t,
2728                cache,
2729                seq_end,
2730            )?;
2731            self.hyper_prime_tail(el, topology, &x, t, n_embd, eps, cache)
2732        }
2733    }
2734
2735    /// ONE call of the mHC prime walk. Carries `tokens.len()` rows of stream state and of every
2736    /// per-layer transient, appends this call's rows to the mixers' own state, and advances
2737    /// `cache.pos`. `seq_end` is the REQUEST's absolute end, passed in rather than recomputed,
2738    /// so no arithmetic here is a function of how the prompt was split.
2739    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2740    fn prime_chunk_hyper(
2741        &self,
2742        e: &Engine,
2743        tokens: &[u32],
2744        cache: &mut Cache,
2745        seq_end: usize,
2746        chunk_off: usize,
2747        overlay: Option<&crate::vision::EmbedOverlay>,
2748    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2749        let topology = *self
2750            .hyper
2751            .as_ref()
2752            .ok_or("prime_chunk_hyper on a model with no HyperConnections topology")?;
2753        let n_embd = self.cfg.n_embd as usize;
2754        let t = tokens.len();
2755        let eps = self.cfg.rms_eps;
2756        let pos: Vec<i32> = (cache.pos as i32..(cache.pos + t) as i32).collect();
2757        let pos_d = e.htod_i32(&pos)?;
2758        // glm5 DFlash2 feature tap: chunked primes write their rows at the chunk's absolute
2759        // offset (cache.pos advances per chunk — the dflash_taps.base precedent).
2760        if let Some(sink) = cache.hc_taps.as_mut() {
2761            sink.base = cache.pos;
2762        }
2763
2764        let mut embedded = self.embed(e, tokens)?;
2765        if let Some(ov) = overlay {
2766            // Mixed-embedding splice (shared with the ppN twin — EmbedOverlay::splice_into):
2767            // image rows overwrite placeholder-token embeddings inside this chunk's
2768            // prompt-relative window [chunk_off, chunk_off+t), BEFORE stream expansion —
2769            // the reference's splice point (execute_multimodal replaces rows before
2770            // hc_expand).
2771            ov.splice_into(e, &mut embedded, chunk_off, t, n_embd)?;
2772        }
2773        let mut x = crate::hyper::expand(e, &topology, &embedded, t, n_embd)?;
2774
2775        for (il, layer) in self.layers.iter().enumerate() {
2776            let hyper = layer.hyper.as_ref().ok_or_else(|| {
2777                format!("layer {il} carries no hyper-connection weights under an hc plan")
2778            })?;
2779
2780            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.attn, &x, t, n_embd)?;
2781            let mut h = e.uninit(t * n_embd)?;
2782            e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2783            let mixed = match &layer.mixer {
2784                Mixer::Full(fa) => {
2785                    self.full_attn_prime(e, fa, &h, None, &pos_d, t, cache, il, seq_end)?
2786                }
2787                Mixer::Linear(la) => self.linear_attn_prime(e, la, &h, None, t, cache, il)?,
2788                Mixer::Mla(mla) if mla.tp.is_some() => {
2789                    self.mla_tp_attn_cached(e, mla, &h, &pos_d, t, il, cache, false)?
2790                }
2791                Mixer::Mla(mla) => self.mla_attn_cached(e, mla, &h, &pos_d, t, il, cache)?,
2792                Mixer::Kda(la) if la.tp.is_some() => crate::glm5_tp::kda_tp_cached(
2793                    e,
2794                    la,
2795                    &h,
2796                    t,
2797                    eps,
2798                    cache,
2799                    il,
2800                    crate::kda::ConvArm::Prefill,
2801                )?,
2802                Mixer::Kda(la) => crate::kda::kda_prime_cached(e, la, &h, t, eps, cache, il)?,
2803            };
2804            x = crate::hyper::post(e, &topology, &mixed, &x, &mix, t, n_embd)?;
2805
2806            let (y, mix) = crate::hyper::pre(e, &topology, &hyper.mlp, &x, t, n_embd)?;
2807            let mut z = e.uninit(t * n_embd)?;
2808            e.rms_norm(
2809                &y,
2810                layer.post_attn_norm.float_data(),
2811                &mut z,
2812                n_embd,
2813                t,
2814                eps,
2815            )?;
2816            let ffn_out = self.hyper_ffn_branch(e, layer, &z, t, il, true)?;
2817            x = crate::hyper::post(e, &topology, &ffn_out, &x, &mix, t, n_embd)?;
2818            // glm5 DFlash2 feature tap (see hyper_range_prime — the unsplit chunk walk
2819            // taps the same completed-layer-output contraction).
2820            self.glm5_hc_tap(e, cache, &topology, il, &x, t)?;
2821        }
2822
2823        let hiddens =
2824            crate::hyper::collapse(e, &topology, self.hyper_head.as_ref(), &x, t, n_embd)?;
2825        let mut hn = e.uninit(t * n_embd)?;
2826        e.rms_norm(
2827            &hiddens,
2828            self.output_norm.float_data(),
2829            &mut hn,
2830            n_embd,
2831            t,
2832            eps,
2833        )?;
2834        let last = e.view(&hn, t * n_embd);
2835        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2836        let mut hlast = e.uninit(n_embd)?;
2837        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2838        let logits = e.matmul(&self.output, &hlast, 1)?;
2839        let host = e.dtoh(&logits)?;
2840
2841        // h_seed is the PRE-output_norm hidden of the last row (MTP-PLAN §A), taken from the
2842        // collapsed stack so it means the same thing it does on the serial path.
2843        let stack = e.view(&hiddens, t * n_embd);
2844        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2845        let mut h_seed = e.uninit(n_embd)?;
2846        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2847        cache.pos += t;
2848        Ok((host, h_seed, hiddens))
2849    }
2850
2851    /// Prime exit shared by `prime_cache_hyper` and its ppN twin: collapse, output_norm, last
2852    /// row logits, and the pre-output_norm hidden seed taken from the collapsed stack (MTP-PLAN
2853    /// §A) so it means the same thing it does on the serial path.
2854    #[allow(clippy::too_many_arguments)]
2855    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2856    fn hyper_prime_tail(
2857        &self,
2858        e: &Engine,
2859        topology: &crate::hyper::HyperTopology,
2860        x: &CudaSlice<f32>,
2861        t: usize,
2862        n_embd: usize,
2863        eps: f32,
2864        cache: &mut Cache,
2865    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2866        let hiddens = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, t, n_embd)?;
2867        let mut hn = e.uninit(t * n_embd)?;
2868        e.rms_norm(
2869            &hiddens,
2870            self.output_norm.float_data(),
2871            &mut hn,
2872            n_embd,
2873            t,
2874            eps,
2875        )?;
2876        let last = e.view(&hn, t * n_embd);
2877        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
2878        let mut hlast = e.uninit(n_embd)?;
2879        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
2880        let logits = e.matmul(&self.output, &hlast, 1)?;
2881        let host = e.dtoh(&logits)?;
2882        let stack = e.view(&hiddens, t * n_embd);
2883        let seed_row = stack.slice((t - 1) * n_embd..t * n_embd);
2884        let mut h_seed = e.uninit(n_embd)?;
2885        e.copy_view_into(&mut h_seed, 0, &seed_row, n_embd)?;
2886        cache.pos += t;
2887        Ok((host, h_seed, hiddens))
2888    }
2889
2890    /// ppN twin of `decode_step_hyper`: the T=1 step as N stage subgraphs, each on its own
2891    /// stream (and, under `MEMRA_PP_DEVICES`, its own device/engine), with the
2892    /// transport-selected boundary handoff of the `[streams, n_embd]` state at each fence cut.
2893    /// `cache.pos` is snapshotted once and advanced once.
2894    fn decode_step_hyper_ppn(
2895        &self,
2896        e: &Engine,
2897        token: u32,
2898        cache: &mut Cache,
2899        topology: &crate::hyper::HyperTopology,
2900        fence: &[usize],
2901    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2902        let n_embd = self.cfg.n_embd as usize;
2903        let eps = self.cfg.rms_eps;
2904        let pos = cache.pos;
2905        let width = topology.streams * n_embd;
2906
2907        if crate::pp::pp2_streams_off() {
2908            let pos_d = e.htod_i32(&[pos as i32])?;
2909            let embedded = e.htod(&self.embd.gather(n_embd, &[token]))?;
2910            let mut x = crate::hyper::expand(e, topology, &embedded, 1, n_embd)?;
2911            x = self.hyper_range_decode(e, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
2912            for s in 1..fence.len() - 1 {
2913                let boundary_tx = e.clone_dtod(&x)?;
2914                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2915                x = self.hyper_range_decode(
2916                    e,
2917                    topology,
2918                    boundary_rx,
2919                    fence[s],
2920                    fence[s + 1],
2921                    &pos_d,
2922                    pos,
2923                    cache,
2924                )?;
2925            }
2926            return self.hyper_decode_tail(e, topology, &x, n_embd, eps, cache);
2927        }
2928
2929        let rt = crate::pp::PpNRt::get(e)?;
2930        let n_st = fence.len() - 1;
2931        assert_eq!(
2932            rt.n_stages(),
2933            n_st,
2934            "PpNRt stage count {} != fence stages {n_st}",
2935            rt.n_stages()
2936        );
2937        rt.fence_stages_behind(&e.stream())?;
2938
2939        let mut slot = {
2940            let _st0 = rt.enter(0);
2941            let e0 = rt.engine(0, e);
2942            let pos_d = e0.htod_i32(&[pos as i32])?;
2943            let embedded = e0.htod(&self.embd.gather(n_embd, &[token]))?;
2944            let x = crate::hyper::expand(e0, topology, &embedded, 1, n_embd)?;
2945            let x =
2946                self.hyper_range_decode(e0, topology, x, fence[0], fence[1], &pos_d, pos, cache)?;
2947            rt.tx(0, &x, width)?
2948        };
2949        for s in 1..n_st - 1 {
2950            let _st = rt.enter(s);
2951            let es = rt.engine(s, e);
2952            let pos_d = es.htod_i32(&[pos as i32])?;
2953            let x = rt.rx(s - 1, slot, width)?;
2954            let x = self.hyper_range_decode(
2955                es,
2956                topology,
2957                x,
2958                fence[s],
2959                fence[s + 1],
2960                &pos_d,
2961                pos,
2962                cache,
2963            )?;
2964            slot = rt.tx(s, &x, width)?;
2965        }
2966        let _stl = rt.enter(n_st - 1);
2967        let el = rt.engine(n_st - 1, e);
2968        let pos_d = el.htod_i32(&[pos as i32])?;
2969        let x = rt.rx(n_st - 2, slot, width)?;
2970        let x = self.hyper_range_decode(
2971            el,
2972            topology,
2973            x,
2974            fence[n_st - 1],
2975            fence[n_st],
2976            &pos_d,
2977            pos,
2978            cache,
2979        )?;
2980        self.hyper_decode_tail(el, topology, &x, n_embd, eps, cache)
2981    }
2982
2983    /// Decode exit shared by `decode_step_hyper` and its ppN twin. `h_seed` is the COLLAPSED
2984    /// hidden (not the pre-collapse stream state and not the post-norm row): that is what the
2985    /// serial hc step publishes, and the two must not drift.
2986    fn hyper_decode_tail(
2987        &self,
2988        e: &Engine,
2989        topology: &crate::hyper::HyperTopology,
2990        x: &CudaSlice<f32>,
2991        n_embd: usize,
2992        eps: f32,
2993        cache: &mut Cache,
2994    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2995        let h_seed = crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, 1, n_embd)?;
2996        let mut hn = e.uninit(n_embd)?;
2997        e.rms_norm(
2998            &h_seed,
2999            self.output_norm.float_data(),
3000            &mut hn,
3001            n_embd,
3002            1,
3003            eps,
3004        )?;
3005        let logits = e.matmul(&self.output, &hn, 1)?;
3006        let host = e.dtoh(&logits)?;
3007        cache.pos += 1;
3008        Ok((host, h_seed))
3009    }
3010
3011    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
3012    pub fn forward(
3013        &self,
3014        e: &Engine,
3015        tokens: &[u32],
3016    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3017        if self.hyper.is_some() {
3018            return self.forward_hyper(e, tokens, false);
3019        }
3020        if self.is_gemma4_e4b() {
3021            return self.gemma4_e4b_forward(e, tokens, false);
3022        }
3023        if self.uses_gemma_program() {
3024            return self.gemma4_forward(e, tokens, false);
3025        }
3026        let cfg = &self.cfg;
3027        let n_embd = cfg.n_embd as usize;
3028        let t = tokens.len();
3029        let eps = cfg.rms_eps;
3030        let pos: Vec<i32> = (0..t as i32).collect();
3031        let pos_d = e.htod_i32(&pos)?;
3032
3033        let mut x = self.embed(e, tokens)?; // [T, n_embd]
3034
3035        for (il, layer) in self.layers.iter().enumerate() {
3036            // attn_norm
3037            let mut h = e.uninit(t * n_embd)?;
3038            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3039
3040            let mixed = match &layer.mixer {
3041                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
3042                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
3043                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3044                Mixer::Kda(la) => crate::kda::kda_attn(e, la, &h, t, eps)?,
3045            };
3046
3047            // residual 1
3048            let mut x1 = e.uninit(t * n_embd)?;
3049            e.add(&x, &mixed, &mut x1, t * n_embd)?;
3050
3051            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
3052            let mut z = e.uninit(t * n_embd)?;
3053            e.rms_norm(
3054                &x1,
3055                layer.post_attn_norm.float_data(),
3056                &mut z,
3057                n_embd,
3058                t,
3059                eps,
3060            )?;
3061            let ffn_out = match &layer.ffn {
3062                crate::hybrid::Ffn::Dense {
3063                    ffn_gate,
3064                    ffn_up,
3065                    ffn_down,
3066                } => {
3067                    let n_ff = ffn_gate.out_features();
3068                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3069                    let up = g2.pop().unwrap();
3070                    let gate = g2.pop().unwrap();
3071                    let mut act = e.uninit(t * n_ff)?;
3072                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
3073                    // both the dense MLP and the shared expert, and its limit is
3074                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
3075                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
3076                    Self::ffn_act_lim(
3077                        e,
3078                        &self.cfg,
3079                        &gate,
3080                        &up,
3081                        1.0,
3082                        1.0,
3083                        self.cfg.clamp_shexp_at(il as u32),
3084                        &mut act,
3085                        t * n_ff,
3086                    )?;
3087                    e.matmul(ffn_down, &act, t)?
3088                }
3089                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
3090            };
3091            let mut x2 = e.uninit(t * n_embd)?;
3092            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3093            x = x2;
3094        }
3095
3096        let mut hn = e.uninit(t * n_embd)?;
3097        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3098        let logits = e.matmul(&self.output, &hn, t)?;
3099        e.dtoh(&logits)
3100    }
3101
3102    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
3103    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
3104    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
3105    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
3106    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
3107    pub fn forward_last(
3108        &self,
3109        e: &Engine,
3110        tokens: &[u32],
3111    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3112        if self.hyper.is_some() {
3113            return self.forward_hyper(e, tokens, true);
3114        }
3115        if self.uses_gemma_program() {
3116            return self.gemma4_forward(e, tokens, true);
3117        }
3118        let cfg = &self.cfg;
3119        let n_embd = cfg.n_embd as usize;
3120        let t = tokens.len();
3121        let eps = cfg.rms_eps;
3122        let pos: Vec<i32> = (0..t as i32).collect();
3123        let pos_d = e.htod_i32(&pos)?;
3124
3125        let mut x = self.embed(e, tokens)?; // [T, n_embd]
3126        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
3127        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
3128        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
3129        let anat = Self::prime_anatomy_on();
3130        let mut anat_last = if anat {
3131            e.stream().synchronize()?;
3132            Some(std::time::Instant::now())
3133        } else {
3134            None
3135        };
3136        macro_rules! anat_mark {
3137            ($slot:expr) => {
3138                if let Some(ts) = anat_last.as_mut() {
3139                    e.stream().synchronize()?;
3140                    Self::prime_anatomy_slots()[$slot].fetch_add(
3141                        ts.elapsed().as_nanos() as u64,
3142                        std::sync::atomic::Ordering::Relaxed,
3143                    );
3144                    *ts = std::time::Instant::now();
3145                }
3146            };
3147        }
3148        for (il, layer) in self.layers.iter().enumerate() {
3149            let mut h = e.uninit(t * n_embd)?;
3150            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
3151            if probe {
3152                e.stream().synchronize()?;
3153                eprintln!("[probe] L{il} norm ok");
3154            }
3155            anat_mark!(4);
3156            let mixed = match &layer.mixer {
3157                Mixer::Full(fa) => {
3158                    let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
3159                    anat_mark!(0);
3160                    y
3161                }
3162                Mixer::Linear(la) => {
3163                    let y = self.linear_attn(e, la, &h, t)?;
3164                    anat_mark!(1);
3165                    y
3166                }
3167                Mixer::Mla(mla) => self.mla_attn(e, mla, &h, &pos_d, t, il)?,
3168                Mixer::Kda(la) => {
3169                    let y = crate::kda::kda_attn(e, la, &h, t, eps)?;
3170                    // KDA shares the linear-mixer anatomy slot: same mixer class, one bucket.
3171                    anat_mark!(1);
3172                    y
3173                }
3174            };
3175            if probe {
3176                e.stream().synchronize()?;
3177                eprintln!("[probe] L{il} mixer ok");
3178            }
3179            let mut x1 = e.uninit(t * n_embd)?;
3180            e.add(&x, &mixed, &mut x1, t * n_embd)?;
3181            let mut z = e.uninit(t * n_embd)?;
3182            e.rms_norm(
3183                &x1,
3184                layer.post_attn_norm.float_data(),
3185                &mut z,
3186                n_embd,
3187                t,
3188                eps,
3189            )?;
3190            anat_mark!(4);
3191            let ffn_out = match &layer.ffn {
3192                crate::hybrid::Ffn::Dense {
3193                    ffn_gate,
3194                    ffn_up,
3195                    ffn_down,
3196                } => {
3197                    let n_ff = ffn_gate.out_features();
3198                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
3199                    let up = g2.pop().unwrap();
3200                    let gate = g2.pop().unwrap();
3201                    let mut act = e.uninit(t * n_ff)?;
3202                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
3203                    Self::ffn_act_lim(
3204                        e,
3205                        &self.cfg,
3206                        &gate,
3207                        &up,
3208                        1.0,
3209                        1.0,
3210                        self.cfg.clamp_shexp_at(il as u32),
3211                        &mut act,
3212                        t * n_ff,
3213                    )?;
3214                    let y = e.matmul(ffn_down, &act, t)?;
3215                    anat_mark!(3);
3216                    y
3217                }
3218                crate::hybrid::Ffn::Moe(m) => {
3219                    let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
3220                    anat_mark!(2);
3221                    y
3222                }
3223            };
3224            if probe {
3225                e.stream().synchronize()?;
3226                eprintln!("[probe] L{il} ffn ok");
3227            }
3228            let mut x2 = e.uninit(t * n_embd)?;
3229            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
3230            x = x2;
3231        }
3232        if anat {
3233            let s = Self::prime_anatomy_slots();
3234            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
3235            eprintln!(
3236                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
3237                 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
3238                ms(0),
3239                ms(1),
3240                ms(2),
3241                ms(3),
3242                ms(4)
3243            );
3244        }
3245        // norm over all T, then slice the LAST row and run lm_head on that single row.
3246        let mut hn = e.uninit(t * n_embd)?;
3247        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3248        let last = e.view(&hn, t * n_embd); // [T, n_embd]
3249        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
3250        let mut hlast = e.uninit(n_embd)?;
3251        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
3252        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
3253        e.dtoh(&logits)
3254    }
3255
3256    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
3257    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
3258    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
3259    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
3260    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
3261    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
3262    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
3263    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
3264    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
3265    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
3266    ///       argmax gate is the accuracy authority, exactly as for forward_last);
3267    ///   (c) `cache.pos`/KV len/len_d advance by T.
3268    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
3269    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
3270    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
3271    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
3272    ///
3273    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
3274    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
3275    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
3276    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
3277    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
3278    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
3279    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
3280    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
3281    /// differently under load — research/tick-seg-20260807, receipt in
3282    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
3283    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
3284    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
3285    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
3286    /// caller that SPLITS one request across calls passes the remainder.
3287    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3288    pub fn prime_cache(
3289        &self,
3290        e: &Engine,
3291        tokens: &[u32],
3292        cache: &mut Cache,
3293        queued_after: usize,
3294    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3295        self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
3296    }
3297
3298    /// `prime_cache` with a vision embedding overlay (lane/vision): image merger outputs
3299    /// replace the `<|image_pad|>` token embeddings at prompt-relative positions before the
3300    /// trunk walk — the mixed-embedding prime. Text-only callers use `prime_cache` (overlay
3301    /// None, byte-identical path). Scope: the serial chunk walk, the single-engine hyper walk,
3302    /// and the hyper ppN twin (splice at stage-0 embedding intake; lane/glm5-vision-default-on,
3303    /// gated by glm5-hyper-ppn-gate's overlay arm). The serial PP-2 pipelined prime and
3304    /// gemma4 E4B still refuse loudly.
3305    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3306    pub fn prime_cache_overlaid(
3307        &self,
3308        e: &Engine,
3309        tokens: &[u32],
3310        cache: &mut Cache,
3311        queued_after: usize,
3312        overlay: Option<&crate::vision::EmbedOverlay>,
3313    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3314        cache.ensure_usable("prime_cache")?;
3315        if self.hyper.is_some() {
3316            // The mixed-embedding splice lands BEFORE stream expansion (the same point
3317            // the reference's execute_multimodal replaces rows — before hc_expand), so
3318            // the hyper walk needs no overlay-specific arithmetic (lane/glm5-vision).
3319            return self.prime_cache_hyper(e, tokens, cache, queued_after, overlay);
3320        }
3321        let _pp_walk =
3322            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
3323                let rt = crate::pp::PpNRt::get(e)?;
3324                Some(rt.acquire_walk("prime_cache")?)
3325            } else {
3326                None
3327            };
3328        let n_embd = self.cfg.n_embd as usize;
3329        let t = tokens.len();
3330        // MEMRA_PRIME_TROWS=1: prefill through the same-session t-row walk (per-row t=1
3331        // program = the tokenwise-prime ORACLE class) — replaces the host-canonical
3332        // per-token step-TP prime. Text-only fresh primes; anything else falls through.
3333        // MEMRA_STEP_GEMM_PRIME: prime the prompt through the batched GEMM path, CHUNKED.
3334        // The batch entry supplies both halves of the fast prime — the GEMM trunk at m = chunk
3335        // and the grouped NVFP4 MoE — which is why routing only the MoE through the ordinary
3336        // chunk loop measured 26.9 s against 3.7 s here. Chunking keeps the transients bounded:
3337        // a whole 32k prompt in one call would build a 262144-pair CSR and ~4.3 GB of partials
3338        // per rank, the blow-up the chunked prime exists to prevent.
3339        //
3340        // CONTINUATION (lane/gemm-suffix, 2026-08-28): the entry NO LONGER requires
3341        // `cache.pos == 0`. The batch core has been continuation-capable since 7700e0b6
3342        // (positions carry each sequence's base; the fresh-prompt guard narrowed to B > 1),
3343        // and d99b2ea3 named this outer guard as the remaining blocker in its own message.
3344        // Every multi-turn suffix and every tick remainder was paying the walk's measured
3345        // ~7.2 ms/token against this path's ~1.0 ms/token, which is why session-affinity
3346        // reuse measured a 1.012x wash on a growing conversation.
3347        // ONE DEFECT HAD TO BE FIXED FIRST, and it was LIVE before this lift:
3348        // `step35_prime_batch_layers` passed `ts[s]` — the CHUNK's length — as `seq_end`.
3349        // `seq_end` is the REQUEST's absolute end position and it steers step35's SWA arm
3350        // (`seq_end > win`, win = 512 on step37). A chunk SHORTER than the window at a
3351        // NONZERO base therefore selected the UNWINDOWED FA arm over a view that the `off`
3352        // trim leaves at ~win-1+t rows: it attended OUTSIDE the sliding window. That was
3353        // already reachable with no continuation at all — a fresh prompt of 4096+k for k in
3354        // [PRIME_MIN_T, 512) ends in a trailing chunk of exactly that shape. `seq_end` is now
3355        // threaded from here (request-absolute, `+ queued_after`, computed ONCE before the
3356        // chunk loop, chunk-size-invariant exactly as on the walk), which is what makes the
3357        // suffix arm expressible at all rather than merely reachable.
3358        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
3359        let seq_end = if legacy_calllocal {
3360            cache.pos + t
3361        } else {
3362            cache.pos + t + queued_after
3363        };
3364        // MEMRA_STEP_GEMM_PRIME_SUFFIX is the SUFFIX-ONLY seam: off (its default in this
3365        // commit) leaves continuations on the walk while fresh primes stay on the fast path;
3366        // MEMRA_STEP_GEMM_PRIME=0 is the whole-path seam. The `seq_end` threading above is
3367        // deliberately NOT behind either door — it is a correctness fix for the fresh path too.
3368        if overlay.is_none()
3369            && (cache.pos == 0 || step_gemm_prime_suffix_on())
3370            && t >= PRIME_MIN_T
3371            && crate::step_gemm_prime_on()
3372            && self.uses_sliding_gated_moe_program()
3373        {
3374            let n_embd = self.cfg.n_embd as usize;
3375            let base = cache.pos;
3376            let width = crate::cache::PRIME_CHUNK_MAX_TOKENS;
3377            let mut hiddens = e.uninit(t * n_embd)?;
3378            let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3379            let mut start = 0usize;
3380            while start < t {
3381                // A trailing chunk below the walk floor folds into the previous one; every chunk
3382                // this entry sees must clear PRIME_MIN_T on its own.
3383                let mut end = (start + width).min(t);
3384                if t - end > 0 && t - end < PRIME_MIN_T {
3385                    end = t;
3386                }
3387                let mut out = self.step35_prime_cache_batch(
3388                    e,
3389                    &[&tokens[start..end]],
3390                    &mut [cache],
3391                    &[seq_end],
3392                )?;
3393                if out.len() != 1 {
3394                    return Err("B=1 batched prime returned a non-singleton".into());
3395                }
3396                let (logits, h_seed, hidden) = out.remove(0);
3397                e.copy_into(
3398                    &mut hiddens,
3399                    start * n_embd,
3400                    &hidden,
3401                    (end - start) * n_embd,
3402                )?;
3403                last = Some((logits, h_seed));
3404                start = end;
3405            }
3406            let (logits, h_seed) = last.expect("prime produced no chunk");
3407            // ENGAGEMENT RECEIPT, both directions. `base` is the discriminator: base=0 is a
3408            // fresh prime (this line existed before the lift), base>0 is a SUFFIX riding the
3409            // GEMM trunk — the arm this lane added. The declining twin below counts the other
3410            // direction, so a log that shows neither line is an instrument fault, not a pass.
3411            eprintln!(
3412                "[gemm-prime] ENGAGED t={t} base={base} seq_end={seq_end} chunks<={width} (GEMM trunk + grouped MoE)"
3413            );
3414            return Ok((logits, h_seed, hiddens));
3415        }
3416        if self.uses_sliding_gated_moe_program() {
3417            eprintln!(
3418                "[gemm-prime] WALK t={t} base={} seq_end={seq_end} (batched prime declined)",
3419                cache.pos
3420            );
3421        }
3422        if overlay.is_none()
3423            && let Some(out) = self.step35_prime_trows(e, tokens, cache)?
3424        {
3425            return Ok(out);
3426        }
3427        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
3428        // session cache — every chunk (including the first) takes the continuation arm
3429        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
3430        assert!(
3431            t >= PRIME_MIN_T,
3432            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
3433        );
3434        assert!(
3435            cache.pos + t <= cache.max_ctx,
3436            "prime_cache: prompt exceeds cache max_ctx"
3437        );
3438
3439        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
3440        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
3441        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
3442        // each chunk runs the full layer stack with transients sized to the chunk, appending its
3443        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
3444        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
3445        // exactly the state carry it was built for). Full-attn chunks after the first attend to
3446        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
3447        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
3448        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
3449        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
3450        if self.is_gemma4_e4b() || self.uses_gemma_program() {
3451            if self.is_gemma4_e4b() {
3452                if overlay.is_some() {
3453                    return Err(
3454                        "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
3455                    );
3456                }
3457                return self.gemma4_e4b_prime(e, tokens, cache);
3458            }
3459            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
3460            // An overlay takes the masked-prefill arm: image rows splice in unscaled
3461            // (gemma4.cpp:182 — embd batches skip the sqrt(n_embd) scale) and the image
3462            // spans become bidirectional attention islands (lane/gemma-vision).
3463            return self.gemma4_prime(e, tokens, cache, overlay);
3464        }
3465        if crate::pp::prime_pipe_on()
3466            && crate::pp::prime_pp_on()
3467            && !crate::pp::pp2_streams_off()
3468            && crate::pp::pp_cuts(self.layers.len())
3469                .is_some_and(|fence| matches!(fence.len(), 4 | 5))
3470        {
3471            crate::pp::pp_wave_on()
3472                .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3473        }
3474        let ranges = prime_chunk_ranges(t, self.layers.len(), self.gdn_prime_grid_on());
3475        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
3476        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
3477        // the prefill's ARITHMETIC, so two rigs with different values produced different
3478        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
3479        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
3480        // (VERDICT.md) — and it is NOT what docs originally said:
3481        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
3482        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
3483        //     output head), so growing a chunk cannot move an existing row's value.
3484        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
3485        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
3486        //     not describe our leak.
3487        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
3488        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
3489        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
3490        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
3491        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
3492        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
3493        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
3494        // the source — every row is in one numeric class, so the chunk size no longer steers
3495        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
3496        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
3497        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
3498        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
3499        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
3500        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
3501        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
3502        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
3503        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
3504        // across calls, the request still ends at the same absolute position, whatever the tick
3505        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
3506        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
3507        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
3508        // default. Read per call, not cached (the probe flips it in-process between arms). Never
3509        // on in a measured default run.
3510        // `seq_end` (and its MEMRA_PRIME_CALLLOCAL seam) is computed ONCE above the batched
3511        // entry so both prime arms read the identical request-absolute value.
3512        if ranges.len() == 1 {
3513            return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
3514        }
3515        // PIPELINED PP PRIME. PP-2 retains its independently-qualified two-stage schedule;
3516        // PP-3/4 require the explicit MEMRA_PP_WAVE=1 persistent-stage wavefront. The serial
3517        // split stays reachable through MEMRA_PRIME_PIPE=0 and is the exactness oracle.
3518        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
3519            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
3520                if overlay.is_some() {
3521                    return Err(
3522                        "vision embedding overlay + pipelined PP prime unsupported (v1); \
3523                         run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
3524                            .into(),
3525                    );
3526                }
3527                if crate::pp::pp_multi_stream_same_device() {
3528                    return Err(
3529                        "prime chunk pipeline refused with 2 stage streams on one device — \
3530                         that concurrent-stream placement remains quarantined by the deferred \
3531                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
3532                         the serial split."
3533                            .into(),
3534                    );
3535                }
3536                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3537            }
3538            if let Some(fence) =
3539                crate::pp::pp_cuts(self.layers.len()).filter(|f| matches!(f.len(), 4 | 5))
3540            {
3541                let wave_on = crate::pp::pp_wave_on()
3542                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3543                let stages = fence.len() - 1;
3544                if crate::pp::pp_wave_route_enabled(
3545                    wave_on,
3546                    crate::pp::pp2_overlap(),
3547                    stages,
3548                    ranges.len(),
3549                ) {
3550                    if overlay.is_some() {
3551                        return Err(
3552                            "vision embedding overlay + pipelined PP prime unsupported; \
3553                             run the serial prime (MEMRA_PP_WAVE=0 or MEMRA_PRIME_PIPE=0)"
3554                                .into(),
3555                        );
3556                    }
3557                    let rt = crate::pp::PpNRt::get(e)?;
3558                    let double_slot = crate::pp::pp2_overlap();
3559                    crate::pp::pp_wave_eligibility(
3560                        stages,
3561                        double_slot,
3562                        rt.host_bounce_active(),
3563                        rt.repeated_stage_device(),
3564                    )
3565                    .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
3566                    return self
3567                        .prime_cache_ppn_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
3568                }
3569            }
3570        }
3571        let mut hiddens = e.uninit(t * n_embd)?;
3572        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3573        for &(start, end) in &ranges {
3574            // chunked prime writes tap rows at the chunk's absolute offset
3575            if let Some(taps) = cache.dflash_taps.as_mut() {
3576                taps.base = start;
3577            }
3578            let (l, hs, x) =
3579                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
3580            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
3581            last = Some((l, hs));
3582        }
3583        let (logits, h_seed) = last.unwrap();
3584        Ok((logits, h_seed, hiddens))
3585    }
3586
3587    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
3588    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
3589    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
3590    /// norm, lm head, and caller hidden-stack copy as the serial split.
3591    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3592    fn prime_cache_pp2_pipelined(
3593        &self,
3594        e: &Engine,
3595        tokens: &[u32],
3596        cache: &mut Cache,
3597        seq_end: usize,
3598        ranges: &[(usize, usize)],
3599        fence: &[usize],
3600    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3601        debug_assert_eq!(fence.len(), 3);
3602        debug_assert!(ranges.len() >= 2);
3603        let rt = crate::pp::PpNRt::get(e)?;
3604        assert_eq!(
3605            rt.n_stages(),
3606            2,
3607            "prime pipeline requires exactly two PP stages"
3608        );
3609        let n_embd = self.cfg.n_embd as usize;
3610        let t = tokens.len();
3611        let initial_base = cache.pos;
3612        let caller_stream = e.stream();
3613
3614        // #87 reverse publication before any new stage allocation, then prewarm both
3615        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
3616        // after stage 1(N) is queued would synchronize that stream and erase the first
3617        // overlap on a two-chunk prompt.
3618        rt.fence_stages_behind(&caller_stream)?;
3619        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
3620        rt.prepare_overlap_slots(0, max_payload)?;
3621
3622        let mut hiddens = e.uninit(t * n_embd)?;
3623        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3624        let mut stage_caches = PrimeCacheStages::new(cache, fence);
3625        let (cache0, cache1) = stage_caches.pp2_parts();
3626        let (first_start, first_end) = ranges[0];
3627        let mut slot = self.prime_pp2_stage0_enqueue(
3628            e,
3629            rt,
3630            &tokens[first_start..first_end],
3631            cache0,
3632            seq_end,
3633            fence,
3634            initial_base + first_start,
3635            true,
3636        )?;
3637        cache0.pos = initial_base + first_end;
3638
3639        for (i, &(start, end)) in ranges.iter().enumerate() {
3640            let base = initial_base + start;
3641            debug_assert_eq!(
3642                cache1.pos, base,
3643                "stage 1 must drain chunks in original position order"
3644            );
3645            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
3646                let next_base = initial_base + next_start;
3647                debug_assert_eq!(
3648                    cache0.pos, next_base,
3649                    "stage 0 must issue chunks in original position order"
3650                );
3651                let cache0_stage = &mut *cache0;
3652                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
3653                // on one host thread therefore serialize even if the calls are ordered as
3654                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
3655                // stage 1 consumes slot N while stage 0 produces slot N+1.
3656                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
3657                    let stage0 = scope.spawn(move || -> Result<usize, String> {
3658                        let next = self
3659                            .prime_pp2_stage0_enqueue(
3660                                e,
3661                                rt,
3662                                &tokens[next_start..next_end],
3663                                cache0_stage,
3664                                seq_end,
3665                                fence,
3666                                next_base,
3667                                true,
3668                            )
3669                            .map_err(|err| err.to_string())?;
3670                        cache0_stage.pos = initial_base + next_end;
3671                        Ok(next)
3672                    });
3673                    let x = self.prime_pp2_stage1_enqueue(
3674                        e,
3675                        rt,
3676                        slot,
3677                        end - start,
3678                        cache1,
3679                        seq_end,
3680                        fence,
3681                        base,
3682                        true,
3683                    )?;
3684                    let out = {
3685                        rt.bind_stage(1)?;
3686                        let _st1 = rt.enter(1);
3687                        let e1 = rt.engine(1, e);
3688                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3689                    };
3690                    let next = match stage0.join() {
3691                        Ok(result) => {
3692                            result.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?
3693                        }
3694                        Err(payload) => std::panic::resume_unwind(payload),
3695                    };
3696                    Ok((out, Some(next)))
3697                })?
3698            } else {
3699                let x = self.prime_pp2_stage1_enqueue(
3700                    e,
3701                    rt,
3702                    slot,
3703                    end - start,
3704                    cache1,
3705                    seq_end,
3706                    fence,
3707                    base,
3708                    true,
3709                )?;
3710                let out = {
3711                    rt.bind_stage(1)?;
3712                    let _st1 = rt.enter(1);
3713                    let e1 = rt.engine(1, e);
3714                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
3715                };
3716                (out, None)
3717            };
3718
3719            rt.publish_to(1, &caller_stream)?;
3720            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
3721            last = Some((out.0, out.1));
3722            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3723
3724            if let Some(next) = next_slot {
3725                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
3726                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
3727                // Stage 0(N+1) is already queued before this wait is appended, so its
3728                // overlap with stage 1(N) is preserved.
3729                rt.fence_stages_behind(&caller_stream)?;
3730                slot = next;
3731            }
3732        }
3733
3734        debug_assert_eq!(cache0.pos, initial_base + t);
3735        debug_assert_eq!(cache1.pos, initial_base + t);
3736        let (logits, h_seed) = last.unwrap();
3737        stage_caches.commit();
3738        Ok((logits, h_seed, hiddens))
3739    }
3740
3741    /// PP-3/4 prime wavefront: one prompt microchunk is one wave. One scoped host worker owns each
3742    /// non-head stage for the whole walk; the caller thread owns the head stage. Forward boundary
3743    /// messages preserve wave order, while reverse exact-wave acknowledgements are sent only after
3744    /// downstream `rx` has recorded `ev_rx`. An upstream stage therefore cannot cycle back to either
3745    /// shared slot before that slot's current generation has a host-observed release point.
3746    ///
3747    /// PP-2 remains on its independently qualified scheduler above. This path is reachable only
3748    /// through MEMRA_PP_WAVE=1 and the topology gate.
3749    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
3750    fn prime_cache_ppn_pipelined(
3751        &self,
3752        e: &Engine,
3753        tokens: &[u32],
3754        cache: &mut Cache,
3755        seq_end: usize,
3756        ranges: &[(usize, usize)],
3757        fence: &[usize],
3758    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3759        let stages = fence.len().saturating_sub(1);
3760        debug_assert!((3..=4).contains(&stages));
3761        debug_assert!(ranges.len() >= 2);
3762        let rt = crate::pp::PpNRt::get(e)?;
3763        assert_eq!(
3764            rt.n_stages(),
3765            stages,
3766            "prime wavefront PpNRt/fence stage mismatch"
3767        );
3768        let n_embd = self.cfg.n_embd as usize;
3769        let initial_base = cache.pos;
3770        let caller_stream = e.stream();
3771        let primary_context = crate::pp::PrimaryContextRestore::new(e);
3772
3773        rt.fence_stages_behind(&caller_stream)?;
3774        let max_payload = ranges
3775            .iter()
3776            .map(|(start, end)| (end - start) * n_embd)
3777            .max()
3778            .unwrap_or(0);
3779        for boundary in 0..stages - 1 {
3780            rt.prepare_overlap_slots(boundary, max_payload)?;
3781        }
3782
3783        let mut stage_caches = PrimeCacheStages::new(cache, fence);
3784        let waves: Vec<_> = ranges
3785            .iter()
3786            .map(|&(start, end)| PrimePpWave {
3787                start,
3788                end,
3789                tokens: &tokens[start..end],
3790            })
3791            .collect();
3792        let mut forward_senders = Vec::with_capacity(stages - 1);
3793        let mut forward_receivers = Vec::with_capacity(stages - 1);
3794        let mut release_senders = Vec::with_capacity(stages - 1);
3795        let mut release_receivers = Vec::with_capacity(stages - 1);
3796        for _ in 0..stages - 1 {
3797            let (forward_sender, forward_receiver) = std::sync::mpsc::channel();
3798            let (release_sender, release_receiver) = std::sync::mpsc::channel();
3799            forward_senders.push(Some(forward_sender));
3800            forward_receivers.push(Some(forward_receiver));
3801            release_senders.push(Some(release_sender));
3802            release_receivers.push(Some(release_receiver));
3803        }
3804        let mut stage_channels = Vec::with_capacity(stages - 1);
3805        for stage in 0..stages - 1 {
3806            stage_channels.push(Some(PrimePpStageChannels {
3807                incoming: (stage > 0).then(|| forward_receivers[stage - 1].take().unwrap()),
3808                release_upstream: (stage > 0).then(|| release_senders[stage - 1].take().unwrap()),
3809                outgoing: forward_senders[stage].take().unwrap(),
3810                released_downstream: release_receivers[stage].take().unwrap(),
3811            }));
3812        }
3813        let head_incoming = forward_receivers[stages - 2].take().unwrap();
3814        let head_release = release_senders[stages - 2].take().unwrap();
3815        // allow: one-shot composite type; naming it would hide the shape that matters here —
3816        // this is the head stage's per-wave output, indexed by wave.
3817        #[allow(clippy::type_complexity)]
3818        let mut results: Vec<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>> =
3819            std::iter::repeat_with(|| None).take(waves.len()).collect();
3820        let walk_result = std::thread::scope(|scope| -> Result<(), Box<dyn std::error::Error>> {
3821            let waves_ref = &waves;
3822            let mut handles = Vec::with_capacity(stages - 1);
3823            // `stage` is a STAGE ID, not merely an index: it indexes two different containers
3824            // (`stage_channels`, `stage_caches.stages()`) and is passed to the worker as the
3825            // stage it owns. An iterator form would keep only one of the three uses.
3826            #[allow(clippy::needless_range_loop)]
3827            for stage in 0..stages - 1 {
3828                let channels = stage_channels[stage].take().unwrap();
3829                let cache_state = &stage_caches.stages()[stage];
3830                handles.push(scope.spawn(move || -> Result<(), String> {
3831                    let result = self.prime_ppn_wave_worker(
3832                        e,
3833                        rt,
3834                        waves_ref,
3835                        cache_state,
3836                        channels.incoming.as_ref(),
3837                        channels.release_upstream.as_ref(),
3838                        &channels.outgoing,
3839                        &channels.released_downstream,
3840                        stage,
3841                        seq_end,
3842                        fence,
3843                        initial_base,
3844                    );
3845                    if let Err(error) = &result {
3846                        channels.notify_failure(&error.to_string());
3847                    }
3848                    result.map_err(|error| error.to_string())
3849                }));
3850            }
3851
3852            let mut head_panic = None;
3853            let head_result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(
3854                || -> Result<(), Box<dyn std::error::Error>> {
3855                    let mut head_cache = stage_caches.stages()[stages - 1]
3856                        .lock()
3857                        .map_err(|_| "prime PP head cache lock poisoned")?;
3858                    for (wave_index, wave) in waves_ref.iter().enumerate() {
3859                        let incoming = recv_prime_pp_signal(
3860                            &head_incoming,
3861                            PrimePpWaveSlot {
3862                                wave: wave_index,
3863                                slot: 0,
3864                            },
3865                            false,
3866                            "prime PP head input",
3867                        )?;
3868                        results[wave_index] = Some(self.prime_ppn_wave_final(
3869                            e,
3870                            rt,
3871                            wave,
3872                            &mut head_cache,
3873                            incoming,
3874                            &head_release,
3875                            seq_end,
3876                            fence,
3877                            initial_base,
3878                        )?);
3879                    }
3880                    Ok(())
3881                },
3882            )) {
3883                Ok(result) => result,
3884                Err(payload) => {
3885                    head_panic = Some(payload);
3886                    Err("prime PP head-stage host walker panicked".into())
3887                }
3888            };
3889            if let Err(error) = &head_result {
3890                let _ = head_release.send(PrimePpSignal::Error(error.to_string()));
3891            }
3892            let mut first_error = head_result.err().map(|error| error.to_string());
3893            let mut worker_panic = None;
3894            for handle in handles {
3895                match handle.join() {
3896                    Ok(Ok(())) => {}
3897                    Ok(Err(error)) => {
3898                        first_error.get_or_insert(error);
3899                    }
3900                    Err(payload) => {
3901                        if worker_panic.is_none() {
3902                            worker_panic = Some(payload);
3903                        }
3904                    }
3905                }
3906            }
3907            if let Some(payload) = head_panic {
3908                std::panic::resume_unwind(payload);
3909            }
3910            if let Some(payload) = worker_panic {
3911                std::panic::resume_unwind(payload);
3912            }
3913            if let Some(error) = first_error {
3914                return Err(error.into());
3915            }
3916            Ok(())
3917        });
3918        let publish_result = if walk_result.is_ok() {
3919            Some(rt.publish_to(stages - 1, &caller_stream))
3920        } else {
3921            None
3922        };
3923        let restore_result = primary_context.restore();
3924        walk_result?;
3925        if let Some(result) = publish_result {
3926            result?;
3927        }
3928        restore_result?;
3929        let mut hiddens = e.uninit(tokens.len() * n_embd)?;
3930        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
3931        for (wave_index, (wave, result)) in waves.iter().zip(results).enumerate() {
3932            debug_assert_eq!((wave.start, wave.end), ranges[wave_index]);
3933            let out = result.ok_or("prime PP wavefront completed without a head-stage result")?;
3934            e.copy_into(
3935                &mut hiddens,
3936                wave.start * n_embd,
3937                &out.2,
3938                (wave.end - wave.start) * n_embd,
3939            )?;
3940            last = Some((out.0, out.1));
3941            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3942        }
3943        stage_caches.commit();
3944        drop(stage_caches);
3945
3946        static LOGGED: std::sync::Once = std::sync::Once::new();
3947        LOGGED.call_once(|| {
3948            eprintln!(
3949                "[pp-wave] PP{stages} prime wavefront engaged: microchunks={} \
3950                 (experimental, MEMRA_PP_WAVE=1)",
3951                ranges.len(),
3952            );
3953        });
3954        let (logits, h_seed) = last.expect("prime PP wavefront produced no microchunk");
3955        crate::pp::record_pp_wave_tick();
3956        Ok((logits, h_seed, hiddens))
3957    }
3958
3959    #[allow(clippy::too_many_arguments)]
3960    fn prime_ppn_wave_worker(
3961        &self,
3962        e: &Engine,
3963        rt: &crate::pp::PpNRt,
3964        waves: &[PrimePpWave<'_>],
3965        cache: &std::sync::Mutex<Cache>,
3966        incoming: Option<&std::sync::mpsc::Receiver<PrimePpSignal>>,
3967        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
3968        outgoing: &std::sync::mpsc::Sender<PrimePpSignal>,
3969        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
3970        stage: usize,
3971        seq_end: usize,
3972        fence: &[usize],
3973        initial_base: usize,
3974    ) -> Result<(), Box<dyn std::error::Error>> {
3975        debug_assert_eq!(incoming.is_some(), stage > 0);
3976        debug_assert_eq!(release_upstream.is_some(), stage > 0);
3977        let mut cache = cache
3978            .lock()
3979            .map_err(|_| "prime PP cache stage lock poisoned")?;
3980        let mut credits = PrimePpWaveCredits::default();
3981        for (wave_index, wave) in waves.iter().enumerate() {
3982            let incoming = match incoming {
3983                Some(receiver) => Some(recv_prime_pp_signal(
3984                    receiver,
3985                    PrimePpWaveSlot {
3986                        wave: wave_index,
3987                        slot: 0,
3988                    },
3989                    false,
3990                    "prime PP stage input",
3991                )?),
3992                None => None,
3993            };
3994            let sent = self.prime_ppn_wave_stage(
3995                e,
3996                rt,
3997                wave,
3998                &mut cache,
3999                stage,
4000                incoming,
4001                release_upstream,
4002                &mut credits,
4003                released_downstream,
4004                seq_end,
4005                fence,
4006                initial_base,
4007            )?;
4008            send_prime_pp_signal(outgoing, PrimePpSignal::Slot(sent), "prime PP stage output")?;
4009        }
4010        while let Some(expected) = credits.pending.front().copied() {
4011            let released = recv_prime_pp_signal(
4012                released_downstream,
4013                expected,
4014                true,
4015                "prime PP final slot release",
4016            )?;
4017            credits.record_release(released)?;
4018        }
4019        Ok(())
4020    }
4021
4022    #[allow(clippy::too_many_arguments)]
4023    fn prime_ppn_wave_stage(
4024        &self,
4025        e: &Engine,
4026        rt: &crate::pp::PpNRt,
4027        wave: &PrimePpWave<'_>,
4028        cache: &mut Cache,
4029        stage: usize,
4030        incoming: Option<PrimePpWaveSlot>,
4031        release_upstream: Option<&std::sync::mpsc::Sender<PrimePpSignal>>,
4032        credits: &mut PrimePpWaveCredits,
4033        released_downstream: &std::sync::mpsc::Receiver<PrimePpSignal>,
4034        seq_end: usize,
4035        fence: &[usize],
4036        initial_base: usize,
4037    ) -> Result<PrimePpWaveSlot, Box<dyn std::error::Error>> {
4038        debug_assert!(stage + 1 < fence.len() - 1);
4039        let t = wave.end - wave.start;
4040        let base = initial_base + wave.start;
4041        debug_assert_eq!(cache.pos, base, "prime PP stage advanced out of order");
4042        let n_embd = self.cfg.n_embd as usize;
4043        let payload = t * n_embd;
4044        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4045        rt.bind_stage(stage)?;
4046        let _stage = rt.enter(stage);
4047        let engine = rt.engine(stage, e);
4048        let positions_d = engine.htod_i32(&positions)?;
4049        let x = if stage == 0 {
4050            debug_assert!(incoming.is_none());
4051            self.embed(engine, wave.tokens)?
4052        } else {
4053            let incoming = incoming.ok_or("prime PP stage has no incoming boundary slot")?;
4054            let x = rt.rx(stage - 1, incoming.slot, payload)?;
4055            send_prime_pp_signal(
4056                release_upstream.ok_or("prime PP stage has no upstream release channel")?,
4057                PrimePpSignal::Slot(incoming),
4058                "prime PP upstream slot release",
4059            )?;
4060            x
4061        };
4062        let x = {
4063            let _wave_cell = crate::pp::enter_pp_wave_cell();
4064            let _overlap = crate::pp::enter_prime_pipe_stage();
4065            self.prime_layers(
4066                engine,
4067                x,
4068                fence[stage],
4069                fence[stage + 1],
4070                &positions_d,
4071                t,
4072                base,
4073                cache,
4074                seq_end,
4075            )?
4076        };
4077        if let Some(expected) = credits.release_required() {
4078            let released =
4079                recv_prime_pp_signal(released_downstream, expected, true, "prime PP slot credit")?;
4080            credits.record_release(released)?;
4081        }
4082        let sent = PrimePpWaveSlot {
4083            wave: credits.next_wave,
4084            slot: rt.tx_pipelined(stage, &x, payload)?,
4085        };
4086        credits.record_send(sent)?;
4087        cache.pos = base + t;
4088        Ok(sent)
4089    }
4090
4091    #[allow(clippy::too_many_arguments)]
4092    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4093    fn prime_ppn_wave_final(
4094        &self,
4095        e: &Engine,
4096        rt: &crate::pp::PpNRt,
4097        wave: &PrimePpWave<'_>,
4098        cache: &mut Cache,
4099        incoming: PrimePpWaveSlot,
4100        release_upstream: &std::sync::mpsc::Sender<PrimePpSignal>,
4101        seq_end: usize,
4102        fence: &[usize],
4103        initial_base: usize,
4104    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4105        let stage = fence.len() - 2;
4106        let t = wave.end - wave.start;
4107        let base = initial_base + wave.start;
4108        debug_assert_eq!(cache.pos, base, "prime PP head stage advanced out of order");
4109        let n_embd = self.cfg.n_embd as usize;
4110        let payload = t * n_embd;
4111        let positions: Vec<i32> = (base as i32..(base + t) as i32).collect();
4112        rt.bind_stage(stage)?;
4113        let _stage = rt.enter(stage);
4114        let engine = rt.engine(stage, e);
4115        let positions_d = engine.htod_i32(&positions)?;
4116        let x = rt.rx(stage - 1, incoming.slot, payload)?;
4117        send_prime_pp_signal(
4118            release_upstream,
4119            PrimePpSignal::Slot(incoming),
4120            "prime PP head slot release",
4121        )?;
4122        let _wave_cell = crate::pp::enter_pp_wave_cell();
4123        let _overlap = crate::pp::enter_prime_pipe_stage();
4124        let x = self.prime_layers(
4125            engine,
4126            x,
4127            fence[stage],
4128            fence[stage + 1],
4129            &positions_d,
4130            t,
4131            base,
4132            cache,
4133            seq_end,
4134        )?;
4135        self.prime_chunk_epilogue(engine, x, t, cache)
4136    }
4137
4138    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
4139    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
4140    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
4141    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
4142    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
4143    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
4144    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
4145        if Engine::gdn_db_on()
4146            && Engine::gdn_chunked_enabled()
4147            && t >= 16
4148            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
4149            && num_k * 2 == num_v
4150        {
4151            num_k
4152        } else {
4153            num_v
4154        }
4155    }
4156
4157    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
4158    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
4159    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
4160    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
4161    fn f16out_on(e: &Engine, t: usize) -> bool {
4162        crate::f16_ffi::pp_f16_enabled()
4163            && t >= 16
4164            && !e.verify_exact_on()
4165            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
4166    }
4167
4168    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
4169    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
4170    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
4171    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
4172    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
4173    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
4174    /// see one entry, byte-identical behavior.
4175    pub fn prime_slabs_get(
4176        &self,
4177        e: &Engine,
4178        t: usize,
4179        n_embd: usize,
4180        n_ff_max: usize,
4181    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
4182        let mut slabs = self.prime_slabs.lock().unwrap();
4183        let dev = e.ctx().ordinal();
4184        let need_new = match slabs.get(&dev) {
4185            None => true,
4186            Some(sl) => sl.lock().unwrap().t_cap < t,
4187        };
4188        if need_new {
4189            slabs.insert(
4190                dev,
4191                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
4192                    t_cap: t,
4193                    h: e.uninit(t * n_embd)?,
4194                    x1: e.uninit(t * n_embd)?,
4195                    z: e.uninit(t * n_embd)?,
4196                    act: e.uninit(t * n_ff_max)?,
4197                    xa: e.uninit(t * n_embd)?,
4198                    xb: e.uninit(t * n_embd)?,
4199                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
4200                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
4201                    gate: e.uninit(t * n_ff_max)?,
4202                    up: e.uninit(t * n_ff_max)?,
4203                    ffn_out: e.uninit(t * n_embd)?,
4204                    seg_glue: Vec::new(),
4205                    mixed: e.uninit(t * n_embd)?,
4206                    seg_mid: Vec::new(),
4207                    seg_t: 0,
4208                })),
4209            );
4210        }
4211        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
4212    }
4213
4214    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
4215    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
4216    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
4217    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4218    fn prime_chunk(
4219        &self,
4220        e: &Engine,
4221        tokens: &[u32],
4222        cache: &mut Cache,
4223        seq_end: usize,
4224        chunk_off: usize,
4225        overlay: Option<&crate::vision::EmbedOverlay>,
4226    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4227        if crate::pp::pp_host_bounce_active()
4228            && (self.uses_gemma_program() || !crate::pp::prime_pp_on())
4229        {
4230            return Err(
4231                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
4232                 has no active prime stage split and would peer-read remote weights; keep \
4233                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
4234                    .into(),
4235            );
4236        }
4237        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
4238        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
4239        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
4240        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
4241        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
4242        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
4243        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
4244        // loader is off and there is nothing remote to split for.
4245        if !self.uses_gemma_program()
4246            && !crate::pp::pp2_streams_off()
4247            && crate::pp::prime_pp_on()
4248            && let Some(fence) = crate::pp::pp_cuts(self.layers.len())
4249        {
4250            if overlay.is_some() {
4251                return Err("vision embedding overlay + PP prime unsupported (v1); \
4252                         run single-device or MEMRA_PRIME_PP=0"
4253                    .into());
4254            }
4255            return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
4256        }
4257        if crate::pp::pp_host_bounce_active() {
4258            return Err(
4259                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
4260                 refusing an unsplit remote-weight walk"
4261                    .into(),
4262            );
4263        }
4264        let t = tokens.len();
4265        let base = cache.pos;
4266        debug_assert!(
4267            seq_end >= base + t,
4268            "prime_chunk: seq_end must cover this chunk"
4269        );
4270        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
4271        let pos_d = e.htod_i32(&pos)?;
4272
4273        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
4274        if let Some(ov) = overlay {
4275            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
4276            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
4277            // Images larger than one prime chunk straddle boundaries, hence the clipping.
4278            let n_embd = self.cfg.n_embd as usize;
4279            for &(pos, row_off, n_rows) in &ov.spans {
4280                let lo = pos.max(chunk_off);
4281                let hi = (pos + n_rows).min(chunk_off + t);
4282                if lo < hi {
4283                    let src_row = row_off + (lo - pos);
4284                    let view = ov
4285                        .rows
4286                        .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
4287                    e.copy_view_into(
4288                        &mut x_embed,
4289                        (lo - chunk_off) * n_embd,
4290                        &view,
4291                        (hi - lo) * n_embd,
4292                    )?;
4293                }
4294            }
4295        }
4296        let x = self.prime_layers(
4297            e,
4298            x_embed,
4299            0,
4300            self.layers.len(),
4301            &pos_d,
4302            t,
4303            base,
4304            cache,
4305            seq_end,
4306        )?;
4307        self.prime_chunk_epilogue(e, x, t, cache)
4308    }
4309
4310    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
4311    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
4312    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
4313    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
4314    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
4315    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
4316    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
4317    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
4318    ///     the plain add (materialize) and the next stage hoists its own first norm — the
4319    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
4320    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
4321    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
4322    ///     each stage walks through its own resident transients;
4323    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
4324    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
4325    #[allow(clippy::too_many_arguments)]
4326    fn prime_layers(
4327        &self,
4328        e: &Engine,
4329        x_in: CudaSlice<f32>,
4330        lo: usize,
4331        hi: usize,
4332        pos_d: &CudaSlice<i32>,
4333        t: usize,
4334        base: usize,
4335        cache: &mut Cache,
4336        seq_end: usize,
4337    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4338        let cfg = &self.cfg;
4339        let n_embd = cfg.n_embd as usize;
4340        let eps = cfg.rms_eps;
4341        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
4342        // standalone convert launches). Only when the f16 lane serves and T reaches the
4343        // GEMM tier; bit-identical either way.
4344        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
4345        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
4346        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
4347        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
4348        // live prefix is fully overwritten before use; x ping-pongs xa<->xb; the inactive
4349        // capacity tail must stay behind checked views. The hidden-stack return clones the
4350        // final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
4351        let n_ff_max = self
4352            .layers
4353            .iter()
4354            .map(|l| match &l.ffn {
4355                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
4356                _ => n_embd,
4357            })
4358            .max()
4359            .unwrap_or(n_embd)
4360            .max(n_embd);
4361        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
4362        let slab = if use_slabs {
4363            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
4364        } else {
4365            None
4366        };
4367        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
4368        let mut x_own; // fallback storage when slabs are off
4369        type SlabRefs<'a> = (
4370            &'a mut CudaSlice<f32>,
4371            &'a mut CudaSlice<f32>,
4372            &'a mut CudaSlice<f32>,
4373            &'a mut CudaSlice<f32>,
4374            &'a mut CudaSlice<u8>,
4375            &'a mut CudaSlice<u8>,
4376            &'a mut CudaSlice<f32>,
4377            &'a mut CudaSlice<f32>,
4378            &'a mut CudaSlice<f32>,
4379        );
4380        let (mut x_cur, mut x_nxt, sl): (
4381            &mut CudaSlice<f32>,
4382            &mut CudaSlice<f32>,
4383            Option<SlabRefs>,
4384        );
4385        #[allow(clippy::type_complexity)]
4386        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4387        let mut seg: Option<(
4388            &mut Vec<Option<cudarc::driver::CudaGraph>>,
4389            &mut Vec<Option<cudarc::driver::CudaGraph>>,
4390            &mut CudaSlice<f32>,
4391            &mut usize,
4392        )> = None;
4393        let mut x_own2;
4394        match slab_guard.as_mut() {
4395            Some(g) => {
4396                let slabs = &mut **g;
4397                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
4398                let PrimeSlabs {
4399                    xa,
4400                    xb,
4401                    h,
4402                    x1,
4403                    z,
4404                    act,
4405                    h16,
4406                    z16,
4407                    gate,
4408                    up,
4409                    ffn_out,
4410                    seg_glue,
4411                    mixed,
4412                    seg_mid,
4413                    seg_t,
4414                    ..
4415                } = slabs;
4416                x_cur = xa;
4417                x_nxt = xb;
4418                seg = Some((seg_glue, seg_mid, mixed, seg_t));
4419                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
4420            }
4421            None => {
4422                x_own = x_in;
4423                x_own2 = e.uninit(t * n_embd)?;
4424                x_cur = &mut x_own;
4425                x_nxt = &mut x_own2;
4426                sl = None;
4427            }
4428        }
4429        let mut alloc_h;
4430        let mut alloc_x1;
4431        let mut alloc_z;
4432        let mut alloc_act;
4433        let mut alloc_h16;
4434        let mut alloc_z16;
4435        let mut alloc_gate;
4436        let mut alloc_up;
4437        let mut alloc_fo;
4438        let (h, x1, z, act): (
4439            &mut CudaSlice<f32>,
4440            &mut CudaSlice<f32>,
4441            &mut CudaSlice<f32>,
4442            &mut CudaSlice<f32>,
4443        );
4444        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
4445        let (sl_gate, sl_up, sl_fo): (
4446            &mut CudaSlice<f32>,
4447            &mut CudaSlice<f32>,
4448            &mut CudaSlice<f32>,
4449        );
4450        match sl {
4451            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
4452                h = a;
4453                x1 = b;
4454                z = c;
4455                act = d;
4456                h16 = e16;
4457                z16 = f16b;
4458                sl_gate = g;
4459                sl_up = u;
4460                sl_fo = fo;
4461            }
4462            None => {
4463                alloc_h = e.uninit(t * n_embd)?;
4464                alloc_x1 = e.uninit(t * n_embd)?;
4465                alloc_z = e.uninit(t * n_embd)?;
4466                alloc_act = e.uninit(t * n_ff_max)?;
4467                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4468                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
4469                alloc_gate = e.uninit(t * n_ff_max)?;
4470                alloc_up = e.uninit(t * n_ff_max)?;
4471                alloc_fo = e.uninit(t * n_embd)?;
4472                h = &mut alloc_h;
4473                x1 = &mut alloc_x1;
4474                z = &mut alloc_z;
4475                act = &mut alloc_act;
4476                h16 = &mut alloc_h16;
4477                z16 = &mut alloc_z16;
4478                sl_gate = &mut alloc_gate;
4479                sl_up = &mut alloc_up;
4480                sl_fo = &mut alloc_fo;
4481            }
4482        }
4483        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
4484        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
4485        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
4486        // first prime at this t (capture does not execute -> launch right after).
4487        let n_layers = self.layers.len();
4488        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
4489        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
4490        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
4491        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
4492        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
4493        // machinery stays (byte-identical) as their foundation.
4494        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
4495        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
4496        // step35 rides its own mixer through the normal per-layer arm below.
4497        let use_seg = f16fuse
4498            && seg.is_some()
4499            && !self.uses_sliding_gated_moe_program()
4500            && lo == 0
4501            && hi == n_layers
4502            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1")
4503            // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): below the
4504            // driver-free floor this prime call takes the eager fused else-arm (the
4505            // byte-identical twin the opt-in was gated against) instead of replaying
4506            // the S-mid/S-glue segment graphs into an exhausted card. Probe runs only
4507            // when the opt-in flag is armed (short-circuit order).
4508            && {
4509                let ok = crate::spec::graph_launch_headroom_ok(e);
4510                if !ok {
4511                    static NOTED: std::sync::Once = std::sync::Once::new();
4512                    NOTED.call_once(|| crate::spec::graph_replay_suspended_note("prime-seg"));
4513                }
4514                ok
4515            };
4516        if let Some((sg, sm, _, st)) = seg.as_mut()
4517            && **st != t
4518        {
4519            sg.clear();
4520            sg.extend((0..n_layers).map(|_| None));
4521            sm.clear();
4522            sm.extend((0..n_layers).map(|_| None));
4523            **st = t;
4524        }
4525        {
4526            let layer_lo = &self.layers[lo];
4527            if f16fuse {
4528                e.rms_norm_f16out(
4529                    x_cur,
4530                    layer_lo.attn_norm.float_data(),
4531                    h,
4532                    h16,
4533                    n_embd,
4534                    t,
4535                    eps,
4536                )?;
4537            } else {
4538                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
4539            }
4540        }
4541        let anat = Self::prime_anatomy_on();
4542        let mut anat_last = if anat {
4543            e.stream().synchronize()?;
4544            Some(std::time::Instant::now())
4545        } else {
4546            None
4547        };
4548        // Closes the region that just ENDED into `slot`, restarting the clock.
4549        macro_rules! anat_mark {
4550            ($slot:expr) => {
4551                if let Some(ts) = anat_last.as_mut() {
4552                    e.stream().synchronize()?;
4553                    Self::prime_anatomy_slots()[$slot].fetch_add(
4554                        ts.elapsed().as_nanos() as u64,
4555                        std::sync::atomic::Ordering::Relaxed,
4556                    );
4557                    *ts = std::time::Instant::now();
4558                }
4559            };
4560        }
4561        for il in lo..hi {
4562            let layer = &self.layers[il];
4563            let hx16 = if f16fuse { Some(&*h16) } else { None };
4564            if use_seg {
4565                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
4566                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
4567                let (pre, pre16, w_out) = match &layer.mixer {
4568                    Mixer::Full(fa) => {
4569                        let g3 = match hx16 {
4570                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
4571                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
4572                        };
4573                        let (pre, pre16) =
4574                            self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
4575                        (pre, pre16, &fa.wo)
4576                    }
4577                    Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("core-split prime"),
4578                    Mixer::Kda(_) => {
4579                        crate::hybrid::kda_path_unimplemented("core-split captured prime")
4580                    }
4581                    Mixer::Linear(la) => {
4582                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
4583                        let g4 = match hx16 {
4584                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
4585                            None => e.matmul_group(&ws, h, t)?,
4586                        };
4587                        let (pre, pre16) =
4588                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
4589                        (pre, pre16, &la.ssm_out)
4590                    }
4591                };
4592                {
4593                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
4594                    let pre_n = pre.len() / t;
4595                    let xh_pre = match pre16 {
4596                        Some(x) => x,
4597                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
4598                    };
4599                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
4600                        let y = e.matmul(w_out, &pre, t)?;
4601                        e.copy_into(mslab, 0, &y, t * n_embd)?;
4602                    }
4603                    if sm[il].is_none() {
4604                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4605                        let w_post = layer.post_attn_norm.float_data();
4606                        e.stream().synchronize()?;
4607                        e.stream()
4608                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4609                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4610                            e.add(x_cur, mslab, x1, t * n_embd)?;
4611                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
4612                            Ok(())
4613                        })();
4614                        let g = e.stream().end_capture(
4615                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
4616                        r?;
4617                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
4618                    }
4619                    sm[il].as_ref().unwrap().launch()?;
4620                }
4621            } else {
4622                let mixed = match &layer.mixer {
4623                    Mixer::Full(fa) => {
4624                        let y =
4625                            self.full_attn_prime(e, fa, h, hx16, pos_d, t, cache, il, seq_end)?;
4626                        anat_mark!(0);
4627                        y
4628                    }
4629                    Mixer::Linear(la) => {
4630                        let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
4631                        anat_mark!(1);
4632                        y
4633                    }
4634                    Mixer::Mla(mla) => self.mla_attn_cached(e, mla, h, pos_d, t, il, cache)?,
4635                    Mixer::Kda(la) => {
4636                        let y = crate::kda::kda_prime_cached(e, la, h, t, eps, cache, il)?;
4637                        anat_mark!(1);
4638                        y
4639                    }
4640                };
4641                if f16fuse {
4642                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
4643                    // bit-identical) — the standalone add pass disappears.
4644                    e.add_rms_norm_f16out(
4645                        x_cur,
4646                        &mixed,
4647                        layer.post_attn_norm.float_data(),
4648                        x1,
4649                        z,
4650                        z16,
4651                        n_embd,
4652                        t,
4653                        eps,
4654                    )?;
4655                } else {
4656                    e.add(x_cur, &mixed, x1, t * n_embd)?;
4657                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
4658                }
4659                anat_mark!(4);
4660            }
4661            let zx16 = if f16fuse { Some(&*z16) } else { None };
4662            match &layer.ffn {
4663                crate::hybrid::Ffn::Dense {
4664                    ffn_gate,
4665                    ffn_up,
4666                    ffn_down,
4667                } => {
4668                    let n_ff = ffn_gate.out_features();
4669                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
4670                    // the allocating group + copy when a mirror is missing.
4671                    let mut into_ok = false;
4672                    if let Some(xh) = zx16 {
4673                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
4674                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
4675                    }
4676                    if !into_ok {
4677                        let mut g2 = match zx16 {
4678                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
4679                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
4680                        };
4681                        let up_y = g2.pop().unwrap();
4682                        let gate_y = g2.pop().unwrap();
4683                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
4684                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
4685                    }
4686                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
4687                    // operand in-epilogue; non-silu activations keep the standalone convert.
4688                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
4689                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
4690                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
4691                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
4692                    {
4693                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
4694                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
4695                        Some(a16)
4696                    } else {
4697                        Self::ffn_act_lim(
4698                            e,
4699                            &self.cfg,
4700                            sl_gate,
4701                            sl_up,
4702                            1.0,
4703                            1.0,
4704                            d_lim,
4705                            act,
4706                            t * n_ff,
4707                        )?;
4708                        None
4709                    };
4710                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
4711                    let xh_act = match act16 {
4712                        Some(x) => x,
4713                        None => e.f16_act(act, t * n_ff, n_ff)?,
4714                    };
4715                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
4716                        let y = e.matmul(ffn_down, &*act, t)?;
4717                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4718                    }
4719                }
4720                crate::hybrid::Ffn::Moe(m) => {
4721                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
4722                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
4723                    anat_mark!(2);
4724                }
4725            }
4726            if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
4727                anat_mark!(3);
4728            }
4729            if use_seg && il + 1 < hi {
4730                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
4731                let w_next = self.layers[il + 1].attn_norm.float_data();
4732                let (sg, _, _, _) = seg.as_mut().unwrap();
4733                if sg[il].is_none() {
4734                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
4735                    e.stream().synchronize()?;
4736                    e.stream()
4737                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
4738                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
4739                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4740                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
4741                        Ok(())
4742                    })();
4743                    let g = e.stream().end_capture(
4744                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
4745                    );
4746                    r?;
4747                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
4748                }
4749                sg[il].as_ref().unwrap().launch()?;
4750            } else {
4751                if il + 1 < hi {
4752                    let w_next = self.layers[il + 1].attn_norm.float_data();
4753                    if f16fuse {
4754                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
4755                    } else {
4756                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4757                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
4758                    }
4759                } else {
4760                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
4761                }
4762            }
4763            anat_mark!(4);
4764            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
4765            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
4766            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
4767            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
4768            // unset (the default) costs one OnceLock read per layer.
4769            if let Some(path) = Self::prime_trace_path() {
4770                let row = base + t - 1;
4771                let host = e.dtoh(x_nxt)?;
4772                let last = &host[(t - 1) * n_embd..t * n_embd];
4773                use std::io::Write as _;
4774                let mut f = std::fs::OpenOptions::new()
4775                    .create(true)
4776                    .append(true)
4777                    .open(path)?;
4778                let mut h64: u64 = 0xcbf29ce484222325;
4779                for v in last {
4780                    h64 ^= v.to_bits() as u64;
4781                    h64 = h64.wrapping_mul(0x100000001b3);
4782                }
4783                writeln!(
4784                    f,
4785                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
4786                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
4787                    last[0], last[1], last[2]
4788                )?;
4789            }
4790            // dflash/dspark tap (no-op when no sink armed): post-layer residual rows for
4791            // drafter conditioning — the qwen twin of the gemma4 tap sites.
4792            self.dflash_tap(e, cache, il, x_nxt, t)?;
4793            std::mem::swap(&mut x_cur, &mut x_nxt);
4794        }
4795        if anat {
4796            let s = Self::prime_anatomy_slots();
4797            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
4798            eprintln!(
4799                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
4800                 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
4801                ms(0),
4802                ms(1),
4803                ms(2),
4804                ms(3),
4805                ms(4)
4806            );
4807        }
4808        // hidden-stack return: clone the final x out of the slab
4809        let mut x = e.uninit(t * n_embd)?;
4810        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
4811        drop(slab_guard);
4812        Ok(x)
4813    }
4814
4815    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
4816    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
4817    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
4818    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
4819    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4820    fn prime_chunk_epilogue(
4821        &self,
4822        e: &Engine,
4823        x: CudaSlice<f32>,
4824        t: usize,
4825        cache: &mut Cache,
4826    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4827        let n_embd = self.cfg.n_embd as usize;
4828        let eps = self.cfg.rms_eps;
4829        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
4830        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
4831        // the post-norm copy happens after hn exists).
4832        let mut h_seed = e.uninit(n_embd)?;
4833        if !crate::spec::spec_hpost() {
4834            e.copy_view_into(
4835                &mut h_seed,
4836                0,
4837                &x.slice((t - 1) * n_embd..t * n_embd),
4838                n_embd,
4839            )?;
4840        }
4841        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
4842        let mut hn = e.uninit(t * n_embd)?;
4843        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4844        if crate::spec::spec_hpost() {
4845            e.copy_view_into(
4846                &mut h_seed,
4847                0,
4848                &hn.slice((t - 1) * n_embd..t * n_embd),
4849                n_embd,
4850            )?;
4851        }
4852        let last = e.view(&hn, t * n_embd);
4853        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
4854        let mut hlast = e.uninit(n_embd)?;
4855        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
4856        let logits = e.matmul(&self.output, &hlast, 1)?;
4857        cache.pos += t;
4858        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
4859        // post-norm stack hn (MEMRA_SPEC_HPOST).
4860        Ok((
4861            e.dtoh(&logits)?,
4862            h_seed,
4863            if crate::spec::spec_hpost() { hn } else { x },
4864        ))
4865    }
4866
4867    /// Post-final-norm hidden state of one row of a prime-returned hidden stack — the
4868    /// embedding-pooling read (lane/embed-serve). `hiddens` is `prime_cache`'s third
4869    /// return: the pre-norm stack by default, but ALREADY post-norm under
4870    /// MEMRA_SPEC_HPOST (see `prime_chunk_epilogue`), so the norm is applied only in
4871    /// the default shape. Returns the host f32 row (`n_embd` wide).
4872    pub fn hidden_postnorm_row(
4873        &self,
4874        e: &Engine,
4875        hiddens: &CudaSlice<f32>,
4876        row: usize,
4877    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4878        let n_embd = self.cfg.n_embd as usize;
4879        let mut x1 = e.uninit(n_embd)?;
4880        e.copy_view_into(
4881            &mut x1,
4882            0,
4883            &hiddens.slice(row * n_embd..(row + 1) * n_embd),
4884            n_embd,
4885        )?;
4886        if crate::spec::spec_hpost() {
4887            return e.dtoh(&x1);
4888        }
4889        let mut hn = e.uninit(n_embd)?;
4890        e.rms_norm(
4891            &x1,
4892            self.output_norm.float_data(),
4893            &mut hn,
4894            n_embd,
4895            1,
4896            self.cfg.rms_eps,
4897        )?;
4898        e.dtoh(&hn)
4899    }
4900
4901    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
4902    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
4903    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
4904    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
4905    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
4906    /// prefill kernels. Structure mirrors the verify split exactly:
4907    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
4908    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
4909    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
4910    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
4911    ///                  there via the sharded loader) → `publish_to`
4912    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
4913    /// round's stage-freed buffers must not be reused under the caller's queued reads);
4914    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
4915    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
4916    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
4917    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
4918    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
4919    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
4920    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
4921    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
4922    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
4923    /// and its liveness counter is bumped here — the gate goes green with this function.
4924    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4925    fn prime_chunk_ppn(
4926        &self,
4927        e: &Engine,
4928        tokens: &[u32],
4929        cache: &mut Cache,
4930        seq_end: usize,
4931        fence: &[usize],
4932    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4933        let rt = crate::pp::PpNRt::get(e)?;
4934        let n_st = fence.len() - 1;
4935        assert_eq!(
4936            rt.n_stages(),
4937            n_st,
4938            "PpNRt stage count {} != fence stages {n_st}",
4939            rt.n_stages()
4940        );
4941        let n_embd = self.cfg.n_embd as usize;
4942        let t = tokens.len();
4943        let base = cache.pos;
4944        debug_assert!(
4945            seq_end >= base + t,
4946            "prime_chunk_ppn: seq_end must cover this chunk"
4947        );
4948        let payload = t * n_embd;
4949        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
4950        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
4951        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
4952        let caller_stream = e.stream();
4953        rt.fence_stages_behind(&caller_stream)?;
4954
4955        if n_st == 2 {
4956            let slot =
4957                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
4958            let x =
4959                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
4960            let out = {
4961                rt.bind_stage(1)?;
4962                let _st1 = rt.enter(1);
4963                let e1 = rt.engine(1, e);
4964                self.prime_chunk_epilogue(e1, x, t, cache)?
4965            };
4966            rt.publish_to(1, &caller_stream)?;
4967            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4968            return Ok(out);
4969        }
4970
4971        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
4972
4973        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
4974        let mut slot = {
4975            let _st0 = rt.enter(0);
4976            let e0 = rt.engine(0, e);
4977            let pos_d = e0.htod_i32(&pos)?;
4978            let x = self.embed(e0, tokens)?;
4979            let x =
4980                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
4981            rt.tx(0, &x, payload)?
4982            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
4983        };
4984
4985        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
4986        for s in 1..n_st - 1 {
4987            let _st = rt.enter(s);
4988            let es = rt.engine(s, e);
4989            let pos_d = es.htod_i32(&pos)?;
4990            let x = rt.rx(s - 1, slot, payload)?;
4991            let x = self.prime_layers(
4992                es,
4993                x,
4994                fence[s],
4995                fence[s + 1],
4996                &pos_d,
4997                t,
4998                base,
4999                cache,
5000                seq_end,
5001            )?;
5002            slot = rt.tx(s, &x, payload)?;
5003        }
5004
5005        // ---- LAST STAGE: RX + final range + the shared epilogue ----
5006        let _stl = rt.enter(n_st - 1);
5007        let el = rt.engine(n_st - 1, e);
5008        let pos_d = el.htod_i32(&pos)?;
5009        let x = rt.rx(n_st - 2, slot, payload)?;
5010        let x = self.prime_layers(
5011            el,
5012            x,
5013            fence[n_st - 1],
5014            fence[n_st],
5015            &pos_d,
5016            t,
5017            base,
5018            cache,
5019            seq_end,
5020        )?;
5021        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
5022        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
5023        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
5024        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
5025        // stage stream host-side, but the law is stated in events, not in a dtoh side
5026        // effect a later deferred form would remove.
5027        rt.publish_to(n_st - 1, &caller_stream)?;
5028        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5029        Ok(out)
5030    }
5031
5032    #[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
5033    fn prime_pp2_stage0_enqueue(
5034        &self,
5035        e: &Engine,
5036        rt: &crate::pp::PpNRt,
5037        tokens: &[u32],
5038        cache: &mut Cache,
5039        seq_end: usize,
5040        fence: &[usize],
5041        base: usize,
5042        pipelined: bool,
5043    ) -> Result<usize, Box<dyn std::error::Error>> {
5044        let t = tokens.len();
5045        let n_embd = self.cfg.n_embd as usize;
5046        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5047        rt.bind_stage(0)?;
5048        let _st0 = rt.enter(0);
5049        let e0 = rt.engine(0, e);
5050        let pos_d = e0.htod_i32(&pos)?;
5051        let x = self.embed(e0, tokens)?;
5052        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5053        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
5054        if pipelined {
5055            rt.tx_pipelined(0, &x, t * n_embd)
5056        } else {
5057            rt.tx(0, &x, t * n_embd)
5058        }
5059    }
5060
5061    #[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
5062    fn prime_pp2_stage1_enqueue(
5063        &self,
5064        e: &Engine,
5065        rt: &crate::pp::PpNRt,
5066        slot: usize,
5067        t: usize,
5068        cache: &mut Cache,
5069        seq_end: usize,
5070        fence: &[usize],
5071        base: usize,
5072        pipelined: bool,
5073    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5074        let n_embd = self.cfg.n_embd as usize;
5075        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
5076        rt.bind_stage(1)?;
5077        let _st1 = rt.enter(1);
5078        let e1 = rt.engine(1, e);
5079        let pos_d = e1.htod_i32(&pos)?;
5080        let x = rt.rx(0, slot, t * n_embd)?;
5081        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
5082        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
5083    }
5084
5085    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
5086    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
5087    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
5088    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
5089    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
5090    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
5091    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
5092    /// bookkeeping still runs on the host per call — the real replay path moves the write
5093    /// slot to the len_d device counter (increment 3).
5094    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
5095    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
5096    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
5097    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
5098    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
5099    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
5100    #[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
5101    pub fn prime_chunk_captured(
5102        &self,
5103        e: &Engine,
5104        x_in: &CudaSlice<f32>,
5105        pos_d: &CudaSlice<i32>,
5106        t: usize,
5107        cache: &mut Cache,
5108        len_d: &CudaSlice<i32>,
5109        logits_out: &mut CudaSlice<f32>,
5110        h_seed_out: &mut CudaSlice<f32>,
5111    ) -> Result<(), Box<dyn std::error::Error>> {
5112        self.refuse_hyper("prime_chunk_captured")?;
5113        cache.ensure_usable("prime_chunk_captured")?;
5114        let cfg = &self.cfg;
5115        let n_embd = cfg.n_embd as usize;
5116        let eps = cfg.rms_eps;
5117        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
5118        let mut x = e.uninit(t * n_embd)?;
5119        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
5120        for (il, layer) in self.layers.iter().enumerate() {
5121            let mut h = e.uninit(t * n_embd)?;
5122            let mut hx16: Option<CudaSlice<u8>> = None;
5123            if f16fuse {
5124                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5125                e.rms_norm_f16out(
5126                    &x,
5127                    layer.attn_norm.float_data(),
5128                    &mut h,
5129                    &mut b16,
5130                    n_embd,
5131                    t,
5132                    eps,
5133                )?;
5134                hx16 = Some(b16);
5135            } else {
5136                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5137            }
5138            let mixed = match &layer.mixer {
5139                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
5140                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
5141                // come from the caller (see step35_attn_pre_wo's doc note).
5142                Mixer::Full(fa) => {
5143                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
5144                }
5145                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("captured-graph prime"),
5146                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("captured prime chunk"),
5147                Mixer::Linear(la) => {
5148                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
5149                    let g4 = match hx16.as_ref() {
5150                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
5151                        None => e.matmul_group(&ws, &h, t)?,
5152                    };
5153                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
5154                }
5155            };
5156            let mut x1 = e.uninit(t * n_embd)?;
5157            e.add(&x, &mixed, &mut x1, t * n_embd)?;
5158            let mut z = e.uninit(t * n_embd)?;
5159            let mut zx16: Option<CudaSlice<u8>> = None;
5160            if f16fuse {
5161                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
5162                e.rms_norm_f16out(
5163                    &x1,
5164                    layer.post_attn_norm.float_data(),
5165                    &mut z,
5166                    &mut b16,
5167                    n_embd,
5168                    t,
5169                    eps,
5170                )?;
5171                zx16 = Some(b16);
5172            } else {
5173                e.rms_norm(
5174                    &x1,
5175                    layer.post_attn_norm.float_data(),
5176                    &mut z,
5177                    n_embd,
5178                    t,
5179                    eps,
5180                )?;
5181            }
5182            let ffn_out = match &layer.ffn {
5183                crate::hybrid::Ffn::Dense {
5184                    ffn_gate,
5185                    ffn_up,
5186                    ffn_down,
5187                } => {
5188                    let n_ff = ffn_gate.out_features();
5189                    let mut g2 = match &zx16 {
5190                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
5191                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
5192                    };
5193                    let up = g2.pop().unwrap();
5194                    let gate = g2.pop().unwrap();
5195                    let mut act = e.uninit(t * n_ff)?;
5196                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
5197                    Self::ffn_act_lim(
5198                        e,
5199                        &self.cfg,
5200                        &gate,
5201                        &up,
5202                        1.0,
5203                        1.0,
5204                        self.cfg.clamp_shexp_at(il as u32),
5205                        &mut act,
5206                        t * n_ff,
5207                    )?;
5208                    e.matmul(ffn_down, &act, t)?
5209                }
5210                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
5211            };
5212            let mut x2 = e.uninit(t * n_embd)?;
5213            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5214            x = x2;
5215        }
5216        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
5217        if !crate::spec::spec_hpost() {
5218            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
5219        }
5220        let mut hn = e.uninit(t * n_embd)?;
5221        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5222        if crate::spec::spec_hpost() {
5223            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
5224        }
5225        let mut hlast = e.uninit(n_embd)?;
5226        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
5227        let logits = e.matmul(&self.output, &hlast, 1)?;
5228        let nv = logits.len();
5229        e.copy_into(logits_out, 0, &logits, nv)?;
5230        Ok(())
5231    }
5232
5233    fn step35_prime_batch_on() -> bool {
5234        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
5235    }
5236
5237    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
5238    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
5239    #[allow(clippy::too_many_arguments)]
5240    /// `seq_ends[s]`: sequence s's REQUEST-absolute end position — NOT `ts[s]`. It is the
5241    /// only thing step35's SWA arm keys on, so a chunk-local value here decides the attention
5242    /// kernel from the chunk size (and, below the 512-row window at a nonzero base, drops the
5243    /// window mask entirely). See the batched entry's note in `prime_cache_overlaid`.
5244    #[allow(clippy::too_many_arguments)]
5245    fn step35_prime_batch_layers(
5246        &self,
5247        e: &Engine,
5248        mut x: CudaSlice<f32>,
5249        lo: usize,
5250        hi: usize,
5251        ts: &[usize],
5252        offs: &[usize],
5253        seq_ends: &[usize],
5254        pos_ds: &[CudaSlice<i32>],
5255        caches: &mut [&mut Cache],
5256    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5257        let cfg = &self.cfg;
5258        let n_embd = cfg.n_embd as usize;
5259        let eps = cfg.rms_eps;
5260        let b = ts.len();
5261        let total: usize = ts.iter().sum();
5262        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
5263
5264        let split = |e: &Engine,
5265                     y: &CudaSlice<f32>,
5266                     dim: usize|
5267         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5268            let mut out = Vec::with_capacity(b);
5269            for s in 0..b {
5270                let mut ys = e.uninit(ts[s] * dim)?;
5271                e.copy_view_into(
5272                    &mut ys,
5273                    0,
5274                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5275                    ts[s] * dim,
5276                )?;
5277                out.push(ys);
5278            }
5279            Ok(out)
5280        };
5281
5282        // MEMRA_PRIME_PROF=1: per-phase wall inside the prime, sync-bounded (absolute time
5283        // inflates; the SPLIT is the signal). Two inspection passes failed to find where a
5284        // 3.8 s/4096-token chunk goes against a ~0.55 s compute budget, and nsys cannot capture
5285        // through the server's worker, so the walk measures itself.
5286        let prof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
5287        let mut ph = [0f64; 4]; // 0 norm+qkv, 1 attn, 2 o_proj+norm, 3 moe
5288        let mark = |e: &Engine, acc: usize, t0: &mut std::time::Instant, ph: &mut [f64; 4]| {
5289            if prof {
5290                let _ = e.stream().synchronize();
5291                ph[acc] += t0.elapsed().as_secs_f64() * 1e3;
5292                *t0 = std::time::Instant::now();
5293            }
5294        };
5295        let mut pt = std::time::Instant::now();
5296        for il in lo..hi {
5297            let layer = &self.layers[il];
5298            let Mixer::Full(fa) = &layer.mixer else {
5299                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
5300            };
5301
5302            let mut h = e.uninit(total * n_embd)?;
5303            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5304            if f16fuse {
5305                e.rms_norm_f16out(
5306                    &x,
5307                    layer.attn_norm.float_data(),
5308                    &mut h,
5309                    &mut hx16,
5310                    n_embd,
5311                    total,
5312                    eps,
5313                )?;
5314            } else {
5315                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
5316            }
5317
5318            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
5319            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
5320            // application stay verbatim.
5321            let gate_w = fa
5322                .attn_gate
5323                .as_ref()
5324                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
5325            let mut g4 = if f16fuse {
5326                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
5327            } else {
5328                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
5329            };
5330            let gate = g4.pop().unwrap();
5331            let mut parts: Vec<Vec<CudaSlice<f32>>> =
5332                (0..b).map(|_| Vec::with_capacity(3)).collect();
5333            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
5334                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
5335                    parts[s].push(ys);
5336                }
5337            }
5338            let gates = split(e, &gate, gate_w.out_features())?;
5339            let geometry = self.step35_geom(il);
5340            let hd = geometry.head_dim_k as usize;
5341            let nh = geometry.n_head as usize;
5342            let mut ag_cat = e.uninit(total * nh * hd)?;
5343            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
5344                mark(e, 0, &mut pt, &mut ph);
5345                let ag = self.step35_attn_pre_wo(
5346                    e,
5347                    fa,
5348                    g3s,
5349                    None,
5350                    Some(&gate),
5351                    &pos_ds[s],
5352                    ts[s],
5353                    Some(&mut *caches[s]),
5354                    il,
5355                    seq_ends[s],
5356                )?;
5357                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
5358            }
5359            mark(e, 1, &mut pt, &mut ph);
5360            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
5361
5362            let mut x1 = e.uninit(total * n_embd)?;
5363            let mut z = e.uninit(total * n_embd)?;
5364            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5365            if f16fuse {
5366                e.add_rms_norm_f16out(
5367                    &x,
5368                    &mixed,
5369                    layer.post_attn_norm.float_data(),
5370                    &mut x1,
5371                    &mut z,
5372                    &mut zx16,
5373                    n_embd,
5374                    total,
5375                    eps,
5376                )?;
5377            } else {
5378                e.add(&x, &mixed, &mut x1, total * n_embd)?;
5379                e.rms_norm(
5380                    &x1,
5381                    layer.post_attn_norm.float_data(),
5382                    &mut z,
5383                    n_embd,
5384                    total,
5385                    eps,
5386                )?;
5387            }
5388
5389            mark(e, 2, &mut pt, &mut ph);
5390            let ffn_out = match &layer.ffn {
5391                crate::hybrid::Ffn::Dense {
5392                    ffn_gate,
5393                    ffn_up,
5394                    ffn_down,
5395                } => {
5396                    let n_ff = ffn_gate.out_features();
5397                    let mut g2 = if f16fuse {
5398                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
5399                    } else {
5400                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
5401                    };
5402                    let up = g2.pop().unwrap();
5403                    let gate = g2.pop().unwrap();
5404                    let mut act = e.uninit(total * n_ff)?;
5405                    let d_lim = cfg.clamp_shexp_at(il as u32);
5406                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
5407                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
5408                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
5409                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
5410                            Some(y) => y,
5411                            None => e.matmul(ffn_down, &act, total)?,
5412                        }
5413                    } else {
5414                        Self::ffn_act_lim(
5415                            e,
5416                            cfg,
5417                            &gate,
5418                            &up,
5419                            1.0,
5420                            1.0,
5421                            d_lim,
5422                            &mut act,
5423                            total * n_ff,
5424                        )?;
5425                        e.matmul(ffn_down, &act, total)?
5426                    }
5427                }
5428                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
5429            };
5430            let mut x2 = e.uninit(total * n_embd)?;
5431            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
5432            x = x2;
5433            mark(e, 3, &mut pt, &mut ph);
5434        }
5435        if prof {
5436            eprintln!(
5437                "[prime-prof] t={total} layers={} norm+qkv={:.0}ms attn={:.0}ms o_proj={:.0}ms moe={:.0}ms",
5438                hi - lo,
5439                ph[0],
5440                ph[1],
5441                ph[2],
5442                ph[3]
5443            );
5444        }
5445        Ok(x)
5446    }
5447
5448    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5449    fn step35_prime_batch_epilogue(
5450        &self,
5451        e: &Engine,
5452        x: CudaSlice<f32>,
5453        ts: &[usize],
5454        offs: &[usize],
5455        caches: &mut [&mut Cache],
5456    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5457        let n_embd = self.cfg.n_embd as usize;
5458        let total: usize = ts.iter().sum();
5459        let mut hn = e.uninit(total * n_embd)?;
5460        e.rms_norm(
5461            &x,
5462            self.output_norm.float_data(),
5463            &mut hn,
5464            n_embd,
5465            total,
5466            self.cfg.rms_eps,
5467        )?;
5468
5469        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
5470        let mut out = Vec::with_capacity(ts.len());
5471        for s in 0..ts.len() {
5472            let mut hidden = e.uninit(ts[s] * n_embd)?;
5473            e.copy_view_into(
5474                &mut hidden,
5475                0,
5476                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
5477                ts[s] * n_embd,
5478            )?;
5479            let last0 = (offs[s] + ts[s] - 1) * n_embd;
5480            let mut h_seed = e.uninit(n_embd)?;
5481            e.copy_view_into(
5482                &mut h_seed,
5483                0,
5484                &hidden_src.slice(last0..last0 + n_embd),
5485                n_embd,
5486            )?;
5487            // Exactness-first: the serial reference runs the output head at m=1.
5488            let mut hlast = e.uninit(n_embd)?;
5489            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
5490            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
5491            caches[s].pos += ts[s];
5492            out.push((logits, h_seed, hidden));
5493        }
5494        Ok(out)
5495    }
5496
5497    /// `seq_ends[s]` = sequence s's REQUEST-absolute end position (`cache.pos + prompt_len
5498    /// + queued_after`, computed once before any chunk loop). Only step35's SWA arm reads it,
5499    /// and it must NOT be this chunk's own length: see the note on the batched entry in
5500    /// `prime_cache_overlaid` for the window the chunk-local value opened.
5501    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5502    fn step35_prime_cache_batch(
5503        &self,
5504        e: &Engine,
5505        prompts: &[&[u32]],
5506        caches: &mut [&mut Cache],
5507        seq_ends: &[usize],
5508    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5509        assert_eq!(
5510            seq_ends.len(),
5511            prompts.len(),
5512            "step35 batched prime: one seq_end per sequence"
5513        );
5514        validate_step_prime_batch_modes(
5515            step_tp_prefill_enabled()?,
5516            step_ep_grouped_prefill_enabled()?,
5517        )?;
5518        if crate::pp::pp_host_bounce_active()
5519            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
5520        {
5521            return Err(
5522                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
5523                 stage split; refusing an unsplit remote-weight walk"
5524                    .into(),
5525            );
5526        }
5527        if !Self::step35_prime_batch_on() {
5528            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
5529        }
5530        // Continuation chunks are admitted (positions above carry each sequence's base). The
5531        // remaining restriction is genuine: a CROSS-REQUEST batch mixing sequences at different
5532        // positions still needs per-request queued_after to place its KV, so B > 1 keeps the
5533        // fresh-prompt rule.
5534        if prompts.len() > 1 && caches.iter().any(|c| c.pos != 0) {
5535            return Err(
5536                "step35 batched prime supports continuation only at B=1; a cross-request batch \
5537                 at mixed positions requires per-request queued_after"
5538                    .into(),
5539            );
5540        }
5541
5542        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5543        for &t in &ts {
5544            assert!(
5545                t >= PRIME_MIN_T,
5546                "step35 batched prime needs T >= {PRIME_MIN_T}"
5547            );
5548        }
5549        for (s, c) in caches.iter().enumerate() {
5550            // POS-INCLUSIVE, like the walk's assert: a continuation chunk's rows land at
5551            // c.pos.., so the fresh-only `ts[s] <= max_ctx` form under-checked it.
5552            assert!(
5553                c.pos + ts[s] <= c.max_ctx,
5554                "step35 batched prime exceeds cache max_ctx"
5555            );
5556            assert!(
5557                seq_ends[s] >= c.pos + ts[s],
5558                "step35 batched prime: seq_end must cover this chunk"
5559            );
5560        }
5561        let mut transaction = CacheTaintGuard::arm(caches);
5562        // MEMRA_STEP35_PRIME_BATCH_TSEND=1: CANARY SEAM restoring the pre-fix chunk-local
5563        // `seq_end` (this chunk's own length, which `ts[s]` used to supply here). It is suffix-
5564        // and chunk-VARIANT by construction, so the suffix byte-identity gate MUST break under
5565        // it. That is how the defect is DEMONSTRATED rather than argued: one binary, one seam,
5566        // the legacy arm fails cold-vs-rewound identity and the default arm passes. Read per
5567        // call; never on in a measured default run.
5568        let legacy_tsend = std::env::var("MEMRA_STEP35_PRIME_BATCH_TSEND").as_deref() == Ok("1");
5569        let seq_ends_eff: Vec<usize> = if legacy_tsend {
5570            ts.clone()
5571        } else {
5572            seq_ends.to_vec()
5573        };
5574        let offs: Vec<usize> = ts
5575            .iter()
5576            .scan(0usize, |a, &t| {
5577                let o = *a;
5578                *a += t;
5579                Some(o)
5580            })
5581            .collect();
5582        let total: usize = ts.iter().sum();
5583        let payload = total * self.cfg.n_embd as usize;
5584        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5585        // Positions start at each sequence's CURRENT cache position, not 0, so this entry can
5586        // prime a continuation chunk. The attention core already supports it: step35_attn_pre_wo
5587        // with Some(cache) is PRIME mode — it appends this chunk's post-rope K / raw V and
5588        // attends THROUGH the cache view — so only the hardcoded 0..t and the guard below ever
5589        // restricted it to fresh prompts.
5590        let positions: Vec<Vec<i32>> = ts
5591            .iter()
5592            .zip(caches.iter())
5593            .map(|(&t, c)| {
5594                let base = c.pos as i32;
5595                (0..t as i32).map(|i| base + i).collect()
5596            })
5597            .collect();
5598        let upload_positions =
5599            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
5600                positions
5601                    .iter()
5602                    .map(|p| e.htod_i32(p))
5603                    .collect::<Result<_, _>>()
5604            };
5605
5606        static ONCE: std::sync::Once = std::sync::Once::new();
5607        ONCE.call_once(|| {
5608            eprintln!(
5609                "[step35-prime-batch] first concat prime: B={} tokens={total}",
5610                prompts.len()
5611            );
5612        });
5613
5614        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
5615            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5616                let rt = crate::pp::PpNRt::get(e)?;
5617                let n_st = fence.len() - 1;
5618                assert_eq!(
5619                    rt.n_stages(),
5620                    n_st,
5621                    "step35 prime batch stage count mismatch"
5622                );
5623                let caller_stream = e.stream();
5624                rt.fence_stages_behind(&caller_stream)?;
5625
5626                let mut slot = {
5627                    let _st0 = rt.enter(0);
5628                    let e0 = rt.engine(0, e);
5629                    let pos_ds = upload_positions(e0)?;
5630                    let x = self.embed(e0, &cat_tokens)?;
5631                    let x = self.step35_prime_batch_layers(
5632                        e0,
5633                        x,
5634                        fence[0],
5635                        fence[1],
5636                        &ts,
5637                        &offs,
5638                        &seq_ends_eff,
5639                        &pos_ds,
5640                        caches,
5641                    )?;
5642                    rt.tx(0, &x, payload)?
5643                };
5644                for s in 1..n_st - 1 {
5645                    let _st = rt.enter(s);
5646                    let es = rt.engine(s, e);
5647                    let pos_ds = upload_positions(es)?;
5648                    let x = rt.rx(s - 1, slot, payload)?;
5649                    let x = self.step35_prime_batch_layers(
5650                        es,
5651                        x,
5652                        fence[s],
5653                        fence[s + 1],
5654                        &ts,
5655                        &offs,
5656                        &seq_ends_eff,
5657                        &pos_ds,
5658                        caches,
5659                    )?;
5660                    slot = rt.tx(s, &x, payload)?;
5661                }
5662
5663                let _stl = rt.enter(n_st - 1);
5664                let el = rt.engine(n_st - 1, e);
5665                let pos_ds = upload_positions(el)?;
5666                let x = rt.rx(n_st - 2, slot, payload)?;
5667                let x = self.step35_prime_batch_layers(
5668                    el,
5669                    x,
5670                    fence[n_st - 1],
5671                    fence[n_st],
5672                    &ts,
5673                    &offs,
5674                    &seq_ends_eff,
5675                    &pos_ds,
5676                    caches,
5677                )?;
5678                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
5679                rt.publish_to(n_st - 1, &caller_stream)?;
5680                crate::pp::STEP35_PRIME_BATCH_SPLITS
5681                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5682                out
5683            } else {
5684                let pos_ds = upload_positions(e)?;
5685                let x = self.embed(e, &cat_tokens)?;
5686                let x = self.step35_prime_batch_layers(
5687                    e,
5688                    x,
5689                    0,
5690                    self.layers.len(),
5691                    &ts,
5692                    &offs,
5693                    &seq_ends_eff,
5694                    &pos_ds,
5695                    caches,
5696                )?;
5697                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5698            }
5699        } else {
5700            let pos_ds = upload_positions(e)?;
5701            let x = self.embed(e, &cat_tokens)?;
5702            let x = self.step35_prime_batch_layers(
5703                e,
5704                x,
5705                0,
5706                self.layers.len(),
5707                &ts,
5708                &offs,
5709                &seq_ends_eff,
5710                &pos_ds,
5711                caches,
5712            )?;
5713            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
5714        };
5715        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5716        transaction.commit();
5717        Ok(out)
5718    }
5719
5720    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
5721    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
5722    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
5723    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
5724    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
5725    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
5726    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
5727    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
5728    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
5729    /// over the quantized past; Linear: the stateful pad_view twin — the same state
5730    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
5731    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
5732    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
5733    /// back to single-chunk serving).
5734    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
5735    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
5736    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
5737    pub fn prime_cache_batch(
5738        &self,
5739        e: &Engine,
5740        prompts: &[&[u32]],
5741        caches: &mut [&mut Cache],
5742    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5743        self.refuse_hyper("prime_cache_batch")?;
5744        for cache in caches.iter() {
5745            cache.ensure_usable("prime_cache_batch")?;
5746        }
5747        if crate::pp::pp_cuts(self.layers.len()).is_some()
5748            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
5749        {
5750            return Err("pipeline rewrite is not qualified for batched prime".into());
5751        }
5752        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::CarriedPrime) {
5753            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
5754                return Err("neither batched-prime nor eager rewrite is qualified".into());
5755            }
5756            if prompts.len() != caches.len() {
5757                return Err("prime fallback prompt/cache shape mismatch".into());
5758            }
5759            static ONCE: std::sync::Once = std::sync::Once::new();
5760            ONCE.call_once(|| {
5761                eprintln!(
5762                    "[rewrite] carried-prime.v1 unqualified; using individual native eager primes"
5763                );
5764            });
5765            let mut transaction = CacheTaintGuard::arm(caches);
5766            let result: Result<Vec<_>, Box<dyn std::error::Error>> = prompts
5767                .iter()
5768                .copied()
5769                .zip(caches.iter_mut())
5770                .map(|(prompt, cache)| self.prime_cache(e, prompt, cache, 0))
5771                .collect();
5772            if result.is_ok() {
5773                transaction.commit();
5774            }
5775            return result;
5776        }
5777        let _pp_walk =
5778            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5779                let rt = crate::pp::PpNRt::get(e)?;
5780                Some(rt.acquire_walk("prime_cache_batch")?)
5781            } else {
5782                None
5783            };
5784        let cfg = &self.cfg;
5785        let n_embd = cfg.n_embd as usize;
5786        let eps = cfg.rms_eps;
5787        let b = prompts.len();
5788        assert!(b >= 1 && b == caches.len());
5789        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
5790        let carried = pos0s.iter().any(|&p| p > 0);
5791        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
5792        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
5793        // generic concat attn core below (uniform geometry, no per-layer swa window, no
5794        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
5795        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
5796        if self.uses_gemma_program() {
5797            return Err(
5798                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
5799                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
5800                    .into(),
5801            );
5802        }
5803        // Step35 has a dedicated concat walk: the generic core below cannot express its
5804        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
5805        if self.uses_sliding_gated_moe_program() {
5806            // The cross-request driver hands whole requests (no chunk loop of its own), so each
5807            // sequence's request-absolute end IS its base plus its prompt length — the value
5808            // `ts[s]` happened to equal for the fresh B>=1 batches this caller admits, which is
5809            // why this arm is bit-for-bit unchanged by the seq_end threading.
5810            let seq_ends: Vec<usize> = caches
5811                .iter()
5812                .zip(prompts.iter())
5813                .map(|(c, p)| c.pos + p.len())
5814                .collect();
5815            return self.step35_prime_cache_batch(e, prompts, caches, &seq_ends);
5816        }
5817        if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
5818            let rt = crate::pp::PpNRt::get(e)?;
5819            if rt.cross_device() {
5820                return Err(
5821                    "prime_cache_batch: generic dense concat prime has no cross-device PP split; use individual prime_cache calls"
5822                        .into(),
5823                );
5824            }
5825        }
5826        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
5827        for &t in &ts {
5828            assert!(
5829                t >= PRIME_MIN_T,
5830                "prime_cache_batch needs T >= {PRIME_MIN_T}"
5831            );
5832        }
5833        for (s, c) in caches.iter().enumerate() {
5834            assert!(
5835                c.pos + ts[s] <= c.max_ctx,
5836                "prime_cache_batch: prompt exceeds cache max_ctx"
5837            );
5838        }
5839        let mut transaction = CacheTaintGuard::arm(caches);
5840        let total: usize = ts.iter().sum();
5841        let offs: Vec<usize> = ts
5842            .iter()
5843            .scan(0usize, |a, &t| {
5844                let o = *a;
5845                *a += t;
5846                Some(o)
5847            })
5848            .collect();
5849        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
5850        let pos_ds: Vec<CudaSlice<i32>> = ts
5851            .iter()
5852            .zip(&pos0s)
5853            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
5854            .collect::<Result<_, _>>()?;
5855        // split a concat [total, dim] buffer into per-seq copies
5856        let split = |e: &Engine,
5857                     y: &CudaSlice<f32>,
5858                     dim: usize|
5859         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
5860            let mut out = Vec::with_capacity(b);
5861            for s in 0..b {
5862                let mut ys = e.uninit(ts[s] * dim)?;
5863                e.copy_view_into(
5864                    &mut ys,
5865                    0,
5866                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
5867                    ts[s] * dim,
5868                )?;
5869                out.push(ys);
5870            }
5871            Ok(out)
5872        };
5873
5874        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
5875        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
5876        for (il, layer) in self.layers.iter().enumerate() {
5877            let mut h = e.uninit(total * n_embd)?;
5878            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
5879            e.rms_norm_f16out(
5880                &x,
5881                layer.attn_norm.float_data(),
5882                &mut h,
5883                &mut hx16,
5884                n_embd,
5885                total,
5886                eps,
5887            )?;
5888            // mixer: projection GROUP on the concat (m = total), stateful core per seq
5889            let mut mixed = e.uninit(total * n_embd)?;
5890            match &layer.mixer {
5891                Mixer::Full(fa) => {
5892                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
5893                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
5894                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
5895                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
5896                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
5897                    // back to the per-seq dispatch.
5898                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
5899                    let (n_head, n_head_kv, head_dim) = (
5900                        geometry.n_head as usize,
5901                        geometry.n_head_kv as usize,
5902                        geometry.head_dim_k as usize,
5903                    );
5904                    let fa_scale = geometry.attention_scale();
5905                    let use_favl = !carried
5906                        && (2..=8).contains(&b)
5907                        && (head_dim == 256 || head_dim == 128)
5908                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
5909                        && std::env::var("MEMRA_NOFA").is_err()
5910                        && std::env::var("MEMRA_FA_FLOOR").is_err()
5911                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
5912                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
5913                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
5914                    if use_favl {
5915                        let (qf_w, kf_w, vf_w) = (
5916                            fa.wq.out_features(),
5917                            fa.wk.out_features(),
5918                            fa.wv.out_features(),
5919                        );
5920                        // Same bounds contract as `Engine::q_gate_split`, applied to the varlen
5921                        // twin's PER-TOKEN stride. `attn_pre_vl8` takes raw device pointers so it
5922                        // cannot check its own extents; `qf_w` is the wq out-features that set
5923                        // them, and `q_gate_split_vl` reads 2*head_dim per head out of it.
5924                        memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
5925                        struct APre {
5926                            q: CudaSlice<f32>,
5927                            gate: Option<CudaSlice<f32>>,
5928                            qn: CudaSlice<f32>,
5929                            kn: CudaSlice<f32>,
5930                        }
5931                        let mut aps = Vec::with_capacity(b);
5932                        for &t in ts.iter().take(b) {
5933                            aps.push(APre {
5934                                q: e.uninit(t * n_head * head_dim)?,
5935                                gate: Some(e.uninit(t * n_head * head_dim)?),
5936                                qn: e.uninit(t * n_head * head_dim)?,
5937                                kn: e.uninit(t * n_head_kv * head_dim)?,
5938                            });
5939                        }
5940                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
5941                            let kvl = caches[0].kv[il].as_ref().unwrap();
5942                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
5943                        };
5944                        let pargs: Vec<crate::AttnPreVl> = (0..b)
5945                            .map(|s| {
5946                                let (o, t) = (offs[s], ts[s]);
5947                                let kvl = caches[s].kv[il].as_ref().unwrap();
5948                                assert!(
5949                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
5950                                    "prime_cache_batch attn vl: fresh + capacity"
5951                                );
5952                                crate::AttnPreVl {
5953                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
5954                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
5955                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
5956                                    q: e.addr_f32(&aps[s].q),
5957                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
5958                                    qn: e.addr_f32(&aps[s].qn),
5959                                    kn: e.addr_f32(&aps[s].kn),
5960                                    kc: e.addr_u8(&kvl.k),
5961                                    vc: e.addr_u8(&kvl.v),
5962                                    t: t as i32,
5963                                    pad: 0,
5964                                }
5965                            })
5966                            .collect();
5967                        e.attn_pre_vl8(
5968                            &pargs,
5969                            fa.q_norm.float_data(),
5970                            fa.k_norm.float_data(),
5971                            head_dim,
5972                            geometry.n_rot as usize,
5973                            n_head,
5974                            n_head_kv,
5975                            self.cfg.rms_eps,
5976                            geometry.rope_base,
5977                            1.0,
5978                            kv_dim_k,
5979                            kv_dim_v,
5980                            ktb,
5981                            vtb,
5982                        )?;
5983                        for s in 0..b {
5984                            let kvl = caches[s].kv[il].as_mut().unwrap();
5985                            kvl.len += ts[s];
5986                            let new_len = kvl.len as i32;
5987                            e.set_i32_one(&mut kvl.len_d, new_len)?;
5988                        }
5989                        let mut attns = Vec::with_capacity(b);
5990                        let mut mirrors = Vec::with_capacity(b);
5991                        for &t in ts.iter().take(b) {
5992                            attns.push(e.uninit(t * n_head * head_dim)?);
5993                            let n = t * n_head_kv * head_dim;
5994                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
5995                        }
5996                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
5997                        // promoted single-seq config is on; else the mma favl.
5998                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
5999                            Ok("0") => false,
6000                            // Same refusal as the single-seq twin (lib.rs fa_prefill): the
6001                            // batched bf16 stage reaches func("f32_to_bf16_bulk"), absent on a
6002                            // portable build.
6003                            Ok("1") => {
6004                                crate::refuse_portable_force(
6005                                    "MEMRA_FA3=1",
6006                                    "the sm_90a fa3/bf16 kernels",
6007                                );
6008                                true
6009                            }
6010                            _ => cfg!(memra_hopper_mma),
6011                        };
6012                        if fa3_on {
6013                            let mut q16s = Vec::with_capacity(b);
6014                            let mut v16s = Vec::with_capacity(b);
6015                            for s in 0..b {
6016                                let t = ts[s];
6017                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
6018                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
6019                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6020                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
6021                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
6022                                e.f32_to_bf16_v(
6023                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
6024                                    &mut v16,
6025                                    t * n_head_kv * head_dim,
6026                                )?;
6027                                q16s.push(q16);
6028                                v16s.push((k16, v16));
6029                            }
6030                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
6031                            let mut kp = qp;
6032                            let mut vp = qp;
6033                            let mut op = [core::ptr::null_mut::<f32>(); 8];
6034                            let mut tsv = [0i32; 8];
6035                            for s in 0..b {
6036                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
6037                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
6038                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
6039                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
6040                                tsv[s] = ts[s] as i32;
6041                            }
6042                            let rc = unsafe {
6043                                crate::fa3_vl_raw(
6044                                    qp.as_ptr(),
6045                                    kp.as_ptr(),
6046                                    vp.as_ptr(),
6047                                    op.as_ptr(),
6048                                    tsv.as_ptr(),
6049                                    b as i32,
6050                                    n_head as i32,
6051                                    n_head_kv as i32,
6052                                    head_dim as i32,
6053                                    fa_scale,
6054                                    e.stream().cu_stream() as *mut core::ffi::c_void,
6055                                )
6056                            };
6057                            if rc != 0 {
6058                                return Err(format!("memra_fa3_vl rc={rc}").into());
6059                            }
6060                        } else {
6061                            let fargs: Vec<crate::FaSeqVl> = (0..b)
6062                                .map(|s| crate::FaSeqVl {
6063                                    q: e.addr_f32(&aps[s].qn),
6064                                    k16: e.addr_u8(&mirrors[s].0),
6065                                    v16: e.addr_u8(&mirrors[s].1),
6066                                    o: e.addr_f32(&attns[s]),
6067                                    kf: e.addr_f32(&aps[s].kn),
6068                                    vf: e.addr_f32v(
6069                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
6070                                    ),
6071                                    t: ts[s] as i32,
6072                                    pad: 0,
6073                                })
6074                                .collect();
6075                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
6076                        }
6077                        for (s, attn) in attns.into_iter().enumerate() {
6078                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
6079                                e,
6080                                attn,
6081                                &aps[s].gate,
6082                                ts[s],
6083                                n_head,
6084                                head_dim,
6085                            )?;
6086                            let mut done = false;
6087                            if let Some(xh) = &ag16 {
6088                                done = e.try_f16_gemm_pre_into_off(
6089                                    &fa.wo,
6090                                    xh,
6091                                    ts[s],
6092                                    &mut mixed,
6093                                    offs[s] * n_embd,
6094                                )?;
6095                            }
6096                            if !done {
6097                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6098                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6099                            }
6100                        }
6101                    } else {
6102                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
6103                            (0..b).map(|_| Vec::new()).collect();
6104                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
6105                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
6106                                parts[s].push(ys);
6107                            }
6108                        }
6109                        for (s, g3s) in parts.into_iter().enumerate() {
6110                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
6111                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
6112                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
6113                            )?;
6114                            let mut done = false;
6115                            if let Some(xh) = &ag16 {
6116                                done = e.try_f16_gemm_pre_into_off(
6117                                    &fa.wo,
6118                                    xh,
6119                                    ts[s],
6120                                    &mut mixed,
6121                                    offs[s] * n_embd,
6122                                )?;
6123                            }
6124                            if !done {
6125                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
6126                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
6127                            }
6128                        }
6129                    }
6130                }
6131                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched cache prime"),
6132                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched prime"),
6133                Mixer::Linear(la) => {
6134                    // task #16: NO split copies (cores read row-offset views of the concat
6135                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
6136                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
6137                    // varlen K5 launch for all sequences.
6138                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6139                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
6140                    let outs =
6141                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
6142                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
6143                        let (o, t) = (offs[s], ts[s]);
6144                        let mut done = false;
6145                        if let Some(xh) = &gn16 {
6146                            done = e.try_f16_gemm_pre_into_off(
6147                                &la.ssm_out,
6148                                xh,
6149                                t,
6150                                &mut mixed,
6151                                o * n_embd,
6152                            )?;
6153                        }
6154                        if !done {
6155                            let m = e.matmul(&la.ssm_out, &gn, t)?;
6156                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
6157                        }
6158                    }
6159                }
6160            }
6161            let mut x1 = e.uninit(total * n_embd)?;
6162            let mut z = e.uninit(total * n_embd)?;
6163            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
6164            e.add_rms_norm_f16out(
6165                &x,
6166                &mixed,
6167                layer.post_attn_norm.float_data(),
6168                &mut x1,
6169                &mut z,
6170                &mut zx16,
6171                n_embd,
6172                total,
6173                eps,
6174            )?;
6175            let ffn_out = match &layer.ffn {
6176                crate::hybrid::Ffn::Dense {
6177                    ffn_gate,
6178                    ffn_up,
6179                    ffn_down,
6180                } => {
6181                    let n_ff = ffn_gate.out_features();
6182                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
6183                    let up = g2.pop().unwrap();
6184                    let gate = g2.pop().unwrap();
6185                    let mut act = e.uninit(total * n_ff)?;
6186                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
6187                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
6188                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
6189                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
6190                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
6191                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
6192                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
6193                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
6194                            Some(y) => y,
6195                            None => e.matmul(ffn_down, &act, total)?,
6196                        }
6197                    } else {
6198                        Self::ffn_act_lim(
6199                            e,
6200                            &self.cfg,
6201                            &gate,
6202                            &up,
6203                            1.0,
6204                            1.0,
6205                            d_lim,
6206                            &mut act,
6207                            total * n_ff,
6208                        )?;
6209                        e.matmul(ffn_down, &act, total)?
6210                    }
6211                }
6212                crate::hybrid::Ffn::Moe(m) => {
6213                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
6214                }
6215            };
6216            let mut x2 = e.uninit(total * n_embd)?;
6217            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
6218            x = x2;
6219        }
6220        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
6221        let mut hn = e.uninit(total * n_embd)?;
6222        e.rms_norm(
6223            &x,
6224            self.output_norm.float_data(),
6225            &mut hn,
6226            n_embd,
6227            total,
6228            eps,
6229        )?;
6230        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
6231        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
6232        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
6233        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
6234        // argmax battery arbitrates, same as every other prefill GEMM change.
6235        let mut hcat = e.uninit(b * n_embd)?;
6236        for s in 0..b {
6237            let last0 = (offs[s] + ts[s] - 1) * n_embd;
6238            e.copy_view_into(
6239                &mut hcat,
6240                s * n_embd,
6241                &hn.slice(last0..last0 + n_embd),
6242                n_embd,
6243            )?;
6244        }
6245        let logits_cat = if b >= 2 {
6246            e.try_f16_gemm(&self.output, &hcat, b)?
6247        } else {
6248            None
6249        };
6250        let logits_host: Option<Vec<f32>> = match &logits_cat {
6251            Some(lc) => Some(e.dtoh(lc)?),
6252            None => None,
6253        };
6254        let n_vocab = self.output.out_features();
6255        let mut hidden_all = if crate::spec::spec_hpost() {
6256            split(e, &hn, n_embd)?
6257        } else {
6258            split(e, &x, n_embd)?
6259        };
6260        let mut out = Vec::with_capacity(b);
6261        for s in 0..b {
6262            let last0 = (offs[s] + ts[s] - 1) * n_embd;
6263            let mut h_seed = e.uninit(n_embd)?;
6264            if !crate::spec::spec_hpost() {
6265                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
6266            } else {
6267                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6268            }
6269            let logits = match &logits_host {
6270                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
6271                None => {
6272                    let mut hlast = e.uninit(n_embd)?;
6273                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
6274                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
6275                }
6276            };
6277            caches[s].pos += ts[s];
6278            out.push((logits, h_seed, hidden_all.remove(0)));
6279        }
6280        transaction.commit();
6281        Ok(out)
6282    }
6283
6284    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
6285    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
6286    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
6287    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
6288    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
6289    ///
6290    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
6291    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
6292    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
6293    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
6294    #[allow(clippy::too_many_arguments)]
6295    fn full_attn_prime(
6296        &self,
6297        e: &Engine,
6298        fa: &FullAttnLayer,
6299        h: &CudaSlice<f32>,
6300        hx: Option<&CudaSlice<u8>>,
6301        pos_d: &CudaSlice<i32>,
6302        t: usize,
6303        cache: &mut Cache,
6304        il: usize,
6305        seq_end: usize,
6306    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6307        if self.uses_sliding_gated_moe_program() {
6308            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
6309        }
6310        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
6311        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
6312        // this single-seq path composes proj+core identically (byte-for-byte the old body).
6313        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
6314        let g3 = match hx {
6315            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
6316            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
6317        };
6318        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
6319    }
6320
6321    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
6322    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
6323    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
6324    #[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
6325    fn full_attn_prime_core(
6326        &self,
6327        e: &Engine,
6328        fa: &FullAttnLayer,
6329        g3: Vec<CudaSlice<f32>>,
6330        pos_d: &CudaSlice<i32>,
6331        t: usize,
6332        cache: &mut Cache,
6333        il: usize,
6334    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6335        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
6336        if let Some(xh) = &ag16
6337            && let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)?
6338        {
6339            return Ok(y);
6340        }
6341        e.matmul(&fa.wo, &attn_g, t)
6342    }
6343
6344    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6345    #[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
6346    fn full_attn_prime_core_inner(
6347        &self,
6348        e: &Engine,
6349        fa: &FullAttnLayer,
6350        g3: Vec<CudaSlice<f32>>,
6351        pos_d: &CudaSlice<i32>,
6352        t: usize,
6353        cache: &mut Cache,
6354        il: usize,
6355    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6356        let cfg = &self.cfg;
6357        let geometry = cfg.full_attention_geometry_at(il as u32);
6358        let n_head = geometry.n_head as usize;
6359        let n_head_kv = geometry.n_head_kv as usize;
6360        let head_dim = geometry.head_dim_k as usize;
6361        let scale = geometry.attention_scale();
6362        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
6363        let AttnPre { q, k, v, gate } = pre;
6364        let mut attn = e.uninit(t * n_head * head_dim)?;
6365        self.full_attn_prime_fa_dispatch(
6366            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
6367        )?;
6368        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
6369    }
6370
6371    /// task #18 (attn side): projections tail through KV append — everything before the
6372    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
6373    /// present BEFORE this chunk's append (base_len; 0 == fresh).
6374    #[allow(clippy::type_complexity)]
6375    #[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
6376    fn full_attn_prime_pre_fa(
6377        &self,
6378        e: &Engine,
6379        fa: &FullAttnLayer,
6380        mut g3: Vec<CudaSlice<f32>>,
6381        pos_d: &CudaSlice<i32>,
6382        t: usize,
6383        cache: &mut Cache,
6384        il: usize,
6385    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
6386        let cfg = &self.cfg;
6387        let geometry = cfg.full_attention_geometry_at(il as u32);
6388        let n_head = geometry.n_head as usize;
6389        let n_head_kv = geometry.n_head_kv as usize;
6390        let head_dim = geometry.head_dim_k as usize;
6391        let eps = cfg.rms_eps;
6392
6393        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
6394        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
6395        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
6396        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6397        let v = g3.pop().unwrap();
6398        let mut k = g3.pop().unwrap();
6399        let qf = g3.pop().unwrap();
6400        let (mut q, gate) = if gated {
6401            let mut q = e.uninit(t * n_head * head_dim)?;
6402            let mut gate = e.uninit(t * n_head * head_dim)?;
6403            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
6404            (q, Some(gate))
6405        } else {
6406            (qf, None)
6407        };
6408
6409        let mut qn = e.uninit(t * n_head * head_dim)?;
6410        e.rms_norm(
6411            &q,
6412            fa.q_norm.float_data(),
6413            &mut qn,
6414            head_dim,
6415            n_head * t,
6416            eps,
6417        )?;
6418        q = qn;
6419        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6420        e.rms_norm(
6421            &k,
6422            fa.k_norm.float_data(),
6423            &mut kn,
6424            head_dim,
6425            n_head_kv * t,
6426            eps,
6427        )?;
6428        k = kn;
6429        let rope_dims = geometry.n_rot as usize;
6430        e.rope_neox(
6431            &mut q,
6432            pos_d,
6433            head_dim,
6434            rope_dims,
6435            n_head,
6436            t,
6437            geometry.rope_base,
6438            1.0,
6439        )?;
6440        e.rope_neox(
6441            &mut k,
6442            pos_d,
6443            head_dim,
6444            rope_dims,
6445            n_head_kv,
6446            t,
6447            geometry.rope_base,
6448            1.0,
6449        )?;
6450
6451        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
6452        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
6453        {
6454            let kvl = cache.kv[il].as_mut().unwrap();
6455            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
6456            e.append_kv_quantized_rows(
6457                &k,
6458                &v,
6459                &mut kvl.k,
6460                &mut kvl.v,
6461                kvl.len,
6462                t,
6463                kvl.kv_dim_k,
6464                kvl.kv_dim_v,
6465                kvl.k_tok_bytes,
6466                kvl.v_tok_bytes,
6467                crate::Engine::kv_fp8_on(),
6468            )?;
6469            kvl.len += t;
6470            let new_len = kvl.len as i32;
6471            e.set_i32_one(&mut kvl.len_d, new_len)?;
6472        }
6473
6474        let base_len = {
6475            let kvl = cache.kv[il].as_ref().unwrap();
6476            kvl.len - t // KV rows present BEFORE this chunk's append above
6477        };
6478        Ok((AttnPre { q, k, v, gate }, base_len))
6479    }
6480
6481    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
6482    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
6483    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
6484    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
6485    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
6486    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
6487    #[allow(clippy::too_many_arguments)]
6488    fn full_attn_prime_fa_dispatch(
6489        &self,
6490        e: &Engine,
6491        q: &CudaSlice<f32>,
6492        k: &CudaSlice<f32>,
6493        v: &CudaSlice<f32>,
6494        attn: &mut CudaSlice<f32>,
6495        base_len: usize,
6496        t: usize,
6497        cache: &mut Cache,
6498        il: usize,
6499        head_dim: usize,
6500        n_head: usize,
6501        n_head_kv: usize,
6502        scale: f32,
6503    ) -> Result<(), Box<dyn std::error::Error>> {
6504        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
6505        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
6506        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
6507        // attend through the quantized cache exactly like every later chunk (quantize-then-
6508        // attend). One numeric class for every row => the chunk size cannot decide where a
6509        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
6510        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
6511        // pin-the-boundary approach).
6512        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
6513        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
6514        // with the fix unconditional, only re-introducing the class edge can prove the gate
6515        // still detects the mechanism. Never on in a measured default run.
6516        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
6517            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6518                e.sdpa_naive(
6519                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6520                )?;
6521            } else {
6522                e.fa_prefill(
6523                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
6524                )?;
6525            }
6526            return Ok(());
6527        }
6528        let kvl = cache.kv[il].as_ref().unwrap();
6529        let t_kv = base_len + t;
6530        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6531        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6532        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
6533        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
6534        // same numeric class, so the uniform contract holds on the fallback too.
6535        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
6536            e.sdpa_naive_quantized_view(
6537                q,
6538                &k_view,
6539                &v_view,
6540                attn,
6541                head_dim,
6542                n_head,
6543                n_head_kv,
6544                t,
6545                t_kv,
6546                scale,
6547                true,
6548                kvl.k_tok_bytes,
6549                kvl.v_tok_bytes,
6550            )?;
6551            return Ok(());
6552        }
6553        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
6554        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
6555        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
6556        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
6557        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
6558        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
6559        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
6560        let deqw = std::env::var("MEMRA_PRIME_DEQW")
6561            .map(|v| v != "0")
6562            .unwrap_or(true);
6563        if deqw {
6564            e.fa_prefill_view_ws(
6565                q,
6566                &k_view,
6567                &v_view,
6568                attn,
6569                head_dim,
6570                n_head,
6571                n_head_kv,
6572                t,
6573                t_kv,
6574                scale,
6575                true,
6576                kvl.k_tok_bytes,
6577                kvl.v_tok_bytes,
6578                crate::Engine::kv_fp8_on(),
6579            )?;
6580        } else {
6581            e.fa_prefill_view(
6582                q,
6583                &k_view,
6584                &v_view,
6585                attn,
6586                head_dim,
6587                n_head,
6588                n_head_kv,
6589                t,
6590                t_kv,
6591                scale,
6592                true,
6593                kvl.k_tok_bytes,
6594                kvl.v_tok_bytes,
6595                crate::Engine::kv_fp8_on(),
6596            )?;
6597        }
6598        Ok(())
6599    }
6600
6601    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
6602    /// (bit-identical composition) and hands wo its fp16 operand directly.
6603    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6604    fn full_attn_prime_post_fa(
6605        &self,
6606        e: &Engine,
6607        attn: CudaSlice<f32>,
6608        gate: &Option<CudaSlice<f32>>,
6609        t: usize,
6610        n_head: usize,
6611        head_dim: usize,
6612    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6613        let (attn_g, ag16) = match gate {
6614            Some(gate) => {
6615                let n = t * n_head * head_dim;
6616                let mut ag = e.uninit(n)?;
6617                if Self::f16out_on(e, t) {
6618                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
6619                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
6620                    (ag, Some(a16))
6621                } else {
6622                    let mut gsig = e.uninit(n)?;
6623                    e.sigmoid(gate, &mut gsig, n)?;
6624                    e.mul(&attn, &gsig, &mut ag, n)?;
6625                    (ag, None)
6626                }
6627            }
6628            None => (attn, None),
6629        };
6630        Ok((attn_g, ag16))
6631    }
6632
6633    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
6634    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
6635    /// carried THROUGH the cache like the spec verify does: carried-ring conv
6636    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
6637    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
6638    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
6639    #[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
6640    fn linear_attn_prime(
6641        &self,
6642        e: &Engine,
6643        la: &LinearAttnLayer,
6644        h: &CudaSlice<f32>,
6645        hx: Option<&CudaSlice<u8>>,
6646        t: usize,
6647        cache: &mut Cache,
6648        il: usize,
6649    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6650        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
6651        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
6652        let g4 = match hx {
6653            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
6654            None => e.matmul_group(&ws, h, t)?,
6655        };
6656        self.linear_attn_prime_core(e, la, g4, t, cache, il)
6657    }
6658
6659    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
6660    fn linear_attn_prime_core(
6661        &self,
6662        e: &Engine,
6663        la: &LinearAttnLayer,
6664        mut g4: Vec<CudaSlice<f32>>,
6665        t: usize,
6666        cache: &mut Cache,
6667        il: usize,
6668    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6669        self.linear_attn_prime_core_pad(e, la, std::mem::take(&mut g4), t, cache, il, None)
6670    }
6671
6672    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
6673    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
6674    /// conv ring writes back from the true tail. None = classic path, byte-identical.
6675    #[allow(clippy::too_many_arguments)]
6676    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6677    fn linear_attn_prime_core_pad_inner(
6678        &self,
6679        e: &Engine,
6680        la: &LinearAttnLayer,
6681        mut g4: Vec<CudaSlice<f32>>,
6682        t: usize,
6683        cache: &mut Cache,
6684        il: usize,
6685        pad_len: Option<&CudaSlice<i32>>,
6686    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
6687        // shim over the view twin (task #16): full-range views of the owned buffers.
6688        let geometry = la.geometry;
6689        let d_state = geometry.key_head_dim as usize;
6690        let num_k = geometry.key_heads as usize;
6691        let num_v = geometry.value_heads as usize;
6692        let key_dim = d_state * num_k;
6693        let value_dim = geometry.value_head_dim as usize * num_v;
6694        let conv_dim = key_dim * 2 + value_dim;
6695        let alpha = g4.pop().unwrap(); // [T, num_v]
6696        let beta_raw = g4.pop().unwrap(); // [T, num_v]
6697        let z = g4.pop().unwrap(); // [T, value_dim]
6698        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
6699        self.linear_attn_prime_core_pad_view(
6700            e,
6701            la,
6702            &qkv_mixed.slice(0..t * conv_dim),
6703            &z.slice(0..t * value_dim),
6704            &beta_raw.slice(0..t * num_v),
6705            &alpha.slice(0..t * num_v),
6706            t,
6707            cache,
6708            il,
6709            pad_len,
6710        )
6711    }
6712
6713    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
6714    /// shared verbatim by the per-seq scan path and the varlen batched path.
6715    #[allow(clippy::too_many_arguments)]
6716    fn linear_attn_gdn_prep(
6717        &self,
6718        e: &Engine,
6719        la: &LinearAttnLayer,
6720        qkv_mixed: &cudarc::driver::CudaView<f32>,
6721        beta_raw: &cudarc::driver::CudaView<f32>,
6722        alpha: &cudarc::driver::CudaView<f32>,
6723        t: usize,
6724        cache: &mut Cache,
6725        il: usize,
6726        pad_len: Option<&CudaSlice<i32>>,
6727    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
6728        let cfg = &self.cfg;
6729        let geometry = la.geometry;
6730        let d_state = geometry.key_head_dim as usize;
6731        let num_k = geometry.key_heads as usize;
6732        let num_v = geometry.value_heads as usize;
6733        let d_conv = geometry.conv_kernel as usize;
6734        let key_dim = d_state * num_k; // 2048
6735        let value_dim = geometry.value_head_dim as usize * num_v;
6736        let conv_dim = key_dim * 2 + value_dim; // 8192
6737        let eps = cfg.rms_eps;
6738        debug_assert!(
6739            t >= d_conv - 1,
6740            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
6741        );
6742
6743        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
6744        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
6745        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
6746        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
6747        let rl = cache.recur[il].as_mut().unwrap();
6748        let hk = Self::gdn_hk(e, t, num_v, num_k);
6749        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
6750        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
6751        let mut q_g = e.uninit(d_state * hk * t)?;
6752        let mut k_g = e.uninit(d_state * hk * t)?;
6753        let mut v_g = e.uninit(d_state * num_v * t)?;
6754        if conv_fuse {
6755            e.ssm_conv1d_gdn_state_pad(
6756                qkv_mixed,
6757                &mut rl.conv_state,
6758                la.ssm_conv1d.float_data(),
6759                &mut q_g,
6760                &mut k_g,
6761                &mut v_g,
6762                conv_dim,
6763                t,
6764                d_conv,
6765                d_state,
6766                num_v,
6767                num_k,
6768                key_dim,
6769                hk,
6770                pad_len,
6771            )?;
6772        } else {
6773            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
6774            e.ssm_conv1d_tm_state_pad_v(
6775                qkv_mixed,
6776                &mut rl.conv_state,
6777                la.ssm_conv1d.float_data(),
6778                &mut conv_out,
6779                conv_dim,
6780                t,
6781                d_conv,
6782                pad_len,
6783            )?;
6784            e.qkv_to_gdn_repack(
6785                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
6786            )?;
6787        }
6788        let mut q_l2 = e.uninit(d_state * hk * t)?;
6789        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
6790        // Emitted only where a consumer exists (the wgmma config) — on other arches the
6791        // alloc + epilogue stores would be pure waste.
6792        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
6793            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6794            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
6795            Some(qb)
6796        } else {
6797            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
6798            None
6799        };
6800        let mut k_l2 = e.uninit(d_state * hk * t)?;
6801        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
6802        let kb16 = if Engine::l2_v2_on(d_state) {
6803            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
6804            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
6805            Some(kb)
6806        } else {
6807            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
6808            None
6809        };
6810        let mut beta = e.uninit(t * num_v)?;
6811        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
6812        let mut g_log = e.uninit(t * num_v)?;
6813        e.gdn_glog_v(
6814            alpha,
6815            la.ssm_dt.float_data(),
6816            la.ssm_a.float_data(),
6817            &mut g_log,
6818            num_v,
6819            t,
6820        )?;
6821        if let Some(len_d) = pad_len {
6822            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
6823        }
6824        Ok(GdnPrep {
6825            hk,
6826            q_l2,
6827            k_l2,
6828            v_g,
6829            beta,
6830            g_log,
6831            kb16,
6832            qb16,
6833        })
6834    }
6835
6836    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
6837    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
6838    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
6839    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
6840    #[allow(clippy::too_many_arguments)]
6841    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
6842    fn linear_attn_prime_core_batch(
6843        &self,
6844        e: &Engine,
6845        la: &LinearAttnLayer,
6846        g4: &[CudaSlice<f32>],
6847        offs: &[usize],
6848        ts: &[usize],
6849        caches: &mut [&mut Cache],
6850        il: usize,
6851    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
6852        let geometry = la.geometry;
6853        let d_state = geometry.key_head_dim as usize;
6854        let num_k = geometry.key_heads as usize;
6855        let num_v = geometry.value_heads as usize;
6856        let d_conv = geometry.conv_kernel as usize;
6857        let key_dim = d_state * num_k;
6858        let value_dim = geometry.value_head_dim as usize * num_v;
6859        let conv_dim = key_dim * 2 + value_dim;
6860        let eps = self.cfg.rms_eps;
6861        let scale = 1.0 / (d_state as f32).sqrt();
6862        let b = ts.len();
6863        let c = Engine::gdn_chunk_size();
6864        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
6865        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
6866        let carried = caches.iter().any(|c| c.pos > 0);
6867        let use_vl = !carried
6868            && (2..=8).contains(&b)
6869            && Engine::gdn_chunked_enabled()
6870            && ts.iter().all(|&t| t >= 16)
6871            && e.gdn_mma_enabled(c)
6872            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
6873        if !use_vl {
6874            return (0..b)
6875                .map(|s| {
6876                    let (o, t) = (offs[s], ts[s]);
6877                    self.linear_attn_prime_core_pad_view(
6878                        e,
6879                        la,
6880                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
6881                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
6882                        &g4[2].slice(o * num_v..(o + t) * num_v),
6883                        &g4[3].slice(o * num_v..(o + t) * num_v),
6884                        t,
6885                        caches[s],
6886                        il,
6887                        None,
6888                    )
6889                })
6890                .collect();
6891        }
6892        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
6893        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
6894        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
6895        struct SeqBufs {
6896            conv_out: CudaSlice<f32>,
6897            q_g: CudaSlice<f32>,
6898            k_g: CudaSlice<f32>,
6899            v_g: CudaSlice<f32>,
6900            q_l2: CudaSlice<f32>,
6901            k_l2: CudaSlice<f32>,
6902            beta: CudaSlice<f32>,
6903            g_log: CudaSlice<f32>,
6904            gn: CudaSlice<f32>,
6905            gn16: CudaSlice<u8>,
6906        }
6907        let f16o = Self::f16out_on(e, 16);
6908        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
6909        let mut sb = Vec::with_capacity(b);
6910        let mut pres = Vec::with_capacity(b);
6911        for &t in ts.iter().take(b) {
6912            sb.push(SeqBufs {
6913                conv_out: e.uninit(conv_dim * t)?,
6914                q_g: e.uninit(d_state * hk * t)?,
6915                k_g: e.uninit(d_state * hk * t)?,
6916                v_g: e.uninit(d_state * num_v * t)?,
6917                q_l2: e.uninit(d_state * hk * t)?,
6918                k_l2: e.uninit(d_state * hk * t)?,
6919                beta: e.uninit(t * num_v)?,
6920                g_log: e.uninit(t * num_v)?,
6921                gn: e.uninit(d_state * num_v * t)?,
6922                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
6923            });
6924            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
6925        }
6926        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
6927            .map(|s| {
6928                let (o, t) = (offs[s], ts[s]);
6929                let rl = caches[s].recur[il].as_ref().unwrap();
6930                crate::GdnPrepVl {
6931                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
6932                    conv_state: e.addr_f32(&rl.conv_state),
6933                    conv_out: e.addr_f32(&sb[s].conv_out),
6934                    q_g: e.addr_f32(&sb[s].q_g),
6935                    k_g: e.addr_f32(&sb[s].k_g),
6936                    v_g: e.addr_f32(&sb[s].v_g),
6937                    q_l2: e.addr_f32(&sb[s].q_l2),
6938                    k_l2: e.addr_f32(&sb[s].k_l2),
6939                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
6940                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
6941                    beta: e.addr_f32(&sb[s].beta),
6942                    g_log: e.addr_f32(&sb[s].g_log),
6943                    o: e.addr_f32(&pres[s].o),
6944                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
6945                    gn: e.addr_f32(&sb[s].gn),
6946                    gn16: e.addr_u8(&sb[s].gn16),
6947                    kb16: if Engine::l2_v2_on(d_state) {
6948                        e.addr_u8(&pres[s].kb16)
6949                    } else {
6950                        0
6951                    },
6952                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
6953                        e.addr_u8(&pres[s].qb16)
6954                    } else {
6955                        0
6956                    },
6957                    t: t as i32,
6958                    pad: 0,
6959                }
6960            })
6961            .collect();
6962        let args: Vec<crate::GdnSeqVl> = (0..b)
6963            .map(|s| {
6964                let rl = caches[s].recur[il].as_ref().unwrap();
6965                crate::GdnSeqVl {
6966                    kb16: e.addr_u8(&pres[s].kb16),
6967                    gcum: e.addr_f32(&pres[s].gcum),
6968                    beta: e.addr_f32(&sb[s].beta),
6969                    u: e.addr_f32(&pres[s].u),
6970                    wb16: e.addr_u8(&pres[s].wb16),
6971                    y: e.addr_u8(&pres[s].y16),
6972                    ssnap: e.addr_u8(&pres[s].ssnap16),
6973                    state_in: e.addr_f32(&rl.ssm_state),
6974                    state_out: e.addr_f32(&rl.ssm_state_alt),
6975                    q: e.addr_f32(&sb[s].q_l2),
6976                    p: e.addr_f32(&pres[s].p),
6977                    o: e.addr_f32(&pres[s].o),
6978                    k: e.addr_f32(&sb[s].k_l2),
6979                    v: e.addr_f32(&sb[s].v_g),
6980                    g: e.addr_f32(&sb[s].g_log),
6981                    a: e.addr_f32(&pres[s].a),
6982                    w: e.addr_f32(&pres[s].w),
6983                    t: ts[s] as i32,
6984                    nc: pres[s].nc as i32,
6985                }
6986            })
6987            .collect();
6988        e.gdn_prep_vl8(
6989            &prep_args,
6990            la.ssm_conv1d.float_data(),
6991            la.ssm_dt.float_data(),
6992            la.ssm_a.float_data(),
6993            conv_dim,
6994            d_conv,
6995            d_state,
6996            num_v,
6997            num_k,
6998            key_dim,
6999            hk,
7000            eps,
7001        )?;
7002        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
7003        // both standalone mirror launches vanish on the default config.
7004        if !Engine::l2_v2_on(d_state) {
7005            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
7006        }
7007        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
7008        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
7009            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
7010            if !Engine::l2_v2_on(d_state) {
7011                for s in 0..b {
7012                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
7013                }
7014            }
7015            let mut wa = [crate::GdnWVl::default(); 8];
7016            for s in 0..b {
7017                wa[s] = crate::GdnWVl {
7018                    qb16: e.addr_u8(&pres[s].qb16),
7019                    pb16: e.addr_u8(&pres[s].pb16),
7020                };
7021            }
7022            Some(crate::GdnWVl8(wa))
7023        } else {
7024            None
7025        };
7026        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
7027        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
7028        if f16o {
7029            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
7030        }
7031        // per-seq state swap (+ non-f16out tail fallback)
7032        let mut out = Vec::with_capacity(b);
7033        for (s, bufs) in sb.into_iter().enumerate() {
7034            let rl = caches[s].recur[il].as_mut().unwrap();
7035            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7036            let (o, t) = (offs[s], ts[s]);
7037            let SeqBufs { mut gn, gn16, .. } = bufs;
7038            if f16o {
7039                out.push((gn, Some(gn16)));
7040            } else {
7041                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
7042                e.gated_rmsnorm_zv(
7043                    &pres[s].o,
7044                    la.ssm_norm.float_data(),
7045                    &z_v,
7046                    &mut gn,
7047                    d_state,
7048                    num_v * t,
7049                    eps,
7050                )?;
7051                out.push((gn, None));
7052            }
7053        }
7054        Ok(out)
7055    }
7056
7057    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
7058    /// views of the CONCAT projection outputs directly (no per-seq split copies).
7059    /// Same kernels, same values, byte-identical to the Vec shim above.
7060    #[allow(clippy::too_many_arguments)]
7061    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
7062    fn linear_attn_prime_core_pad_view(
7063        &self,
7064        e: &Engine,
7065        la: &LinearAttnLayer,
7066        qkv_mixed: &cudarc::driver::CudaView<f32>,
7067        z: &cudarc::driver::CudaView<f32>,
7068        beta_raw: &cudarc::driver::CudaView<f32>,
7069        alpha: &cudarc::driver::CudaView<f32>,
7070        t: usize,
7071        cache: &mut Cache,
7072        il: usize,
7073        pad_len: Option<&CudaSlice<i32>>,
7074    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
7075        let cfg = &self.cfg;
7076        let geometry = la.geometry;
7077        let d_state = geometry.key_head_dim as usize;
7078        let num_v = geometry.value_heads as usize;
7079        let eps = cfg.rms_eps;
7080        let scale = 1.0 / (d_state as f32).sqrt();
7081
7082        let prep =
7083            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
7084
7085        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
7086        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
7087        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
7088        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
7089        // verify keep the sequential kernel).
7090        let mut o = e.uninit(d_state * num_v * t)?;
7091        let rl = cache.recur[il].as_mut().unwrap();
7092        {
7093            let crate::cache::RecurLayer {
7094                ssm_state,
7095                ssm_state_alt,
7096                ..
7097            } = rl;
7098            e.gdn_scan_prefill(
7099                &prep.q_l2,
7100                &prep.k_l2,
7101                &prep.v_g,
7102                &prep.g_log,
7103                &prep.beta,
7104                prep.kb16.as_ref(),
7105                prep.qb16.as_ref(),
7106                ssm_state,
7107                ssm_state_alt,
7108                &mut o,
7109                num_v,
7110                t,
7111                scale,
7112                prep.hk,
7113            )?;
7114        }
7115        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7116
7117        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
7118        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
7119        let mut gn = e.uninit(d_state * num_v * t)?;
7120        let gn16 = if Self::f16out_on(e, t) {
7121            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
7122            e.gated_rmsnorm_f16out_zv(
7123                &o,
7124                la.ssm_norm.float_data(),
7125                z,
7126                &mut gn,
7127                &mut g16,
7128                d_state,
7129                num_v * t,
7130                eps,
7131            )?;
7132            Some(g16)
7133        } else {
7134            e.gated_rmsnorm_zv(
7135                &o,
7136                la.ssm_norm.float_data(),
7137                z,
7138                &mut gn,
7139                d_state,
7140                num_v * t,
7141                eps,
7142            )?;
7143            None
7144        };
7145        Ok((gn, gn16))
7146    }
7147
7148    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
7149    #[allow(clippy::too_many_arguments)]
7150    fn linear_attn_prime_core_pad(
7151        &self,
7152        e: &Engine,
7153        la: &LinearAttnLayer,
7154        g4: Vec<CudaSlice<f32>>,
7155        t: usize,
7156        cache: &mut Cache,
7157        il: usize,
7158        pad_len: Option<&CudaSlice<i32>>,
7159    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7160        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
7161        if let Some(xh) = &gn16
7162            && let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)?
7163        {
7164            return Ok(y);
7165        }
7166        e.matmul(&la.ssm_out, &gn, t)
7167    }
7168
7169    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
7170    ///
7171    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
7172    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
7173    pub fn full_attn(
7174        &self,
7175        e: &Engine,
7176        fa: &FullAttnLayer,
7177        h: &CudaSlice<f32>,
7178        pos_d: &CudaSlice<i32>,
7179        t: usize,
7180        il: usize,
7181    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7182        if self.uses_sliding_gated_moe_program() {
7183            return self.step35_attn(e, fa, h, pos_d, t, il);
7184        }
7185        let cfg = &self.cfg;
7186        let _n_embd = cfg.n_embd as usize;
7187        let geometry = cfg.full_attention_geometry_at(il as u32);
7188        let n_head = geometry.n_head as usize;
7189        let n_head_kv = geometry.n_head_kv as usize;
7190        let head_dim = geometry.head_dim_k as usize;
7191        let eps = cfg.rms_eps;
7192        let scale = geometry.attention_scale();
7193
7194        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
7195        // gate — wq out = n_head*head_dim, no split (see prime-path note).
7196        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7197        // A load-time full-attention TP plan owns the same Q/K/V projections for every
7198        // architecture. Fall back to the original grouped owner-device projection when this
7199        // layer has no TP sidecar.
7200        let mut g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
7201            Some(g3) => g3,
7202            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7203        };
7204        let v = g3.pop().unwrap();
7205        let mut k = g3.pop().unwrap();
7206        let qf = g3.pop().unwrap();
7207        let (mut q, gate) = if gated {
7208            let mut q = e.uninit(t * n_head * head_dim)?;
7209            let mut gate = e.uninit(t * n_head * head_dim)?;
7210            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7211            (q, Some(gate))
7212        } else {
7213            (qf, None)
7214        };
7215
7216        // QK-norm (per head_dim row), then partial RoPE.
7217        let mut qn = e.uninit(t * n_head * head_dim)?;
7218        e.rms_norm(
7219            &q,
7220            fa.q_norm.float_data(),
7221            &mut qn,
7222            head_dim,
7223            n_head * t,
7224            eps,
7225        )?;
7226        q = qn;
7227        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7228        e.rms_norm(
7229            &k,
7230            fa.k_norm.float_data(),
7231            &mut kn,
7232            head_dim,
7233            n_head_kv * t,
7234            eps,
7235        )?;
7236        k = kn;
7237        let rope_dims = geometry.n_rot as usize;
7238        e.rope_neox(
7239            &mut q,
7240            pos_d,
7241            head_dim,
7242            rope_dims,
7243            n_head,
7244            t,
7245            geometry.rope_base,
7246            1.0,
7247        )?;
7248        e.rope_neox(
7249            &mut k,
7250            pos_d,
7251            head_dim,
7252            rope_dims,
7253            n_head_kv,
7254            t,
7255            geometry.rope_base,
7256            1.0,
7257        )?;
7258
7259        // SDPA
7260        let mut attn = e.uninit(t * n_head * head_dim)?;
7261        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
7262        // falls back to naive sdpa.
7263        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
7264            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
7265            e.sdpa_naive(
7266                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7267            )?;
7268        } else {
7269            e.fa_prefill(
7270                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
7271            )?;
7272        }
7273
7274        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
7275        let attn_g = match &gate {
7276            Some(gate) => {
7277                let mut gsig = e.uninit(t * n_head * head_dim)?;
7278                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
7279                let mut ag = e.uninit(t * n_head * head_dim)?;
7280                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
7281                ag
7282            }
7283            None => attn,
7284        };
7285
7286        // O follows the same generic load-time TP plan as Q/K/V.
7287        self.full_attn_o(e, fa, &attn_g, t)
7288    }
7289
7290    /// The absorb/decompress operands (`attn_k_b` / `attn_v_b`) are 3D and are ALWAYS the Float
7291    /// arm — on every checkpoint dtype, not just an f32 fixture. Two guards upstream make that a
7292    /// property rather than a hope: `GpuTensor::load_from_source` refuses any quantized non-2D
7293    /// tensor by name (`row_bytes` is derived from `ne[1]`, the MIDDLE axis of a 3D tensor), and
7294    /// `MlaAttnLayer::load` audits residency at load. A quantized `kv_b_proj` is dequantized at
7295    /// the source by `TransformKind::MlaKeyUpSplit`/`MlaValueUpSplit`. This is the last backstop:
7296    /// fail NAMING the constraint rather than through `float_data()`'s norm-flavoured panic.
7297    fn mla_split_operand<'w>(
7298        w: &'w crate::model::GpuTensor,
7299        name: &str,
7300        il: usize,
7301    ) -> &'w CudaSlice<f32> {
7302        match w {
7303            crate::model::GpuTensor::Float { data, .. } => data,
7304            _ => panic!(
7305                "layer {il}: MLA conversion-split operand {name} is not f32-resident. The 3D \
7306                 (d_nope|kv_rank, kv_rank|d_v, n_head) splits have no quantized resident layout: \
7307                 a quantized 3D tensor mis-derives row_bytes in the generic 2D Quant arm, so the \
7308                 source must dequantize the fused kv_b_proj (TensorTransform::SplitMlaKv). \
7309                 Reaching this means both the loader rank guard and MlaAttnLayer::load's \
7310                 residency audit were bypassed"
7311            ),
7312        }
7313    }
7314
7315    /// MLA (multi-head latent attention) mixer core, ABSORBED form — the one arm that serves
7316    /// prefill, chunked prefill and decode (see `cu/mla_attn.cu` FORM CHOICE).
7317    ///
7318    /// `latent` is the layer's latent KV plane; this call APPENDS its own `t` rows at row
7319    /// `slot` and then attends rows `0..slot + t`, which is exactly the oracle's convention
7320    /// that the queries are the LAST `t_q` rows of the cache (`crate::mla::MlaInputs`).
7321    /// Returns the post-`wo` block output [t, n_embd].
7322    #[allow(clippy::too_many_arguments)]
7323    // allow: the parameter list mirrors the kernel/FFI/call contract (rows_exact is the
7324    // verify-batch matmul-class selector, lane/glm5-verify-batch); bundling into a struct
7325    // is a refactor, not a lint fix
7326    fn mla_attn_core(
7327        &self,
7328        e: &Engine,
7329        mla: &crate::hybrid::MlaAttnLayer,
7330        h: &CudaSlice<f32>,
7331        pos_d: &CudaSlice<i32>,
7332        t: usize,
7333        il: usize,
7334        latent: &mut CudaSlice<f32>,
7335        index_plane: Option<IndexerPlanes<'_>>,
7336        slot: usize,
7337        rows_exact: bool,
7338    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7339        let attn = self.mla_attn_core_pre_wo(
7340            e,
7341            mla,
7342            h,
7343            pos_d,
7344            t,
7345            il,
7346            latent,
7347            index_plane,
7348            slot,
7349            rows_exact,
7350        )?;
7351        // Verify-batch wo seam (lane/glm5-verify-batch): the rows arm routes the output
7352        // projection decode-exact, same as every projection inside the core — the wo
7353        // dispatch moved here with the TP split, its routing did not change.
7354        if rows_exact {
7355            e.matmul_rows_exact(&mla.wo, &attn, t)
7356        } else {
7357            e.matmul(&mla.wo, &attn, t)
7358        }
7359    }
7360
7361    /// [`mla_attn_core`] up to (and excluding) the output projection: returns the
7362    /// per-head attention output `[t, n_head * d_v]`. Split out for the glm5 TP-2 seam,
7363    /// whose column-parallel `wo` runs over the cross-rank GATHERED heads — the plain path
7364    /// is the wrapper above, byte-for-byte the pre-split body (the wo matmul moved,
7365    /// nothing else).
7366    #[allow(clippy::too_many_arguments)]
7367    fn mla_attn_core_pre_wo(
7368        &self,
7369        e: &Engine,
7370        mla: &crate::hybrid::MlaAttnLayer,
7371        h: &CudaSlice<f32>,
7372        pos_d: &CudaSlice<i32>,
7373        t: usize,
7374        il: usize,
7375        latent: &mut CudaSlice<f32>,
7376        index_plane: Option<IndexerPlanes<'_>>,
7377        slot: usize,
7378        rows_exact: bool,
7379    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7380        let g = mla.geom;
7381        let cfg = &self.cfg;
7382        let eps = cfg.rms_eps;
7383        let base = cfg.rope_freq_base;
7384        let (nh, dn, dr, dv, r) = (g.n_head, g.d_nope, g.d_rope, g.d_v, g.kv_rank);
7385        assert_eq!(
7386            g.latent_dim,
7387            r + dr,
7388            "layer {il}: MlaGeom latent_dim disagrees with kv_rank + d_rope"
7389        );
7390        let t_kv = slot + t;
7391        // Verify-batch matmul seam (lane/glm5-verify-batch): rows_exact routes every
7392        // projection through the decode-exact classes so each of the t rows is
7393        // bit-identical to the t=1 decode program (matmul_rows_exact contract); false =
7394        // the unchanged dispatch for every other caller (prime keeps its classes).
7395        let mm = |w: &crate::model::GpuTensor,
7396                  x: &CudaSlice<f32>|
7397         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7398            if rows_exact {
7399                e.matmul_rows_exact(w, x, t)
7400            } else {
7401                e.matmul(w, x, t)
7402            }
7403        };
7404
7405        // --- q path: wq_a -> q_a_norm -> wq_b -> per-head [nope | rope] ---
7406        let q_a = mm(&mla.wq_a, h)?;
7407        let q_lora = mla.wq_b.in_features();
7408        let mut q_an = e.uninit(t * q_lora)?;
7409        e.rms_norm(&q_a, mla.q_a_norm.float_data(), &mut q_an, q_lora, t, eps)?;
7410        let q = mm(&mla.wq_b, &q_an)?;
7411        // Per head the row is [nope | rope] contiguous, so t*nh rows of width dn+dr split with
7412        // the same kernel the latent row uses — the two layouts are the same shape.
7413        let mut q_nope = e.uninit(t * nh * dn)?;
7414        let mut q_pe = e.uninit((t * nh * dr).max(1))?;
7415        e.mla_split_latent(&q, &mut q_nope, &mut q_pe, t * nh, dn, dr)?;
7416        // NoPE (glm5_next, rope_head_dim 0): no rope plane exists. The launcher is a no-op, but
7417        // the allocation above is still non-empty so nothing downstream holds a null slice.
7418        e.mla_rope_interleaved(&mut q_pe, pos_d, t, nh, dr, base)?;
7419
7420        // --- kv path: wkv_a -> [c_kv (rms-normed) | k_pe (roped, NOT normed)] ---
7421        let kv = mm(&mla.wkv_a, h)?;
7422        let mut c_kv = e.uninit(t * r)?;
7423        let mut k_pe = e.uninit((t * dr).max(1))?;
7424        e.mla_split_latent(&kv, &mut c_kv, &mut k_pe, t, r, dr)?;
7425        let mut c_kv_n = e.uninit(t * r)?;
7426        e.rms_norm(&c_kv, mla.kv_a_norm.float_data(), &mut c_kv_n, r, t, eps)?;
7427        e.mla_rope_interleaved(&mut k_pe, pos_d, t, 1, dr, base)?;
7428        e.mla_append_latent(latent, &c_kv_n, &k_pe, slot, t, r, dr)?;
7429
7430        // --- DSA k-pool selection, BEFORE attending: the indexer's own state row for each of
7431        // this call's tokens is appended first, so a query sees itself exactly as the latent
7432        // plane already lets it (the reference concatenates into the indexer cache, then scores).
7433        let gathered = match (&mla.index, index_plane) {
7434            (Some(indexer), Some(plane)) => {
7435                Some(self.mla_kpool_select(e, indexer, h, &q_an, plane, t, slot, il, rows_exact)?)
7436            }
7437            (Some(_), None) => {
7438                return Err(format!(
7439                    "layer {il} declares a DSA k-pool indexer but no indexer state plane was \
7440                     supplied — the ModelPlan must declare StatePlan::LatentKvCache with a \
7441                     non-zero index_width for it"
7442                )
7443                .into());
7444            }
7445            (None, _) => None,
7446        };
7447
7448        // --- absorbed MLA core over the latent plane ---
7449        let wk_b = Self::mla_split_operand(&mla.wk_b, "attn_k_b", il);
7450        let wv_b = Self::mla_split_operand(&mla.wv_b, "attn_v_b", il);
7451
7452        // MEMRA_MLA_TC_PREFILL door (default OFF; flag read per call — the rollback seam).
7453        // Engagement conditions, every one load-bearing:
7454        //   * a gathered selection exists — the door serves the DSA arm only; the dense
7455        //     absorbed arm (GLM-5.2, no indexer) keeps the f32 kernel it was gated on;
7456        //   * d_rope == 0 (NoPE) — the TC kernel treats the latent row as both K and V,
7457        //     which is only the whole truth when there is no rope plane;
7458        //   * kv_rank == 512 — the kernel's stamped head dim (glm5_next / GLM-5.2 class);
7459        //   * t >= 16 — prefill widths only. Decode (t == 1) and short resumes NEVER enter,
7460        //     which is what the decode byte-identity gate proves rather than assumes.
7461        // Anything else falls through to the unchanged f32 kernels below — behavior identical
7462        // to the flag being off.
7463        // A chain returning Ok(None) is a cuBLASLt shape DECLINE (announced once per shape);
7464        // the let-chain then simply does not match and the f32 kernels below serve the call.
7465        if let Some((idx, slots)) = &gathered
7466            && dr == 0
7467            && r == 512
7468            && t >= 16
7469            && !rows_exact // verify-batch stays on the decode-exact classes (t <= 15 anyway)
7470            && !crate::portable_mma_gated()
7471            && mla_tc_prefill_enabled()
7472            && let Some(attn) = self.mla_tc_prefill_chain(
7473                e, wk_b, wv_b, &q_nope, latent, idx, *slots, t, t_kv, nh, dn, dv, r, g.scale,
7474            )?
7475        {
7476            return Ok(attn);
7477        }
7478
7479        let mut q_lat = e.uninit(t * nh * r)?;
7480        e.mla_absorb_q(&q_nope, wk_b, &mut q_lat, t, nh, dn, r)?;
7481        let mut o_lat = e.uninit(t * nh * r)?;
7482        match &gathered {
7483            Some((idx, slots)) => e.mla_attn_gathered(
7484                &q_lat, &q_pe, latent, idx, &mut o_lat, nh, r, dr, t, *slots, g.scale,
7485            )?,
7486            None => e.mla_attn_absorbed(
7487                &q_lat, &q_pe, latent, &mut o_lat, nh, r, dr, t, t_kv, g.scale,
7488            )?,
7489        }
7490        let mut attn = e.uninit(t * nh * dv)?;
7491        e.mla_decompress_v(&o_lat, wv_b, &mut attn, t, nh, dv, r)?;
7492
7493        Ok(attn)
7494    }
7495
7496    /// The MEMRA_MLA_TC_PREFILL chain: absorb and decompress as strided-batched bf16
7497    /// tensor-core GEMMs, attention as the gathered bf16 MMA kernel. Returns `Ok(None)` when
7498    /// cuBLASLt declines a GEMM shape (announced once per shape) so the caller falls back to
7499    /// the f32 kernels; every other failure is a hard error.
7500    ///
7501    /// FORM CHOICE, stated for the record (the dual-form MLA law): every fast engine runs
7502    /// MATERIALIZED (per-head MHA) attention at DENSE prefill and absorbed MQA at decode.
7503    /// glm5_next prefill is NOT dense: the DSA indexer caps every query at topk+tail rows and
7504    /// selects ONE list per query SHARED ACROSS ALL 64 HEADS. That shared list is what makes
7505    /// the ABSORBED form the GEMM-shaped one here — the head axis is the MMA m, the shared
7506    /// latent rows are one B operand per tile — while materializing K/V would give every head
7507    /// its own K plane and destroy exactly that sharing (back to per-(query,head) matvecs on
7508    /// the gathered walk). It is also FlashMLA's own sparse-prefill geometry (q 576/512 over
7509    /// gathered latent rows). Queries whose selection is trivial (visible <= topk: the lists
7510    /// ARE the full causal prefix, emitted by the selector itself) ride the SAME kernel with
7511    /// the identity gather — there is no separate dense program to gate.
7512    ///
7513    /// Transient cost per (layer, chunk) at the census shape (t=2313, t_kv=4626): bf16 q_lat
7514    /// 152 MB + bf16 latent window 4.7 MB + bf16 q_nope/o_lat copies ~230 MB — all freed with
7515    /// the call. (The materialized-K/V alternative would have been 4096 x 64 x (256+256) x 2B
7516    /// = 256 MB/layer-chunk of K/V ALONE, plus the per-head-K program cost above.) Weight
7517    /// bf16 converts (wk_b/wv_b, 8.4M elems each) run per call, ~50 us class; a resident
7518    /// mirror is a later diet, not correctness.
7519    #[allow(clippy::too_many_arguments)]
7520    fn mla_tc_prefill_chain(
7521        &self,
7522        e: &Engine,
7523        wk_b: &CudaSlice<f32>,
7524        wv_b: &CudaSlice<f32>,
7525        q_nope: &CudaSlice<f32>,
7526        latent: &CudaSlice<f32>,
7527        idx: &CudaSlice<i32>,
7528        width: usize,
7529        t: usize,
7530        t_kv: usize,
7531        nh: usize,
7532        dn: usize,
7533        dv: usize,
7534        r: usize,
7535        scale: f32,
7536    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7537        // Once-per-shape decline announce (the bf16_tc_gemm pattern): a door that quietly
7538        // stops engaging reads exactly like a door that never helped.
7539        fn declined(stage: &str, m: usize, n: usize, k: usize, batch: usize) {
7540            type ShapeSet = std::collections::HashSet<(usize, usize, usize, usize)>;
7541            static SAID: std::sync::Mutex<Option<ShapeSet>> = std::sync::Mutex::new(None);
7542            let mut g = SAID.lock().unwrap();
7543            if g.get_or_insert_with(std::collections::HashSet::new)
7544                .insert((m, n, k, batch))
7545            {
7546                eprintln!(
7547                    "[mla-tc-prefill] DECLINED at {stage} m={m} n={n} k={k} batch={batch} \
7548                     (no cuBLASLt heuristic) — this call falls back to the f32 MLA kernels"
7549                );
7550            }
7551        }
7552        // Weights and activations to bf16. The converts require n % 4 == 0; every operand here
7553        // is a multiple of the head dims (dn/dv/r all >= 16 and % 4 == 0 on the shapes the door
7554        // admits), asserted rather than assumed.
7555        for (name, n) in [
7556            ("wk_b", nh * r * dn),
7557            ("wv_b", nh * dv * r),
7558            ("q_nope", t * nh * dn),
7559            ("latent", t_kv * r),
7560        ] {
7561            debug_assert!(
7562                n.is_multiple_of(4),
7563                "mla-tc-prefill: {name} elems {n} % 4 != 0"
7564            );
7565            let _ = (name, n);
7566        }
7567        let wk_bf = e.f32_to_bf16(wk_b, nh * r * dn)?;
7568        let wv_bf = e.f32_to_bf16(wv_b, nh * dv * r)?;
7569        let qn_bf = e.f32_to_bf16(q_nope, t * nh * dn)?;
7570        // absorb: per head h, q_lat[:,h,:] [t, r] = q_nope[:,h,:] [t, dn] @ W_uk[h] [r, dn]^T.
7571        // wk_b is the conversion-split (h, l, p) plane, contiguous in p == the reduction axis:
7572        // per head it IS the [n=r, k=dn] row-major operand. bf16 out feeds the attention kernel.
7573        let mut q_lat_bf = e.alloc_u8_uninit(t * nh * r * 2)?;
7574        if !e.mla_bf16_gemm_sb_bf16out(
7575            &wk_bf,
7576            &qn_bf,
7577            &mut q_lat_bf,
7578            t,
7579            r,
7580            dn,
7581            nh * dn,
7582            dn,
7583            nh * r,
7584            r,
7585            nh,
7586        )? {
7587            declined("absorb", t, r, dn, nh);
7588            return Ok(None);
7589        }
7590        // The latent window rows 0..t_kv (this call's rows were appended above), bf16.
7591        let cache_bf = e.f32_to_bf16(latent, t_kv * r)?;
7592        let mut o_lat = e.uninit(t * nh * r)?;
7593        e.mla_attn_gathered_tc(
7594            &q_lat_bf, &cache_bf, idx, &mut o_lat, nh, r, t, width, scale,
7595        )?;
7596        // decompress: per head h, attn[:,h,:] [t, dv] = o_lat[:,h,:] [t, r] @ W_uv[h] [dv, r]^T.
7597        // wv_b is (h, j, l), contiguous in l == the reduction axis: per head [n=dv, k=r].
7598        let o_bf = e.f32_to_bf16(&o_lat, t * nh * r)?;
7599        let mut attn = e.uninit(t * nh * dv)?;
7600        if !e.mla_bf16_gemm_sb_f32out(
7601            &wv_bf,
7602            &o_bf,
7603            &mut attn,
7604            t,
7605            dv,
7606            r,
7607            nh * r,
7608            r,
7609            nh * dv,
7610            dv,
7611            nh,
7612        )? {
7613            declined("decompress", t, dv, r, nh);
7614            return Ok(None);
7615        }
7616        crate::MLA_TC_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7617        {
7618            static ANNOUNCED: std::sync::Once = std::sync::Once::new();
7619            ANNOUNCED.call_once(|| {
7620                eprintln!(
7621                    "[mla-tc-prefill] engaged: absorb/decompress = strided-batched bf16 TC \
7622                     GEMMs, attention = fa_mla_gathered_bf16 (t={t}, t_kv={t_kv}, nh={nh}, \
7623                     width={width}); dispatches counted in MLA_TC_PREFILL_DISPATCHES"
7624                );
7625            });
7626        }
7627        Ok(Some(attn))
7628    }
7629
7630    /// Layer-scoped wrapper: names the layer in any selection failure.
7631    #[allow(clippy::too_many_arguments)]
7632    fn mla_kpool_select(
7633        &self,
7634        e: &Engine,
7635        indexer: &crate::hybrid::MlaIndexer,
7636        h: &CudaSlice<f32>,
7637        q_resid: &CudaSlice<f32>,
7638        plane: IndexerPlanes<'_>,
7639        t: usize,
7640        slot: usize,
7641        il: usize,
7642        rows_exact: bool,
7643    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7644        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, rows_exact).map_err(
7645            |source| -> Box<dyn std::error::Error> {
7646                format!("layer {il}: DSA k-pool selection failed: {source}").into()
7647            },
7648        )
7649    }
7650
7651    /// DSA k-pool indexer: append this call's packed indexer state, then select the cache rows
7652    /// each query may attend. Returns the per-query position list and its width (`-1` padded).
7653    ///
7654    /// The program is `Glm5NextTextIndexer.forward`
7655    /// (research/glm53-flash-bringup-20260827/modular_glm5_next-ref.py:771), transcribed in
7656    /// `memra_reference::kpool_allowed_tokens`, which is this path's oracle:
7657    ///   1. `k = LayerNorm_affine(wk(x))` — LayerNorm WITH BIAS at eps 1e-5, NOT the model's
7658    ///      RMSNorm at `rms_norm_eps`; `gate = index_kpool_compress_gate(x)`. Both are cached.
7659    ///   2. Every COMPLETE pool of `pool` consecutive cached tokens collapses to one key by a
7660    ///      per-channel softmax over (gate + positional embedding).
7661    ///   3. `score[i][p] = sum_h relu(q[i][h] . pool_key[p] * d^-1/2) * weights_proj(x)[i][h] *
7662    ///      heads^-1/2`, with pools whose last token is invisible to the query masked out.
7663    ///   4. Top `top_k / pool` pools expand back to raw rows; the incomplete tail is appended raw.
7664    ///
7665    /// `q_resid` is `q_a_layernorm(q_a_proj(x))` — the SAME tensor the MLA query up-projection
7666    /// consumes, which is why the indexer is scored here rather than before the core.
7667    #[allow(clippy::too_many_arguments)]
7668    pub fn mla_kpool_indices(
7669        e: &Engine,
7670        indexer: &crate::hybrid::MlaIndexer,
7671        h: &CudaSlice<f32>,
7672        q_resid: &CudaSlice<f32>,
7673        plane: IndexerPlanes<'_>,
7674        t: usize,
7675        slot: usize,
7676    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7677        Self::mla_kpool_indices_ex(e, indexer, h, q_resid, plane, t, slot, false)
7678    }
7679
7680    /// [`Self::mla_kpool_indices`] with the verify-batch matmul-class selector
7681    /// (lane/glm5-verify-batch): `rows_exact` routes the indexer's four projections
7682    /// through the decode-exact classes so each of the t rows is bit-identical to the
7683    /// t=1 decode program; `false` is the unchanged dispatch.
7684    #[allow(clippy::too_many_arguments)]
7685    pub fn mla_kpool_indices_ex(
7686        e: &Engine,
7687        indexer: &crate::hybrid::MlaIndexer,
7688        h: &CudaSlice<f32>,
7689        q_resid: &CudaSlice<f32>,
7690        plane: IndexerPlanes<'_>,
7691        t: usize,
7692        slot: usize,
7693        rows_exact: bool,
7694    ) -> Result<(CudaSlice<i32>, usize), Box<dyn std::error::Error>> {
7695        let mm = |w: &crate::model::GpuTensor,
7696                  x: &CudaSlice<f32>|
7697         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7698            if rows_exact {
7699                e.matmul_rows_exact(w, x, t)
7700            } else {
7701                e.matmul(w, x, t)
7702            }
7703        };
7704        /// `nn.LayerNorm` default epsilon. The indexer's k_norm is a LayerNorm, so it does NOT
7705        /// take the model's `rms_norm_eps` (census: "eps 1e-5, NOT rms_norm_eps"); they coincide
7706        /// numerically on GLM-5.3-Flash and the constant keeps them from being coupled.
7707        const INDEX_NORM_EPS: f32 = 1e-5;
7708
7709        let ig = indexer.geom;
7710        let d = ig.head_dim;
7711        let t_kv = slot + t;
7712        let IndexerPlanes {
7713            state: plane,
7714            pool_keys: pool_key_plane,
7715            ready: pools_ready,
7716            state_ring_rows,
7717            capacity_tokens,
7718        } = plane;
7719
7720        // TAIL RING. The plane's rows are read EXACTLY ONCE, by the pool-key build of the pool
7721        // each row belongs to, so a ring of `ring` rows holds everything still live. The state
7722        // plan does not carry `pool`, so the allocator books physical rows and the EFFECTIVE ring
7723        // is rounded down here: a ring that is not a whole number of pools would split a pool
7724        // across the wrap. ONE POOL is the whole correctness floor (lane/glm53-ring-sizing): the
7725        // drain below serves any `t` from any ring at or above it, so `ring` never bounds a
7726        // prompt.
7727        let ring = if state_ring_rows == 0 {
7728            0
7729        } else {
7730            state_ring_rows / ig.pool * ig.pool
7731        };
7732        if state_ring_rows > 0 && ring == 0 {
7733            return Err(format!(
7734                "indexer tail ring of {state_ring_rows} rows cannot hold one pool of {}; \
7735                 raise MEMRA_DSA_INDEX_RING or set it to 0 for the flat plane",
7736                ig.pool
7737            )
7738            .into());
7739        }
7740        // RESIDENCY TRIPWIRE. `*pools_ready` counts pools whose keys were built over rows that are
7741        // now history. If the cache ever rewound past `slot` without clamping it (see
7742        // `LatentKvLayer::truncate_index_pool_keys`), those keys were built over rows this call is
7743        // about to overwrite — a silent wrong selection. Fail here instead.
7744        if *pools_ready > slot / ig.pool {
7745            return Err(format!(
7746                "resident k-pool key plane claims {} finished pools but the cache holds only {} \
7747                 complete pools before this call ({slot} rows / pool {}) — a rewind reduced the \
7748                 latent length without clamping index_pools_ready",
7749                *pools_ready,
7750                slot / ig.pool,
7751                ig.pool
7752            )
7753            .into());
7754        }
7755
7756        // 1. packed state rows [k | gate], appended at `slot` — the same [a|b] row shape the
7757        //    latent plane uses, so `mla_append_latent` packs it with no new kernel.
7758        let k_raw = mm(&indexer.wk, h)?;
7759        let mut k_norm = e.uninit(t * d)?;
7760        e.layer_norm_bias(
7761            &k_raw,
7762            indexer.k_norm_w.float_data(),
7763            indexer.k_norm_b.float_data(),
7764            &mut k_norm,
7765            d,
7766            t,
7767            INDEX_NORM_EPS,
7768        )?;
7769        let gate = mm(&indexer.kpool_gate, h)?;
7770
7771        // 2. pool keys over every COMPLETE pool in the cache — INCREMENTALLY. A pool's key is a
7772        //    function of its own `pool` state rows (append-only, never rewritten) and the constant
7773        //    `kpool_ape`, so it is final the instant the pool's last row lands. Only pools
7774        //    `[*pools_ready, n_pools)` are built; the rest are already resident and bit-identical
7775        //    to a rebuild. This turns the old O(t_kv * d) per-call pass into O(t * d).
7776        let n_pools = t_kv / ig.pool;
7777        let select_k = ig.select_k(n_pools);
7778        let width = ig.index_width(n_pools);
7779        // Sized to the SESSION's capacity, so a session that primes and then decodes never
7780        // reallocates (a fresh buffer would drop every resident key, and under the ring the rows
7781        // to rebuild them from are gone). `capacity_tokens` is that capacity; a pool covers
7782        // `ig.pool` tokens, so the key plane is `pool` times SHORTER than a flat state plane —
7783        // 32 f32 per token against 256. It is NOT read off `plane.len()` any more: once the state
7784        // plane is a ring, its length is one call's tail, not the context.
7785        // `.max(1)` keeps the slice non-null at t_kv < pool, where no complete pool exists yet.
7786        let capacity_pools = capacity_tokens / ig.pool;
7787        let need = (capacity_pools * d).max(n_pools * d).max(1);
7788        if pool_key_plane.as_ref().is_none_or(|k| k.len() < need) {
7789            *pool_key_plane = Some(e.uninit(need)?);
7790            *pools_ready = 0;
7791        }
7792        let pool_keys = pool_key_plane
7793            .as_mut()
7794            .expect("resident pool-key plane just allocated");
7795
7796        // THE DRAIN. The state plane is written by exactly one kernel and read by exactly one,
7797        // and a row's single read is the pool-key build of the pool that row belongs to. So the
7798        // rows that must be live at any instant are `[*pools_ready * pool, cur)`: everything
7799        // below has been read, everything above is not written yet, and the two kernels can be
7800        // interleaved in sub-ranges of the call instead of run once each over the whole call.
7801        //
7802        // That is what makes the ring size a WORKING-SET choice rather than a bound on `t`:
7803        // `index_ring_take` hands back how many rows fit before the ring must be drained, the
7804        // build drains it, and the loop continues. `k_norm`/`gate` are computed ONCE for the
7805        // whole call above and walked by source-row offset, so the values, their order, and the
7806        // ring addresses they land on are exactly what a single whole-call append produced.
7807        // A flat plane (`ring == 0`) takes the whole call in one iteration, byte for byte.
7808        let ape = indexer.kpool_ape.float_data();
7809        let mut cur = slot;
7810        let mut appended = 0usize;
7811        while appended < t {
7812            let take =
7813                crate::cache::index_ring_take(ring, ig.pool, *pools_ready, cur, t - appended)
7814                    .ok_or_else(|| -> Box<dyn std::error::Error> {
7815                        format!(
7816                            "indexer tail ring lapped: {ring} rows cannot hold the {} rows still \
7817                         owed to unbuilt pools at row {cur} (pools_ready {}, pool {}, slot \
7818                         {slot}, t {t}). The pool-key plane was reset or the cache rewound \
7819                         without clamping index_pools_ready, so rows this call must read were \
7820                         already overwritten. Raise MEMRA_DSA_INDEX_RING, or set \
7821                         MEMRA_DSA_INDEX_RING=0 for the flat plane",
7822                            cur.saturating_sub((*pools_ready).saturating_mul(ig.pool)),
7823                            *pools_ready,
7824                            ig.pool
7825                        )
7826                        .into()
7827                    })?;
7828            debug_assert!(take > 0 && appended + take <= t);
7829            e.mla_index_append(plane, &k_norm, &gate, appended, cur, take, d, d, ring)?;
7830            cur += take;
7831            appended += take;
7832            let ready_now = cur / ig.pool;
7833            e.mla_kpool_pool_keys(
7834                plane,
7835                ape,
7836                pool_keys,
7837                (*pools_ready).min(ready_now),
7838                ready_now,
7839                ig.pool,
7840                d,
7841                ring,
7842            )?;
7843            *pools_ready = ready_now;
7844        }
7845        debug_assert!(t == 0 || *pools_ready == n_pools);
7846        let pool_keys = &*pool_keys;
7847
7848        // 3. score + head mix, 4. top-k -> raw rows + tail.
7849        let q_index = mm(&indexer.wq_b, q_resid)?;
7850        let head_weights = mm(&indexer.weights_proj, h)?;
7851        let mut score = e.uninit((t * n_pools).max(1))?;
7852        e.mla_kpool_score(
7853            &q_index,
7854            pool_keys,
7855            &head_weights,
7856            &mut score,
7857            t,
7858            ig.heads,
7859            d,
7860            n_pools,
7861            ig.pool,
7862            slot,
7863            (d as f32).powf(-0.5),
7864            (ig.heads as f32).powf(-0.5),
7865        )?;
7866        let mut idx = e.uninit_i32(t * width)?;
7867        e.mla_kpool_select(
7868            &score,
7869            &mut idx,
7870            t,
7871            n_pools,
7872            ig.pool,
7873            select_k,
7874            width,
7875            slot,
7876            ig.always_select_tail,
7877        )?;
7878        Ok((idx, width))
7879    }
7880
7881    /// STATELESS MLA arm (`HybridModel::forward`): the latent plane lives for this call only,
7882    /// sized to the request. Same math as the cached arm — it is the same core with slot 0.
7883    pub fn mla_attn(
7884        &self,
7885        e: &Engine,
7886        mla: &crate::hybrid::MlaAttnLayer,
7887        h: &CudaSlice<f32>,
7888        pos_d: &CudaSlice<i32>,
7889        t: usize,
7890        il: usize,
7891    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7892        if mla.tp.is_some() {
7893            return Err(format!(
7894                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the stateless \
7895                 mixer path is unwired for a head shard"
7896            )
7897            .into());
7898        }
7899        let mut latent = e.uninit(t * mla.geom.latent_dim)?;
7900        let mut index_plane = match mla.index.as_ref() {
7901            Some(indexer) => Some(e.uninit(t * indexer.geom.state_width())?),
7902            None => None,
7903        };
7904        // No residency across calls here — the planes die with the call, so `ready` starts at 0 and
7905        // every pool is built exactly once, which is what the cached arm also does on its prime.
7906        let mut pool_keys = None;
7907        let mut pools_ready = 0usize;
7908        let planes = index_plane.as_mut().map(|state| IndexerPlanes {
7909            state,
7910            pool_keys: &mut pool_keys,
7911            ready: &mut pools_ready,
7912            // Per-call plane, sized to the request: no ring, capacity is the request itself.
7913            state_ring_rows: 0,
7914            capacity_tokens: t,
7915        });
7916        self.mla_attn_core(e, mla, h, pos_d, t, il, &mut latent, planes, 0, false)
7917    }
7918
7919    /// STATEFUL MLA arm (prime and T=1 decode): appends into the session's latent plane and
7920    /// attends the whole history. `cache.latent[il]` is allocated by the `LatentKvCache` arm of
7921    /// the cache allocator; a `None` here means the ModelPlan and the loaded mixer disagree.
7922    #[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
7923    pub fn mla_attn_cached(
7924        &self,
7925        e: &Engine,
7926        mla: &crate::hybrid::MlaAttnLayer,
7927        h: &CudaSlice<f32>,
7928        pos_d: &CudaSlice<i32>,
7929        t: usize,
7930        il: usize,
7931        cache: &mut Cache,
7932    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7933        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, false)
7934    }
7935
7936    /// [`Self::mla_attn_cached`] on the VERIFY-BATCH matmul classes (lane/glm5-verify-batch):
7937    /// the SAME core at t=K+1 rows with every internal projection routed decode-exact
7938    /// (`matmul_rows_exact`), so row r of the batched call is bit-identical to the t=1
7939    /// `mla_attn_cached` call the per-row verify walk makes at position pos0+r. Causality
7940    /// within the batch is per-query by construction: the kpool selection masks pools
7941    /// invisible to each query and appends each query's OWN raw tail
7942    /// (`first_pos + t + 1`), and the gathered attention walks each query's own idx list
7943    /// (-1 padding arithmetic-invariant). Held by `glm5_tparallel_verify_gpu` gates 1+2
7944    /// running the batched arm. ONLY the glm5 verify-batch walk calls this.
7945    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
7946    pub fn mla_attn_cached_rows_exact(
7947        &self,
7948        e: &Engine,
7949        mla: &crate::hybrid::MlaAttnLayer,
7950        h: &CudaSlice<f32>,
7951        pos_d: &CudaSlice<i32>,
7952        t: usize,
7953        il: usize,
7954        cache: &mut Cache,
7955    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7956        self.mla_attn_cached_inner(e, mla, h, pos_d, t, il, cache, true)
7957    }
7958
7959    #[allow(clippy::too_many_arguments)] // allow: mirrors mla_attn_cached's contract
7960    fn mla_attn_cached_inner(
7961        &self,
7962        e: &Engine,
7963        mla: &crate::hybrid::MlaAttnLayer,
7964        h: &CudaSlice<f32>,
7965        pos_d: &CudaSlice<i32>,
7966        t: usize,
7967        il: usize,
7968        cache: &mut Cache,
7969        rows_exact: bool,
7970    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7971        // glm5 TP fail-closed choke point, covering BOTH plain entries (decode/prime AND
7972        // the verify-batch rows arm): a TP-sharded layer holds heads/2 and a per-rank
7973        // latent replica — running it on the plain path would compute a silently-halved
7974        // mixer against the wrong plane, so it refuses by name instead.
7975        if mla.tp.is_some() {
7976            return Err(format!(
7977                "layer {il}: MLA layer is glm5-TP-sharded (MEMRA_GLM5_TP): the plain mixer \
7978                 path is unwired for a head shard — only the TP decode/prime walk may \
7979                 execute it (rows_exact={rows_exact})"
7980            )
7981            .into());
7982        }
7983        // Read before the layer borrow: it sizes the resident pool-key plane, which the ring'd
7984        // state plane's own length can no longer stand in for.
7985        let max_ctx = cache.max_ctx;
7986        let layer = cache.latent[il].as_mut().ok_or_else(|| {
7987            format!(
7988                "layer {il} is Mixer::Mla but the cache has no latent plane — the ModelPlan \
7989                 must declare StatePlan::LatentKvCache for it"
7990            )
7991        })?;
7992        let attn =
7993            self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?;
7994        // Verify-batch wo seam: the rows arm keeps its decode-exact output projection —
7995        // the wo dispatch moved here with the TP split, its routing did not change.
7996        if rows_exact {
7997            e.matmul_rows_exact(&mla.wo, &attn, t)
7998        } else {
7999            e.matmul(&mla.wo, &attn, t)
8000        }
8001    }
8002
8003    /// The stateful MLA call against ONE latent plane, up to (and excluding) the output
8004    /// projection. The plain path wraps it above (canonical plane + `wo`); the glm5 TP-2
8005    /// walk calls it once per rank (root shard on the canonical plane, peer shard on the
8006    /// replicated peer plane) and joins the halves through the column-parallel `wo`.
8007    #[allow(clippy::too_many_arguments)]
8008    pub(crate) fn mla_attn_cached_pre_wo(
8009        &self,
8010        e: &Engine,
8011        mla: &crate::hybrid::MlaAttnLayer,
8012        h: &CudaSlice<f32>,
8013        pos_d: &CudaSlice<i32>,
8014        t: usize,
8015        il: usize,
8016        layer: &mut memra_kv::LatentKvLayer,
8017        max_ctx: usize,
8018        rows_exact: bool,
8019    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8020        let slot = layer.len;
8021        let width = layer.width;
8022        assert_eq!(
8023            width, mla.geom.latent_dim,
8024            "layer {il}: cache latent width {width} != MlaGeom latent_dim {}",
8025            mla.geom.latent_dim
8026        );
8027        let capacity = layer.rows.len() / width;
8028        if slot + t > capacity {
8029            return Err(format!(
8030                "layer {il}: latent cache overflow — {slot} + {t} rows exceeds capacity {capacity}"
8031            )
8032            .into());
8033        }
8034        if mla.index.is_some() && layer.index_rows.is_none() {
8035            return Err(format!(
8036                "layer {il} loaded a DSA k-pool indexer but its latent cache carries no indexer \
8037                 state plane — StatePlan::LatentKvCache declared index_width 0 for a layer whose \
8038                 SparseIndexPlan is Own {{ kpool: Some(..) }}"
8039            )
8040            .into());
8041        }
8042        // Both planes are borrowed for the whole core call, so len bookkeeping happens after.
8043        // ONE `len` covers both: they are appended in the same call and must never drift.
8044        // The resident pool-key plane rides along: it is state that must SURVIVE the call, so the
8045        // core writes `ready` back through the borrow and it is restored with the buffers.
8046        let mut rows = std::mem::replace(&mut layer.rows, e.uninit(0)?);
8047        let mut index_rows = layer.index_rows.take();
8048        let mut pool_keys = layer.index_pool_keys.take();
8049        let mut pools_ready = layer.index_pools_ready;
8050        let index_ring_rows = layer.index_ring_rows.unwrap_or(0);
8051        let planes = index_rows.as_mut().map(|state| IndexerPlanes {
8052            state,
8053            pool_keys: &mut pool_keys,
8054            ready: &mut pools_ready,
8055            state_ring_rows: index_ring_rows,
8056            capacity_tokens: max_ctx,
8057        });
8058        let out =
8059            self.mla_attn_core_pre_wo(e, mla, h, pos_d, t, il, &mut rows, planes, slot, rows_exact);
8060        layer.rows = rows;
8061        layer.index_rows = index_rows;
8062        layer.index_pool_keys = pool_keys;
8063        // A FAILED core leaves `len` where it was, so the resident plane must go back too: it may
8064        // have advanced over pools built from rows a retry is about to rewrite with different
8065        // inputs. Clamping here (rather than letting the next call's tripwire fire) makes
8066        // retry-after-error correct instead of merely loud.
8067        layer.index_pools_ready = if out.is_ok() {
8068            pools_ready
8069        } else if let Some(indexer) = mla.index.as_ref() {
8070            pools_ready.min(layer.len / indexer.geom.pool)
8071        } else {
8072            pools_ready
8073        };
8074        let out = out?;
8075        // Resolve the layer's RESIDENT pool copy (the state plan does not carry `pool`; the
8076        // latent-plane snapshot path reads this field to address the tail ring). A nonzero
8077        // resident value that disagrees with the loaded geometry is corruption, not a race:
8078        // there is exactly one geometry per loaded layer.
8079        if let Some(indexer) = mla.index.as_ref() {
8080            let pool = indexer.geom.pool;
8081            if layer.index_pool != 0 && layer.index_pool != pool {
8082                return Err(format!(
8083                    "layer {il}: resident indexer pool {} != loaded geometry pool {pool}",
8084                    layer.index_pool,
8085                )
8086                .into());
8087            }
8088            layer.index_pool = pool;
8089        }
8090        layer.len = slot + t;
8091        let len_i32 = i32::try_from(layer.len).map_err(|_| "latent length exceeds i32 mirror")?;
8092        // Door H (`MEMRA_GLM5_HTOD_DIET`): the async `i32_set_k` launch instead of this
8093        // SYNCHRONIZING pageable 4-byte copy — 11 of these per round, one per MLA trunk layer.
8094        e.i32_mirror_store(&mut layer.len_d, len_i32)?;
8095        Ok(out)
8096    }
8097
8098    /// The glm5 TP MLA walk for one prime/decode call (`mla` is the ROOT head shard; its
8099    /// sidecar carries the peer shards + runtime). Replicated per-token work runs on EVERY
8100    /// rank from identical inputs (wq_a/wkv_a/indexer/k-pool selection — identical bytes by
8101    /// determinism on uniform hardware, gate-held); each rank attends with its heads over
8102    /// its OWN latent replica; the attention parts are gathered through the armed transport
8103    /// and each rank's COLUMN-parallel `wo` slice computes its slice of the output with the
8104    /// plain matvec kernel — no cross-rank arithmetic anywhere.
8105    #[allow(clippy::too_many_arguments)]
8106    pub(crate) fn mla_tp_attn_cached(
8107        &self,
8108        e: &Engine,
8109        mla: &crate::hybrid::MlaAttnLayer,
8110        h: &CudaSlice<f32>,
8111        pos_d: &CudaSlice<i32>,
8112        t: usize,
8113        il: usize,
8114        cache: &mut Cache,
8115        rows_exact: bool,
8116    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8117        let tp = mla
8118            .tp
8119            .as_ref()
8120            .ok_or("mla_tp_attn_cached called on an unsharded layer")?;
8121        let rt = &tp.rt;
8122        let ranks = tp.ranks();
8123        let g = mla.geom; // SHARD geometry: n_head = full/ranks
8124        let hl = g.n_head;
8125        let dv = g.d_v;
8126        let full_heads = tp.full_heads;
8127        let n_embd = tp.n_embd;
8128        let hh = n_embd / ranks;
8129        let max_ctx = cache.max_ctx;
8130
8131        // HOP 1 — fan-out of the mixer input and positions to every peer rank. Both move the
8132        // WHOLE buffer, exactly as the v1 arm did, so the transport arms move identical
8133        // byte ranges (lane/glm5-tp-transport).
8134        let hop = rt.hop(e);
8135        let h_peers = crate::glm5_tp_transport::fanout_f32(&hop, h, h.len())?;
8136        let pos_peers = crate::glm5_tp_transport::fanout_i32(&hop, pos_d, pos_d.len())?;
8137
8138        // Peer replica planes (lazily geometry-cloned from the canonical plane).
8139        {
8140            let canonical = cache.latent[il].as_ref().ok_or_else(|| {
8141                format!("layer {il}: glm5 TP MLA walk found no canonical latent plane")
8142            })?;
8143            crate::glm5_tp::ensure_mla_peer_latent(
8144                rt,
8145                canonical,
8146                &mut cache.glm5_tp_latent_peer[il],
8147            )?;
8148        }
8149
8150        // Peer passes first (each rank's heads over its replica), then root (canonical
8151        // plane unchanged) — v1's issue order at two ranks. `rows_exact` threads the
8152        // caller's matmul class through every rank: false = the prime/decode walk
8153        // (byte-for-byte the pre-composition arm), true = the spec x TP verify walk
8154        // (lane/glm5-composition) riding the same rows-exact classes as the unsharded
8155        // verify walk.
8156        let mut attn: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
8157        for r in 1..ranks {
8158            let layer = &mut cache.glm5_tp_latent_peer[il].as_mut().unwrap()[r - 1];
8159            attn[r] = Some(self.mla_attn_cached_pre_wo(
8160                &rt.peers[r - 1],
8161                &tp.peers[r - 1],
8162                &h_peers[r - 1],
8163                &pos_peers[r - 1],
8164                t,
8165                il,
8166                layer,
8167                max_ctx,
8168                rows_exact,
8169            )?);
8170        }
8171        attn[0] = {
8172            let layer = cache.latent[il].as_mut().unwrap();
8173            Some(self.mla_attn_cached_pre_wo(e, mla, h, pos_d, t, il, layer, max_ctx, rows_exact)?)
8174        };
8175
8176        // HOP 2 — gather the per-head parts into the FULL [t, full_heads*dv] layout on
8177        // every rank. `full_heads * dv == ranks * (hl * dv)` by the shard map.
8178        let part = hl * dv;
8179        debug_assert_eq!(full_heads * dv, ranks * part);
8180        let attn_refs: Vec<&CudaSlice<f32>> = attn
8181            .iter()
8182            .map(|a| a.as_ref().expect("filled above"))
8183            .collect();
8184        let fulls = crate::glm5_tp_transport::gather_parts(&hop, &attn_refs, t, part)?;
8185
8186        // Column-parallel wo slices + output concat (pure movement). The verify walk's
8187        // wo rides the rows-exact class, exactly like the unsharded verify walk's wo.
8188        let mut ys = Vec::with_capacity(ranks);
8189        if rows_exact {
8190            ys.push(e.matmul_rows_exact(&mla.wo, &fulls[0], t)?);
8191            for r in 1..ranks {
8192                ys.push(rt.peers[r - 1].matmul_rows_exact(&tp.peers[r - 1].wo, &fulls[r], t)?);
8193            }
8194        } else {
8195            ys.push(e.matmul(&mla.wo, &fulls[0], t)?);
8196            for r in 1..ranks {
8197                ys.push(rt.peers[r - 1].matmul(&tp.peers[r - 1].wo, &fulls[r], t)?);
8198            }
8199        }
8200        // HOP 3 — concat the column parts into the mixer output on ROOT.
8201        debug_assert_eq!(n_embd, ranks * hh);
8202        let y_refs: Vec<&CudaSlice<f32>> = ys.iter().collect();
8203        crate::glm5_tp_transport::concat_parts_on_root(&hop, &y_refs, t, hh)
8204    }
8205
8206    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
8207    pub fn linear_attn(
8208        &self,
8209        e: &Engine,
8210        la: &LinearAttnLayer,
8211        h: &CudaSlice<f32>,
8212        t: usize,
8213    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8214        let cfg = &self.cfg;
8215        let _n_embd = cfg.n_embd as usize;
8216        let geometry = la.geometry;
8217        let d_state = geometry.key_head_dim as usize;
8218        let num_k = geometry.key_heads as usize;
8219        let num_v = geometry.value_heads as usize;
8220        let d_conv = geometry.conv_kernel as usize;
8221        let head_k = d_state;
8222        let head_v = geometry.value_head_dim as usize;
8223        let key_dim = head_k * num_k; // 2048
8224        let value_dim = head_v * num_v; // 4096
8225        let conv_dim = key_dim * 2 + value_dim; // 8192
8226        let eps = cfg.rms_eps;
8227        let scale = 1.0 / (d_state as f32).sqrt();
8228
8229        // projections
8230        // grouped: one f16 activation convert feeds all four projections (matmul_group)
8231        let mut g4 = e.matmul_group(
8232            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
8233            h,
8234            t,
8235        )?;
8236        let alpha = g4.pop().unwrap(); // [T, num_v]
8237        let beta_raw = g4.pop().unwrap(); // [T, num_v]
8238        let z = g4.pop().unwrap(); // [T, value_dim]
8239        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
8240
8241        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
8242        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
8243        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
8244        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
8245        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
8246        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
8247        let _ = (head_k, head_v);
8248        let mut q_g = e.uninit(d_state * num_v * t)?;
8249        let mut k_g = e.uninit(d_state * num_v * t)?;
8250        let mut v_g = e.uninit(d_state * num_v * t)?;
8251        e.ssm_conv1d_gdn(
8252            &qkv_mixed,
8253            la.ssm_conv1d.float_data(),
8254            &mut q_g,
8255            &mut k_g,
8256            &mut v_g,
8257            conv_dim,
8258            t,
8259            d_conv,
8260            d_state,
8261            num_v,
8262            num_k,
8263            key_dim,
8264        )?;
8265        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
8266        let mut q_l2 = e.uninit(d_state * num_v * t)?;
8267        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8268        let mut k_l2 = e.uninit(d_state * num_v * t)?;
8269        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8270        let v_gd = v_g;
8271
8272        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
8273        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
8274        let mut beta = e.uninit(t * num_v)?;
8275        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8276        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
8277        let mut g_log = e.uninit(t * num_v)?;
8278        e.gdn_glog(
8279            &alpha,
8280            la.ssm_dt.float_data(),
8281            la.ssm_a.float_data(),
8282            &mut g_log,
8283            num_v,
8284            t,
8285        )?;
8286
8287        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
8288        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
8289        let mut state_out = e.zeros(d_state * d_state * num_v)?;
8290        let mut o = e.uninit(d_state * num_v * t)?;
8291        e.gdn_scan_prefill(
8292            &q_l2,
8293            &k_l2,
8294            &v_gd,
8295            &g_log,
8296            &beta,
8297            None,
8298            None,
8299            &state_in,
8300            &mut state_out,
8301            &mut o,
8302            num_v,
8303            t,
8304            scale,
8305            num_v,
8306        )?;
8307
8308        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
8309        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
8310        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
8311        // o rows are (t*num_v+vh) too. Good.
8312        let mut gn = e.uninit(d_state * num_v * t)?;
8313        e.gated_rmsnorm(
8314            &o,
8315            la.ssm_norm.float_data(),
8316            &z,
8317            &mut gn,
8318            d_state,
8319            num_v * t,
8320            eps,
8321        )?;
8322
8323        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
8324        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
8325        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
8326        let out = e.matmul(&la.ssm_out, &gn, t)?;
8327        Ok(out)
8328    }
8329}
8330
8331impl HybridModel {
8332    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
8333    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
8334    ///
8335    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
8336    /// different 860160-byte block than the same expert of layer 7).
8337    ///
8338    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
8339    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
8340    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
8341    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
8342    pub fn moe_ffn_il(
8343        &self,
8344        e: &Engine,
8345        m: &MoeWeights,
8346        z: &CudaSlice<f32>,
8347        t: usize,
8348        il: u16,
8349    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8350        Self::moe_ffn_inner(
8351            e,
8352            m,
8353            z,
8354            None,
8355            t,
8356            &self.cfg,
8357            il,
8358            self.max_moe_block(),
8359            false,
8360            None,
8361            self.uses_sliding_gated_moe_program(),
8362            false,
8363        )
8364    }
8365
8366    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
8367    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
8368    pub fn moe_ffn_il_prefill(
8369        &self,
8370        e: &Engine,
8371        m: &MoeWeights,
8372        z: &CudaSlice<f32>,
8373        t: usize,
8374        il: u16,
8375    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8376        Self::moe_ffn_inner(
8377            e,
8378            m,
8379            z,
8380            None,
8381            t,
8382            &self.cfg,
8383            il,
8384            self.max_moe_block(),
8385            true,
8386            Some(&self.step_grouped_prefill),
8387            self.uses_sliding_gated_moe_program(),
8388            false,
8389        )
8390    }
8391
8392    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
8393    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
8394    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
8395    pub fn moe_ffn_il_zq8(
8396        &self,
8397        e: &Engine,
8398        m: &MoeWeights,
8399        z: &CudaSlice<f32>,
8400        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8401        t: usize,
8402        il: u16,
8403    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8404        Self::moe_ffn_inner(
8405            e,
8406            m,
8407            z,
8408            zq8,
8409            t,
8410            &self.cfg,
8411            il,
8412            self.max_moe_block(),
8413            false,
8414            None,
8415            self.uses_sliding_gated_moe_program(),
8416            false,
8417        )
8418    }
8419
8420    /// Verify-rows twin of [`Self::moe_ffn_il_zq8`] (lane/glm5-vrest): the SAME routing and
8421    /// dispatch decisions with the pairs-shaped batched routed-expert arm armed. Only the
8422    /// verify walk's batched arm (`MEMRA_GLM5_VERIFY_BATCH`, t>=2) calls this; every
8423    /// unqualified shape inside falls closed to the byte-identical sequential loop.
8424    pub(crate) fn moe_ffn_il_zq8_vrows(
8425        &self,
8426        e: &Engine,
8427        m: &MoeWeights,
8428        z: &CudaSlice<f32>,
8429        t: usize,
8430        il: u16,
8431    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8432        Self::moe_ffn_inner(
8433            e,
8434            m,
8435            z,
8436            None,
8437            t,
8438            &self.cfg,
8439            il,
8440            self.max_moe_block(),
8441            false,
8442            None,
8443            self.uses_sliding_gated_moe_program(),
8444            true,
8445        )
8446    }
8447
8448    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
8449    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
8450    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
8451    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
8452    ///
8453    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
8454    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
8455    pub(crate) fn moe_ffn(
8456        e: &Engine,
8457        m: &MoeWeights,
8458        z: &CudaSlice<f32>,
8459        t: usize,
8460        cfg: &ModelConfig,
8461        il: u16,
8462        max_block: usize,
8463    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8464        Self::moe_ffn_inner(
8465            e, m, z, None, t, cfg, il, max_block, false, None, false, false,
8466        )
8467    }
8468
8469    #[allow(clippy::too_many_arguments)]
8470    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
8471    pub(crate) fn moe_ffn_inner(
8472        e: &Engine,
8473        m: &MoeWeights,
8474        z: &CudaSlice<f32>,
8475        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
8476        t: usize,
8477        cfg: &ModelConfig,
8478        il: u16,
8479        max_block: usize,
8480        prefill: bool,
8481        grouped_prefill: Option<&std::sync::Mutex<crate::hybrid::StepEpGroupedPrefill>>,
8482        sliding_gated_moe: bool,
8483        vrows: bool,
8484    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8485        let worker_io = crate::spill_pread::worker_enabled();
8486        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
8487        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
8488            e.with_moe_cache(max_block, |cache, _| {
8489                cache.begin_forward_epoch(il, t);
8490                if worker_io {
8491                    cache.begin_worker_scope();
8492                }
8493                Ok(())
8494            })?;
8495        }
8496        if let Some(ep) = &m.glm5_ep {
8497            // glm5 TP-2 EP walk (MEMRA_GLM5_TP): whole-expert halves, root router, slot-ordered
8498            // canonical combine. Every other arm of this function is unreachable for an
8499            // EP-armed layer by construction. `prefill` keys the EP grouped-prime arm
8500            // (MEMRA_GLM5_EP_GROUPED_PRIME) exactly as it keys the plain grouped arm below.
8501            return Self::moe_ffn_glm5_ep(e, m, ep, z, zq8, t, cfg, il, prefill);
8502        }
8503        if m.step_ep.is_some() || m.step_tp.is_some() {
8504            let moe = cfg
8505                .moe
8506                .as_ref()
8507                .ok_or("Step distributed execution requires MoE model metadata")?;
8508            let n_embd = cfg.n_embd as usize;
8509            let n_expert = moe.expert_count as usize;
8510            let n_used = moe.expert_used_count as usize;
8511            let sigmoid = cfg
8512                .sigmoid_router()
8513                .ok_or("Step distributed execution requires the Step sigmoid router")?;
8514            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
8515            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
8516            let grouped_prefill_requested = prefill && step_ep_grouped_prefill_enabled()?;
8517            if grouped_prefill_requested && !step_tp_prefill_enabled()? {
8518                return Err(
8519                    "MEMRA_STEP_EP_GROUPED_PREFILL=1 requires MEMRA_STEP_TP_PREFILL=1".into(),
8520                );
8521            }
8522            if grouped_prefill_requested && !step_grouped_prefill_shape(true, prefill, t) {
8523                return Err(format!(
8524                    "Step grouped prefill tokens {t} are outside the qualified {}..={} range",
8525                    PRIME_MIN_T,
8526                    crate::cache::PRIME_CHUNK_MAX_TOKENS,
8527                )
8528                .into());
8529            }
8530            let grouped_decode_shape = step_grouped_decode_shape(prefill, t);
8531            let grouped_prefill_shape =
8532                step_grouped_prefill_shape(grouped_prefill_requested, prefill, t);
8533            if let Some(ep) = m.step_ep.as_ref().filter(|ep| {
8534                ep.grouped_decode.is_some() && (grouped_decode_shape || grouped_prefill_shape)
8535            }) {
8536                let (selected, route_weights) =
8537                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8538                crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
8539                Self::trace_moe_routes(il, t, &selected, &route_weights)?;
8540                Self::trace_moe_input(e, il, t, n_embd, z)?;
8541                let selected = selected
8542                    .iter()
8543                    .map(|&expert| expert as usize)
8544                    .collect::<Vec<_>>();
8545
8546                // The narrow route readback above orders the owning-stage producer. The grouped
8547                // runtime then copies the resident root activation into its persistent rank inputs.
8548                e.stream().synchronize()?;
8549                let execute = |state: &mut crate::hybrid::StepEpGroupedDecode| {
8550                    state.projection.set_activation_limit(ep.activation_limit)?;
8551                    ep.runtime
8552                        .refresh_step_grouped_expert_parallel_gate_from_root_device(
8553                            ep.experts.e4m3()?,
8554                            &mut state.projection,
8555                            z,
8556                            t,
8557                            &selected,
8558                        )?;
8559                    ep.runtime.refresh_step_grouped_expert_parallel_combine(
8560                        &state.projection,
8561                        &mut state.combine,
8562                        &route_weights,
8563                    )?;
8564                    ep.runtime.execute_step_grouped_expert_parallel_gate(
8565                        ep.experts.e4m3()?,
8566                        &mut state.projection,
8567                    )?;
8568                    ep.runtime.execute_step_grouped_expert_parallel_combine(
8569                        &state.projection,
8570                        &mut state.combine,
8571                    )?;
8572                    let mut output = ep.runtime.copy_step_grouped_expert_parallel_combine_root(
8573                        &state.projection,
8574                        &state.combine,
8575                        e,
8576                    )?;
8577                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8578                    if prefill {
8579                        // A shared plan may be reused by the next layer on a different runtime
8580                        // stream. Complete the owning-stage copy before its source is overwritten.
8581                        e.stream().synchronize()?;
8582                    }
8583                    eprintln!(
8584                        "[step-tp-ep-grouped] execute layer={il} tokens={t} devices={:?} \
8585                         attention_layout=tensor-parallel expert_layout=expert-parallel \
8586                         expert_transport={} native_p2p=true route_control=host-narrow \
8587                         input=root-device projection_workspaces=persistent \
8588                         combine=root-device output=owning-stage-device \
8589                         prefill={prefill} batched_decode=false capacity={} \
8590                         performance_claim=false",
8591                        ep.devices,
8592                        ep.runtime.transport_label(),
8593                        state.projection.max_tokens(),
8594                    );
8595                    Ok::<_, Box<dyn std::error::Error>>(output)
8596                };
8597
8598                if grouped_prefill_shape {
8599                    let grouped_prefill = grouped_prefill
8600                        .ok_or("Step grouped prefill has no model-scoped executor")?;
8601                    let mut shared = grouped_prefill
8602                        .lock()
8603                        .map_err(|_| "Step grouped prefill state lock is poisoned")?;
8604                    let needs_prepare = shared.state.as_ref().is_none_or(|state| {
8605                        state.devices != ep.devices
8606                            || state.grouped.projection.max_tokens() < t
8607                            || state.grouped.projection.input_width() != n_embd
8608                            || state.grouped.projection.expert_width()
8609                                != moe.expert_ff_length as usize
8610                    });
8611                    if needs_prepare {
8612                        let seed_input = vec![0.0f32; n_embd];
8613                        let seed_selected = &selected[..n_used];
8614                        let seed_weights = &route_weights[..n_used];
8615                        let projection = ep
8616                            .runtime
8617                            .prepare_step_grouped_expert_parallel_gate_with_capacity(
8618                                ep.experts.e4m3()?,
8619                                &seed_input,
8620                                1,
8621                                seed_selected,
8622                                ep.activation_limit,
8623                                t,
8624                            )?;
8625                        let combine = ep.runtime.prepare_step_grouped_expert_parallel_combine(
8626                            &projection,
8627                            seed_weights,
8628                        )?;
8629                        shared.state = Some(crate::hybrid::StepEpGroupedPrefillState {
8630                            devices: ep.devices.clone(),
8631                            grouped: crate::hybrid::StepEpGroupedDecode {
8632                                projection,
8633                                combine,
8634                            },
8635                        });
8636                        eprintln!(
8637                            "[step-tp-ep-grouped-prefill] prepare capacity={t} devices={:?} \
8638                             shared_across_layers=true performance_claim=false",
8639                            ep.devices,
8640                        );
8641                    }
8642                    return execute(
8643                        &mut shared
8644                            .state
8645                            .as_mut()
8646                            .expect("Step grouped prefill state prepared above")
8647                            .grouped,
8648                    );
8649                }
8650
8651                let mut grouped = ep
8652                    .grouped_decode
8653                    .as_ref()
8654                    .expect("grouped decode presence checked above")
8655                    .lock()
8656                    .map_err(|_| "Step grouped decode state lock is poisoned")?;
8657                return execute(&mut grouped);
8658            }
8659            if grouped_prefill_shape {
8660                return Err(
8661                    "Step grouped prefill requires native-P2P expert-owner device arithmetic"
8662                        .into(),
8663                );
8664            }
8665            // MEMRA_STEP_TP_DEV_ROUTER=1 (t=1): route on device and feed the device-routed
8666            // expert program — the per-layer host logits readback (the last per-layer host
8667            // sync) disappears. Selection tie-breaking may differ from the host router:
8668            // numeric-class door, run-gen argmax gate + boot battery.
8669            // STEP TP2 GEMM PRIME (2026-08-27, TTFT lane): a prime chunk's routed MoE goes
8670            // through ONE grouped f16 GEMM per projection over the resident NVFP4 banks —
8671            // the per-token device routes below cost 240 s at m=4092 (measured), the grouped
8672            // lane's sizing rows run 170-270 TFLOP/s. Router selections come from the same
8673            // sigmoid host oracle the EP arm uses; shexp rides the canonical grouped add.
8674            // t>=16 alone keys the branch: the batch prime reaches here through moe_ffn_il,
8675            // whose `prefill` is FALSE (only the _prefill twin sets it), and no other step37
8676            // route runs t>=16 — verify walks t<=8, decode t=1. Requiring `prefill` made the
8677            // first gate arm skip this branch entirely and wake the generic f16g arm instead
8678            // (48 s + kq_gemm_sk rc=1001, 2026-08-27).
8679            if t >= 16
8680                && crate::step_gemm_prime_on()
8681                && let Some(tp) = &m.step_tp
8682                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8683            {
8684                // MEMRA_PRIME_PROF=1 sub-split of the moe bucket. The phase timer put
8685                // 1788 ms of a 3093 ms chunk here, but forcing the 32-row tile form (4x
8686                // more weight dequant) moved it only 5% — so the grouped GEMM is not
8687                // obviously what dominates. The router below is a HOST oracle: sigmoid +
8688                // top-8 over 288 experts for every one of 4096 tokens, per layer, which
8689                // is a D2H copy and a full pipeline drain 42 times per chunk. Attribute
8690                // it before optimizing the kernel it sits in front of.
8691                let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1") && t >= 16;
8692                let mut mt = std::time::Instant::now();
8693                let (selected, route_weights) =
8694                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
8695                let sel_i32: Vec<i32> = selected.iter().map(|&x| x as i32).collect();
8696                let d_router = if mprof {
8697                    let _ = e.stream().synchronize();
8698                    let v = mt.elapsed().as_secs_f64() * 1e3;
8699                    mt = std::time::Instant::now();
8700                    v
8701                } else {
8702                    0.0
8703                };
8704                // MEMRA_MOE_DETERM=1: run the WHOLE grouped routine twice on identical
8705                // inputs and diff. The standalone harness cleared the grouped GEMM kernels
8706                // (8 invocations, both lanes, maxdiff 0.0 over 20.9M elements) but it does
8707                // not model the cross-device join/scatter or the o_proj-style reduction,
8708                // and the loader refuses both topologies (TP1, same-device) that would
8709                // isolate those by env. This tests the un-excluded region directly, in
8710                // the place it actually runs.
8711                //
8712                // The prime is nondeterministic: same prompt, one forward, temperature=0,
8713                // max_tokens=1, and the first token varies across reps. That blocks
8714                // MEMRA_PP_BF16's correctness receipt and invalidates every byte-identity
8715                // gate taken through the server. This probe also yields the jitter
8716                // MAGNITUDE, which any tolerance band needs.
8717                let mdet =
8718                    std::env::var("MEMRA_MOE_DETERM").as_deref() == Ok("1") && t >= 16 && il < 4;
8719                if mdet {
8720                    let a = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8721                        bank,
8722                        e,
8723                        z,
8724                        t,
8725                        &sel_i32,
8726                        &route_weights,
8727                        n_used,
8728                        tp.activation_limit,
8729                    )?;
8730                    let b = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8731                        bank,
8732                        e,
8733                        z,
8734                        t,
8735                        &sel_i32,
8736                        &route_weights,
8737                        n_used,
8738                        tp.activation_limit,
8739                    )?;
8740                    let (ha, hb) = (e.dtoh(&a)?, e.dtoh(&b)?);
8741                    let mut md = 0.0f32;
8742                    let mut ndiff = 0usize;
8743                    for (x, y) in ha.iter().zip(hb.iter()) {
8744                        let d = (x - y).abs();
8745                        if d > 0.0 {
8746                            ndiff += 1;
8747                        }
8748                        if d > md {
8749                            md = d;
8750                        }
8751                    }
8752                    eprintln!(
8753                        "[moe-determ] il={il} t={t} maxdiff={md:.3e} \
8754                                 differing={ndiff}/{} -> {}",
8755                        ha.len(),
8756                        if ndiff == 0 {
8757                            "IDENTICAL"
8758                        } else {
8759                            "NONDETERMINISTIC"
8760                        }
8761                    );
8762                }
8763                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_prime_grouped(
8764                    bank,
8765                    e,
8766                    z,
8767                    t,
8768                    &sel_i32,
8769                    &route_weights,
8770                    n_used,
8771                    tp.activation_limit,
8772                )?;
8773                let d_gemm = if mprof {
8774                    let _ = e.stream().synchronize();
8775                    let v = mt.elapsed().as_secs_f64() * 1e3;
8776                    mt = std::time::Instant::now();
8777                    v
8778                } else {
8779                    0.0
8780                };
8781                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
8782                if mprof {
8783                    let _ = e.stream().synchronize();
8784                    let d_shared = mt.elapsed().as_secs_f64() * 1e3;
8785                    // Per LAYER, not accumulated: the four trunk phases already carry the
8786                    // per-chunk totals, and one line per layer is what shows whether the
8787                    // cost is flat across layers or concentrated in a few.
8788                    eprintln!(
8789                        "[moe-prof] il={il} t={t} router={d_router:.1}ms \
8790                                 gemm={d_gemm:.1}ms shared={d_shared:.1}ms"
8791                    );
8792                }
8793                return Ok(output);
8794            }
8795            if t == 1
8796                && crate::tp::step_nvfp4_dev_routes_enabled()?
8797                && crate::tp::step_tp_dev_router_enabled()?
8798                && let Some(tp) = &m.step_tp
8799                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
8800            {
8801                let (sf, route_norm) = sigmoid;
8802                // MEMRA_ROUTES_PRESTAGE=1: issue the rank input pull + quantize NOW,
8803                // before the router — the rank streams overlap the gemv+topk.
8804                // MEMRA_DEV1_ROUTER=1 rides the prestage hook: rank1 routes LOCALLY
8805                // from its own z copy (replicated deterministic router — identical
8806                // bits in, identical sel/w out) and starts its sweep without
8807                // waiting the root's sel broadcast.
8808                static D1_ROUTER: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8809                let d1_router = *D1_ROUTER
8810                    .get_or_init(|| std::env::var("MEMRA_DEV1_ROUTER").as_deref() == Ok("1"));
8811                if d1_router {
8812                    let (sf_h, rn_h) = sigmoid;
8813                    let n_ex = m.gate_exps.n_expert;
8814                    let act_ct = m.active_count();
8815                    let _ = tp.runtime.nvfp4_routes_prestage_with(
8816                        bank,
8817                        e,
8818                        z,
8819                        |rank1, in1, sel1, w1| {
8820                            let mut guard = DEV1_ROUTER_REPS
8821                                .lock()
8822                                .map_err(|_| "dev1 router replica lock")?;
8823                            let (reps, scratch) =
8824                                guard.get_or_insert_with(|| (Default::default(), None));
8825                            if !reps.contains_key(&il) {
8826                                use cudarc::driver::DevicePtr;
8827                                let (g1, p1, a1) = (
8828                                    rank1.htod(&vec![0.0f32; n_ex * n_embd])?,
8829                                    rank1.htod(&vec![0.0f32; n_ex])?,
8830                                    rank1.alloc_u8_uninit(n_ex)?,
8831                                );
8832                                for (src, dst_len, dst) in [
8833                                    (
8834                                        {
8835                                            let s = e.stream();
8836                                            let (p, _g) = m.gate_inp.float_data().device_ptr(&s);
8837                                            p
8838                                        },
8839                                        n_ex * n_embd * 4,
8840                                        {
8841                                            let s = rank1.stream();
8842                                            let (p, _g) = g1.device_ptr(&s);
8843                                            p
8844                                        },
8845                                    ),
8846                                    (
8847                                        {
8848                                            let s = e.stream();
8849                                            let (p, _g) = m.exp_probs_b_dev.device_ptr(&s);
8850                                            p
8851                                        },
8852                                        n_ex * 4,
8853                                        {
8854                                            let s = rank1.stream();
8855                                            let (p, _g) = p1.device_ptr(&s);
8856                                            p
8857                                        },
8858                                    ),
8859                                    (
8860                                        {
8861                                            let s = e.stream();
8862                                            let (p, _g) = m.active_experts_dev.device_ptr(&s);
8863                                            p
8864                                        },
8865                                        n_ex,
8866                                        {
8867                                            let s = rank1.stream();
8868                                            let (p, _g) = a1.device_ptr(&s);
8869                                            p
8870                                        },
8871                                    ),
8872                                ] {
8873                                    crate::tp::raw_copy_bytes(dst, src, dst_len, rank1)?;
8874                                }
8875                                rank1.stream().synchronize()?;
8876                                reps.insert(il, (g1, p1, a1));
8877                            }
8878                            if scratch.is_none() {
8879                                *scratch = Some(rank1.htod(&vec![0.0f32; n_ex])?);
8880                            }
8881                            let (g1, p1, a1) = reps.get(&il).expect("armed above");
8882                            let logits1 = scratch.as_mut().expect("armed above");
8883                            rank1.router_gemv_into(g1, in1, logits1, n_embd, n_ex, 1)?;
8884                            rank1.moe_router_sigmoid_topk_into(
8885                                logits1, 1, n_ex, n_used, act_ct, p1, a1, sf_h, rn_h, sel1, w1,
8886                            )?;
8887                            Ok(true)
8888                        },
8889                    )?;
8890                } else {
8891                    let _ = tp.runtime.nvfp4_routes_prestage(bank, e, z)?;
8892                }
8893                // Persistent selection buffers: the allocating topk built two fresh
8894                // slices per layer; sel/w land in process-static rows instead
8895                // (host-op diet — same kernel, same bytes).
8896                #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
8897                static SELW: std::sync::Mutex<Option<(usize, CudaSlice<i32>, CudaSlice<f32>)>> =
8898                    std::sync::Mutex::new(None);
8899                let mut selw = SELW.lock().map_err(|_| "selw lock poisoned")?;
8900                if selw.as_ref().is_none_or(|(d, ..)| *d != e.ctx().ordinal()) {
8901                    *selw = Some((
8902                        e.ctx().ordinal(),
8903                        e.htod_i32(&vec![0i32; n_used])?,
8904                        e.htod(&vec![0.0f32; n_used])?,
8905                    ));
8906                }
8907                let (_, sel_d, w_d) = selw.as_mut().expect("armed above");
8908                e.moe_router_sigmoid_topk_into(
8909                    &logits,
8910                    t,
8911                    n_expert,
8912                    n_used,
8913                    m.active_count(),
8914                    &m.exp_probs_b_dev,
8915                    &m.active_experts_dev,
8916                    sf,
8917                    route_norm,
8918                    sel_d,
8919                    w_d,
8920                )?;
8921                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
8922                // FAIL-CLOSED for the route taps: this walk keeps the selection
8923                // device-side, so `trace_moe_routes` (MEMRA_MOE_TRACE /
8924                // MEMRA_MOE_WEIGHT_TRACE) never sees its rows. Every other MoE walk is
8925                // either host-routed (the taps ride the existing readback) or forced to
8926                // the host-visible path by observation mode — this one is neither. A
8927                // trace that silently misses whole layers poisons any placement mint
8928                // built on it (LAW:coactivation-expert-placement measurement leg), so an
8929                // armed tap refuses by name instead of dropping rows.
8930                if std::env::var("MEMRA_MOE_TRACE").is_ok()
8931                    || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
8932                {
8933                    return Err("MEMRA_MOE_TRACE/MEMRA_MOE_WEIGHT_TRACE cannot trace the \
8934                         device-routed step TP walk (selection never returns to host; \
8935                         tracing would add a new sync). Route through the host-router \
8936                         arm — refused rather than silently dropping rows"
8937                        .into());
8938                }
8939                // MEMRA_SHEXP_OVERLAP=1: issue the shared expert from the routes
8940                // PREJOIN hook so it executes while the peer rank drains its sweep
8941                // (fills dev0's join wait); apply adds the identical values after.
8942                static SHEXP_OV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8943                let shexp_ov = *SHEXP_OV
8944                    .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
8945                // MEMRA_SHEXP_DEV1=1 (supersedes the dev0 overlap): the shared
8946                // expert runs on rank1 — the idle device — same kernels, same
8947                // split program, down row root-resident: bit-identical.
8948                static SHEXP_D1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8949                let shexp_d1 = *SHEXP_D1
8950                    .get_or_init(|| std::env::var("MEMRA_SHEXP_DEV1").as_deref() == Ok("1"))
8951                    && tp.runtime.rank_engine(1).is_some();
8952                // MOE TAIL FUSION M1 (MEMRA_TAIL_ADD3=0 reverts): pre-arm the
8953                // overlap ws + ones row and hand their RAW pointers to the routed
8954                // run — the join add folds the shexp apply into one launch.
8955                static TAIL3: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8956                let tail3 =
8957                    *TAIL3.get_or_init(|| std::env::var("MEMRA_TAIL_ADD3").as_deref() != Ok("0"));
8958                let mut ov_issued = false;
8959                let mut d1_issued = false;
8960                let mut tail_folded = false;
8961                let mut output = if shexp_d1 {
8962                    let rank1 = tp.runtime.rank_engine(1).expect("checked above");
8963                    tp.runtime
8964                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin(
8965                            bank,
8966                            e,
8967                            z,
8968                            sel_d,
8969                            w_d,
8970                            n_used,
8971                            tp.activation_limit,
8972                            || {
8973                                d1_issued =
8974                                    Self::shexp_dev1_issue(e, rank1, m, z, cfg, il, n_embd)?;
8975                                Ok(())
8976                            },
8977                        )?
8978                } else if shexp_ov {
8979                    // Raw sh/ones pointers for the fused tail (persistent statics;
8980                    // pointers stable, no lock held across the routed call). The
8981                    // sh CONTENT is written by the prejoin-issued kernels earlier
8982                    // on e's stream — stream order covers the fused add.
8983                    let post_add = if tail3 {
8984                        Self::shexp_overlap_tail_ptrs(e, m, cfg, n_embd)?
8985                    } else {
8986                        None
8987                    };
8988                    let used_post = post_add.is_some();
8989                    let out = tp
8990                        .runtime
8991                        .run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3(
8992                            bank,
8993                            e,
8994                            z,
8995                            sel_d,
8996                            w_d,
8997                            n_used,
8998                            tp.activation_limit,
8999                            || {
9000                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9001                                Ok(())
9002                            },
9003                            post_add,
9004                        )?;
9005                    // ov_issued false with post_add armed = an early-return arm
9006                    // (the GRAPH door) skipped the prejoin AND ignored post_add —
9007                    // fall through to the normal shexp add (battery v22 receipt:
9008                    // the strict error here failed every graph-door boot).
9009                    if used_post && ov_issued {
9010                        tail_folded = true; // apply folded into the join add
9011                    }
9012                    out
9013                } else {
9014                    tp.runtime.run_tensor_parallel_routes_nvfp4_device_routed(
9015                        bank,
9016                        e,
9017                        z,
9018                        sel_d,
9019                        w_d,
9020                        n_used,
9021                        tp.activation_limit,
9022                    )?
9023                };
9024                if output.len() != t * n_embd {
9025                    return Err(format!(
9026                        "Step tp routed output has {} values, expected {t}x{n_embd}",
9027                        output.len()
9028                    )
9029                    .into());
9030                }
9031                if tail_folded {
9032                    // shexp already folded into the join add (MOE TAIL FUSION M1)
9033                } else if d1_issued {
9034                    Self::shexp_dev1_apply(e, &mut output, n_embd)?;
9035                } else if ov_issued {
9036                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9037                } else {
9038                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9039                }
9040                static DR_LOGGED: std::sync::atomic::AtomicU64 =
9041                    std::sync::atomic::AtomicU64::new(0);
9042                let layer_bit = 1u64 << (il as u64 % 64);
9043                if DR_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9044                    == 0
9045                {
9046                    eprintln!(
9047                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
9048                                 expert_transport={} native_p2p={} router=device \
9049                                 activation=host-canonical accumulation=host-canonical \
9050                                 output=e-device io=device performance_claim=false \
9051                                 (logged once per layer)",
9052                        tp.devices,
9053                        tp.runtime.transport_label(),
9054                        tp.runtime.native_p2p(),
9055                    );
9056                }
9057                return Ok(output);
9058            }
9059            let automatic_ep_device_router = crate::tp::parallel_ep_device_router_enabled()?;
9060            let automatic_ep_q8_act = crate::tp::parallel_ep_q8_act_enabled()?;
9061            let automatic_ep_q8_scope = crate::tp::parallel_ep_q8_scope()?;
9062            let automatic_ep_q8_active =
9063                automatic_ep_q8_act && t <= crate::tp::NVFP4_EP_Q8_BATCH_CAP;
9064            if automatic_ep_q8_scope.is_some() && !automatic_ep_q8_act {
9065                return Err(
9066                    "MEMRA_PARALLEL_EP_Q8_SCOPE requires MEMRA_PARALLEL_EP_Q8_ACT=1".into(),
9067                );
9068            }
9069            if automatic_ep_q8_act && !automatic_ep_device_router {
9070                return Err(
9071                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires MEMRA_PARALLEL_EP_DEVICE_ROUTER=1".into(),
9072                );
9073            }
9074            if automatic_ep_q8_act && m.step_ep.as_ref().is_none_or(|ep| !ep.nvfp4_device_routes) {
9075                return Err(
9076                    "MEMRA_PARALLEL_EP_Q8_ACT=1 requires automatic W4A16 whole-expert EP".into(),
9077                );
9078            }
9079            if t <= crate::tp::NVFP4_EP_DEVICE_ROUTER_BATCH_CAP
9080                && automatic_ep_device_router
9081                && let Some(ep) = &m.step_ep
9082                && ep.nvfp4_device_routes
9083            {
9084                let bank = match &ep.experts {
9085                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9086                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
9087                        return Err("W4A16 device-routed EP reached an E4M3 expert bank".into());
9088                    }
9089                };
9090                let pairs = t
9091                    .checked_mul(n_used)
9092                    .ok_or("W4A16 device-routed EP pair count overflow")?;
9093                let capacity = crate::tp::NVFP4_EP_DEVICE_BATCH_CAP * n_used;
9094                /// Per-device persistent route scratch: device ordinal -> (armed capacity in
9095                /// pairs, selected-expert rows, route-weight rows). Named because the nested
9096                /// form is unreadable at this depth, not to hide it.
9097                type EpSelwByDevice =
9098                    std::collections::HashMap<usize, (usize, CudaSlice<i32>, CudaSlice<f32>)>;
9099                static EP_SELW: std::sync::Mutex<Option<EpSelwByDevice>> =
9100                    std::sync::Mutex::new(None);
9101                let mut selw = EP_SELW
9102                    .lock()
9103                    .map_err(|_| "automatic EP device-router workspace lock poisoned")?;
9104                let device = e.ctx().ordinal();
9105                let workspaces = selw.get_or_insert_with(Default::default);
9106                if workspaces
9107                    .get(&device)
9108                    .is_none_or(|(cap, ..)| *cap < capacity)
9109                {
9110                    workspaces.insert(
9111                        device,
9112                        (
9113                            capacity,
9114                            e.htod_i32(&vec![0i32; capacity])?,
9115                            e.htod(&vec![0.0f32; capacity])?,
9116                        ),
9117                    );
9118                }
9119                let (_, sel_d, w_d) = workspaces.get_mut(&device).expect("armed above");
9120                let (sf, route_norm) = sigmoid;
9121                e.moe_router_sigmoid_topk_into(
9122                    &logits,
9123                    t,
9124                    n_expert,
9125                    n_used,
9126                    m.active_count(),
9127                    &m.exp_probs_b_dev,
9128                    &m.active_experts_dev,
9129                    sf,
9130                    route_norm,
9131                    sel_d,
9132                    w_d,
9133                )?;
9134                crate::moesd::record_device_routes(e, il, n_expert, n_used, sel_d)?;
9135                static SHEXP_OV_AUTO: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9136                let shexp_ov = t == 1
9137                    && *SHEXP_OV_AUTO
9138                        .get_or_init(|| std::env::var("MEMRA_SHEXP_OVERLAP").as_deref() == Ok("1"));
9139                let mut ov_issued = false;
9140                let mut output = if shexp_ov {
9141                    ep.runtime
9142                        .run_routed_experts_nvfp4_w4a16_device_routed_prejoin(
9143                            bank,
9144                            e,
9145                            z,
9146                            sel_d,
9147                            w_d,
9148                            t,
9149                            n_used,
9150                            ep.activation_limit,
9151                            || {
9152                                ov_issued = Self::shexp_overlap_issue(e, m, z, cfg, il, n_embd)?;
9153                                Ok(())
9154                            },
9155                        )?
9156                } else {
9157                    ep.runtime.run_routed_experts_nvfp4_w4a16_device_routed(
9158                        bank,
9159                        e,
9160                        z,
9161                        sel_d,
9162                        w_d,
9163                        t,
9164                        n_used,
9165                        ep.activation_limit,
9166                    )?
9167                };
9168                if output.len() != t * n_embd {
9169                    return Err(format!(
9170                        "W4A16 device-routed EP output has {} values, expected \
9171                         {t}x{n_embd}={}",
9172                        output.len(),
9173                        t * n_embd,
9174                    )
9175                    .into());
9176                }
9177                if ov_issued {
9178                    Self::shexp_overlap_apply(e, &mut output, n_embd)?;
9179                } else {
9180                    Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9181                }
9182                static DEVICE_ROUTER_LOGGED: std::sync::atomic::AtomicU64 =
9183                    std::sync::atomic::AtomicU64::new(0);
9184                let layer_bit = 1u64 << (il as u64 % 64);
9185                if DEVICE_ROUTER_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9186                    & layer_bit
9187                    == 0
9188                {
9189                    eprintln!(
9190                        "[parallel-ep] execute layer={il} tokens={t} devices={:?} \
9191                         router=device expert_transport={} native_p2p={} \
9192                         activation=bf16-rounded accumulation={} output=e-device \
9193                         performance_claim=false (logged once per layer)",
9194                        ep.devices,
9195                        ep.runtime.transport_label(),
9196                        ep.runtime.native_p2p(),
9197                        if automatic_ep_q8_active {
9198                            "token-slot-order-q8"
9199                        } else {
9200                            "token-slot-order"
9201                        },
9202                    );
9203                }
9204                debug_assert!(pairs <= capacity);
9205                return Ok(output);
9206            }
9207            // MEMRA_STEP_TP_TIMING=1: cumulative cost of the host routing seam (the dtoh here
9208            // drains every e-stream op queued since the layer's FFN entry, so this bills the
9209            // router matmul + glue too — the decode-bucket ffn residue decomposes here).
9210            static ROUTE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9211            static ROUTE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9212            let route_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
9213            let route_started = route_timing.then(std::time::Instant::now);
9214            let (selected, route_weights, input) = Self::moe_route_sigmoid_with_input(
9215                e,
9216                &logits,
9217                z,
9218                t,
9219                n_embd,
9220                n_expert,
9221                n_used,
9222                m.exp_probs_b.as_deref(),
9223                sigmoid,
9224                m.active_experts.as_deref(),
9225            )?;
9226            if let Some(started) = route_started {
9227                use std::sync::atomic::Ordering;
9228                let ns = ROUTE_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
9229                    + started.elapsed().as_nanos() as u64;
9230                let calls = ROUTE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
9231                if calls.is_multiple_of(430) {
9232                    eprintln!(
9233                        "[moe-route-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
9234                        ns as f64 / 1.0e6,
9235                        ns as f64 / calls as f64 / 1.0e3,
9236                    );
9237                }
9238            }
9239            crate::moesd::record_host_routes(il, n_expert, n_used, &selected)?;
9240            Self::trace_moe_routes(il, t, &selected, &route_weights)?;
9241            Self::trace_moe_input(e, il, t, n_embd, z)?;
9242            let selected = selected
9243                .iter()
9244                .map(|&expert| expert as usize)
9245                .collect::<Vec<_>>();
9246            if t <= crate::tp::NVFP4_EP_DEVICE_BATCH_CAP
9247                && let Some(ep) = &m.step_ep
9248                && ep.nvfp4_device_routes
9249            {
9250                let bank = match &ep.experts {
9251                    crate::hybrid::StepEpExpertBank::Nvfp4(bank) => bank,
9252                    crate::hybrid::StepEpExpertBank::E4m3(_) => {
9253                        return Err("W4A16 NVFP4 device EP reached an E4M3 expert bank".into());
9254                    }
9255                };
9256                let mut output = ep.runtime.run_routed_experts_nvfp4_w4a16_device_io(
9257                    bank,
9258                    e,
9259                    z,
9260                    t,
9261                    &selected,
9262                    &route_weights,
9263                    n_used,
9264                    ep.activation_limit,
9265                )?;
9266                if output.len() != t * n_embd {
9267                    return Err(format!(
9268                        "W4A16 NVFP4 EP routed output has {} values, expected \
9269                         {t}x{n_embd}={}",
9270                        output.len(),
9271                        t * n_embd,
9272                    )
9273                    .into());
9274                }
9275                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9276                static W4A16_EP_LOGGED: std::sync::atomic::AtomicU64 =
9277                    std::sync::atomic::AtomicU64::new(0);
9278                let layer_bit = 1u64 << (il as u64 % 64);
9279                if W4A16_EP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed)
9280                    & layer_bit
9281                    == 0
9282                {
9283                    eprintln!(
9284                        "[step-ep] execute layer={il} tokens={t} devices={:?} \
9285                         expert_transport={} native_p2p={} activation=bf16-rounded \
9286                        accumulation={} output=e-device \
9287                         performance_claim=false (logged once per layer)",
9288                        ep.devices,
9289                        ep.runtime.transport_label(),
9290                        ep.runtime.native_p2p(),
9291                        if t == 1 {
9292                            "owner-grouped-rank-order"
9293                        } else {
9294                            "token-slot-order"
9295                        },
9296                    );
9297                }
9298                return Ok(output);
9299            }
9300            // Device-IO routes (t=1): the layer input goes to the ranks as a device row and the
9301            // combined output comes back as an e-context row — no host round-trip, no host
9302            // stream sync. Program bytes identical to the host-IO twin (dtoh/htod and dtod
9303            // both preserve f32 bits), gated by greedy token identity.
9304            if t == 1
9305                && crate::tp::step_nvfp4_dev_routes_enabled()?
9306                && let Some(tp) = &m.step_tp
9307                && let crate::hybrid::StepTpExpertBank::Nvfp4(bank) = &tp.experts
9308            {
9309                let mut output = tp.runtime.run_tensor_parallel_routes_nvfp4_device_io(
9310                    bank,
9311                    e,
9312                    z,
9313                    &selected,
9314                    &route_weights,
9315                    n_used,
9316                    tp.activation_limit,
9317                )?;
9318                if output.len() != t * n_embd {
9319                    return Err(format!(
9320                        "Step tp routed output has {} values, expected {t}x{n_embd}",
9321                        output.len()
9322                    )
9323                    .into());
9324                }
9325                Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9326                static IO_LOGGED: std::sync::atomic::AtomicU64 =
9327                    std::sync::atomic::AtomicU64::new(0);
9328                let layer_bit = 1u64 << (il as u64 % 64);
9329                if IO_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9330                    == 0
9331                {
9332                    eprintln!(
9333                        "[step-tp] execute layer={il} tokens={t} devices={:?} \
9334                                 expert_transport={} native_p2p={} activation=host-canonical \
9335                                 accumulation=host-canonical output=e-device io=device \
9336                                 performance_claim=false (logged once per layer)",
9337                        tp.devices,
9338                        tp.runtime.transport_label(),
9339                        tp.runtime.native_p2p(),
9340                    );
9341                }
9342                return Ok(output);
9343            }
9344            let (routed, mode, devices, transport, native_p2p) = if let Some(tp) = &m.step_tp {
9345                (
9346                    match &tp.experts {
9347                        crate::hybrid::StepTpExpertBank::E4m3(bank) => {
9348                            tp.runtime.run_tensor_parallel_routes(
9349                                bank,
9350                                &input,
9351                                t,
9352                                &selected,
9353                                &route_weights,
9354                                n_used,
9355                            )?
9356                        }
9357                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => {
9358                            if t == 1 && crate::tp::step_nvfp4_dev_routes_enabled()? {
9359                                tp.runtime.run_tensor_parallel_routes_nvfp4_device(
9360                                    bank,
9361                                    &input,
9362                                    &selected,
9363                                    &route_weights,
9364                                    n_used,
9365                                    tp.activation_limit,
9366                                )?
9367                            } else {
9368                                tp.runtime.run_tensor_parallel_routes_nvfp4(
9369                                    bank,
9370                                    &input,
9371                                    t,
9372                                    &selected,
9373                                    &route_weights,
9374                                    n_used,
9375                                    tp.activation_limit,
9376                                )?
9377                            }
9378                        }
9379                    },
9380                    "tp",
9381                    &tp.devices,
9382                    tp.runtime.transport_label(),
9383                    tp.runtime.native_p2p(),
9384                )
9385            } else {
9386                let ep = m
9387                    .step_ep
9388                    .as_ref()
9389                    .ok_or("Step distributed runtime has no EP or TP state")?;
9390                (
9391                    match &ep.experts {
9392                        crate::hybrid::StepEpExpertBank::E4m3(bank) => {
9393                            ep.runtime.run_routed_experts(
9394                                bank,
9395                                &input,
9396                                t,
9397                                &selected,
9398                                &route_weights,
9399                                n_used,
9400                                ep.activation_limit,
9401                            )?
9402                        }
9403                        crate::hybrid::StepEpExpertBank::Nvfp4(bank) => {
9404                            ep.runtime.run_routed_experts_nvfp4(
9405                                bank,
9406                                &input,
9407                                t,
9408                                &selected,
9409                                &route_weights,
9410                                n_used,
9411                                ep.activation_limit,
9412                            )?
9413                        }
9414                    },
9415                    if ep.configured_by_tp { "tp-ep" } else { "ep" },
9416                    &ep.devices,
9417                    ep.runtime.transport_label(),
9418                    ep.runtime.native_p2p(),
9419                )
9420            };
9421            if routed.len() != t * n_embd {
9422                return Err(format!(
9423                    "Step {mode} routed output has {} values, expected {t}x{n_embd}",
9424                    routed.len()
9425                )
9426                .into());
9427            }
9428            let mut output = e.htod(&routed)?;
9429            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut output)?;
9430            // Once per layer per process: the topology contract line is a boot receipt, not a
9431            // per-token trace — 4520 of these per 64-token run measured as real decode wall.
9432            static STEP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9433            let layer_bit = 1u64 << (il as u64 % 64);
9434            if STEP_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit
9435                == 0
9436            {
9437                eprintln!(
9438                    "[step-{mode}] execute layer={il} tokens={t} devices={devices:?} \
9439                     expert_transport={transport} native_p2p={native_p2p} \
9440                     activation={} accumulation={} output={} \
9441                     performance_claim=false (logged once per layer)",
9442                    if let Some(ep) = &m.step_ep {
9443                        ep.runtime.expert_activation_label()
9444                    } else {
9445                        "host-canonical"
9446                    },
9447                    if let Some(ep) = &m.step_ep {
9448                        ep.runtime.expert_accumulation_label()
9449                    } else {
9450                        "host-canonical"
9451                    },
9452                    if let Some(ep) = &m.step_ep {
9453                        ep.runtime.expert_output_label()
9454                    } else {
9455                        "host-accumulated"
9456                    },
9457                );
9458                if let Some(ep) = &m.step_ep
9459                    && let Some(limit) = ep.activation_limit
9460                {
9461                    eprintln!(
9462                        "[step-ep-clamp] execute layer={il} tokens={t} routed_clamp={limit} \
9463                             formula=min-silu-times-clamped-up performance_claim=false"
9464                    );
9465                }
9466            }
9467            return Ok(output);
9468        }
9469        if Self::sigmoid_resident_dev_eligible(e, m, cfg, sliding_gated_moe) {
9470            let moe = cfg.moe.as_ref().unwrap();
9471            let n_expert = moe.expert_count as usize;
9472            let n_used = moe.expert_used_count as usize;
9473            let sigmoid = cfg.sigmoid_router().unwrap();
9474            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9475            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
9476            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
9477        }
9478        // GROUPED MoE PREFILL, sigmoid-router class (glm5_next), MEMRA_MOE_GROUPED_PREFILL
9479        // default ON since 2026-08-29 (owner-accepted flip; =0 rollback seam; receipts on the
9480        // flag helper). The prefill-gap attribution (research/glm53-flash-bringup-20260827/
9481        // prefill-gap-20260829/PREFILL-GAP.md §1.1) measured this arch prefilling every prompt
9482        // token through the decode program: 49 launches per token-layer, ~8.4M launches and
9483        // 4.76 GB of expert-weight VRAM re-reads per token across 42 layers per 4096-token
9484        // chunk, because every batched arm is predicate-denied for sigmoid-router archs. This
9485        // arm is the composition of qualified ingredients: the m-invariant router + sigmoid
9486        // host oracle (routing sel/w BIT-identical to the sequential arm by construction),
9487        // host token-sort by expert (the moe_align_block_size shape), one grouped NVFP4
9488        // tensor-core GEMM per projection (the step37 grouped-prime kernel class via
9489        // `moe_f16_grouped`, generalized to the single-device resident slab), the PRE-clamped
9490        // SwiGLU epilogue and the per-expert weight_scale_2 macro fold the fused-epilogue lane
9491        // gated for this family. Keyed on `prefill` (only the _prefill twin sets it) so decode,
9492        // spec verify and the exact-16 batched-decode tier keep their dispatch class, and on
9493        // `t > MOE_DEV_MAX_T` so t<=16 stays on the per-token program (grouped/pairs prefill
9494        // classes start at 17, same seam as the softmax pairs arm).
9495        // ENGAGEMENT RECEIPT: the announce below prints in BOTH arms (flag on and off), once
9496        // per process, so an A/B grep distinguishes engagement without the line itself being
9497        // an arm-local cost (the step37 engagement-receipt trap: prove the path RAN before
9498        // attributing a number to it).
9499        if prefill && t > MOE_DEV_MAX_T && cfg.sigmoid_router().is_some() && cfg.glm5.is_some() {
9500            // Once per process PER FLAG VALUE (bit 0 = off, bit 1 = on): a server boot prints
9501            // exactly one line, and a gate process that flips the flag shows both arms.
9502            static GPF_ANNOUNCED: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
9503            let enabled = moe_grouped_prefill_enabled();
9504            let bit = 1u8 << u8::from(enabled);
9505            if GPF_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
9506                eprintln!(
9507                    "[moe-grouped-prefill] flag={} t={t} il={il} (announce printed in both \
9508                     arms; engagement is the per-layer execute line + the dispatch counter)",
9509                    if enabled { "on" } else { "off" },
9510                );
9511            }
9512            if enabled
9513                && let Some(out) = Self::moe_ffn_grouped_prefill_sigmoid(e, m, z, t, cfg, il)?
9514            {
9515                return Ok(out);
9516            }
9517        }
9518        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
9519        // current caller into this research arm; the naked default stays on the established path.
9520        if t > 1 && moe_grouped_enabled(cfg, prefill) {
9521            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
9522            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
9523            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
9524            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
9525            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
9526            if std::env::var("MEMRA_MOE_GATE").is_ok() {
9527                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
9528                let g_host = e.dtoh(&grouped_out)?;
9529                let s_host = e.dtoh(&seq_out)?;
9530                let g_bytes: &[u8] = unsafe {
9531                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
9532                };
9533                let s_bytes: &[u8] = unsafe {
9534                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
9535                };
9536                if g_bytes == s_bytes {
9537                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
9538                } else {
9539                    let diffs = g_host
9540                        .iter()
9541                        .zip(s_host.iter())
9542                        .enumerate()
9543                        .filter(|(_, (a, b))| a != b)
9544                        .count();
9545                    let maxdiff = g_host
9546                        .iter()
9547                        .zip(s_host.iter())
9548                        .map(|(a, b)| (a - b).abs())
9549                        .fold(0.0f32, f32::max);
9550                    panic!(
9551                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
9552                        g_host.len()
9553                    );
9554                }
9555            }
9556            return Ok(grouped_out);
9557        }
9558        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block, vrows)
9559    }
9560
9561    fn sigmoid_resident_dev_eligible(
9562        e: &Engine,
9563        m: &MoeWeights,
9564        cfg: &ModelConfig,
9565        sliding_gated_moe: bool,
9566    ) -> bool {
9567        let Some(moe) = cfg.moe.as_ref() else {
9568            return false;
9569        };
9570        // Cached once per process: this predicate runs per MoE layer per decode step, and five
9571        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
9572        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9573        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
9574            std::env::var("MEMRA_MOE_STATS").is_ok()
9575                || std::env::var("MEMRA_MOE_TRACE").is_ok()
9576                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9577                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
9578                || std::env::var("MEMRA_MOE_GATE").is_ok()
9579        });
9580        let resident_layout_supported = m.dev_exps.as_ref().is_some_and(|dev| {
9581            if dev.dev != e.ctx().ordinal() {
9582                return false;
9583            }
9584            let q8 = moe_q8_enabled_for_model(cfg, m);
9585            let fp8 = dev.fp8_blk.is_some()
9586                && m.gate_exps.qtype == crate::QT_F8_E4M3_BLK
9587                && m.up_exps.qtype == crate::QT_F8_E4M3_BLK
9588                && m.down_exps.qtype == crate::QT_F8_E4M3_BLK;
9589            q8 || fp8
9590        });
9591        sliding_gated_moe
9592            && sigmoid_router_enabled()
9593            && moe_dev_enabled()
9594            && moe_slab_enabled()
9595            && !observation_mode
9596            && moe.expert_used_count <= 8
9597            && m.has_uniform_expert_layout()
9598            && m.gate_exps.macros.is_none()
9599            && m.up_exps.macros.is_none()
9600            && m.down_exps.macros.is_none()
9601            && !m.has_macros
9602            && resident_layout_supported
9603    }
9604
9605    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
9606    pub(crate) fn moe_ffn_sequential(
9607        e: &Engine,
9608        m: &MoeWeights,
9609        z: &CudaSlice<f32>,
9610        t: usize,
9611        cfg: &ModelConfig,
9612        il: u16,
9613        max_block: usize,
9614    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9615        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block, false)
9616    }
9617
9618    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
9619    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
9620    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
9621    fn moe_router_logits(
9622        e: &Engine,
9623        m: &MoeWeights,
9624        z: &CudaSlice<f32>,
9625        t: usize,
9626        cfg: &ModelConfig,
9627    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9628        if t < PRIME_MIN_T {
9629            // Decode and speculative verify use one fixed per-row reduction program.
9630            if crate::router_kernel_on() {
9631                e.router_gemv(
9632                    m.gate_inp.float_data(),
9633                    z,
9634                    cfg.n_embd as usize,
9635                    m.gate_exps.n_expert,
9636                    t,
9637                )
9638            } else {
9639                e.matmul_decode_exact(&m.gate_inp, z, t)
9640            }
9641        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
9642            e.router_gemv(
9643                m.gate_inp.float_data(),
9644                z,
9645                cfg.n_embd as usize,
9646                m.gate_exps.n_expert,
9647                t,
9648            )
9649        } else {
9650            e.matmul(&m.gate_inp, z, t)
9651        }
9652    }
9653
9654    /// Append the host-visible router selection for one layer/forward when calibration tracing is
9655    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
9656    /// trace is independent of the dispatch optimization selected for the forward.
9657    /// `MEMRA_MOE_WEIGHT_TRACE` is also the co-activation measurement input of
9658    /// LAW:coactivation-expert-placement (lane/glm5-ep-place: rows ride this existing host
9659    /// readback — zero new device syncs; `glm5-tp-gate` arm T holds the ON-identity +
9660    /// row-count bar on the glm5 walks).
9661    fn trace_moe_routes(
9662        il: u16,
9663        t: usize,
9664        sel_all: &[u32],
9665        weights: &[f32],
9666    ) -> Result<(), Box<dyn std::error::Error>> {
9667        use std::io::Write as _;
9668        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
9669            let mut f = std::fs::OpenOptions::new()
9670                .create(true)
9671                .append(true)
9672                .open(path)?;
9673            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
9674            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
9675        }
9676        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
9677            let mut f = std::fs::OpenOptions::new()
9678                .create(true)
9679                .append(true)
9680                .open(path)?;
9681            let pairs: Vec<String> = sel_all
9682                .iter()
9683                .zip(weights)
9684                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
9685                .collect();
9686            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
9687        }
9688        Ok(())
9689    }
9690
9691    #[allow(clippy::too_many_arguments)]
9692    fn trace_sigmoid_router_logits(
9693        e: &Engine,
9694        il: u16,
9695        t: usize,
9696        n_expert: usize,
9697        n_used: usize,
9698        logits: &CudaSlice<f32>,
9699        m: &MoeWeights,
9700        (scaling_factor, route_norm): (f32, bool),
9701    ) -> Result<(), Box<dyn std::error::Error>> {
9702        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
9703            return Ok(());
9704        }
9705        let logits = e.dtoh(logits)?;
9706        let active: Vec<u8> = m
9707            .active_experts
9708            .as_ref()
9709            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
9710            .unwrap_or_else(|| vec![1; n_expert]);
9711        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
9712        crate::sigrouter_contract::capture_served_logits(
9713            il as u32,
9714            t,
9715            n_expert,
9716            n_used,
9717            scaling_factor,
9718            route_norm,
9719            &active,
9720            &bias,
9721            &logits,
9722        )?;
9723        Ok(())
9724    }
9725
9726    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
9727    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
9728    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
9729    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
9730    fn trace_moe_input(
9731        e: &Engine,
9732        il: u16,
9733        t: usize,
9734        n_embd: usize,
9735        z: &CudaSlice<f32>,
9736    ) -> Result<(), Box<dyn std::error::Error>> {
9737        use std::io::Write as _;
9738        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
9739            return Ok(());
9740        };
9741        let values = active_matrix_values(z.len(), t, n_embd, "MoE input trace activation")?;
9742        let host = e.dtoh_view(&z.slice(0..values))?;
9743        let bytes = unsafe {
9744            std::slice::from_raw_parts(
9745                host.as_ptr().cast::<u8>(),
9746                host.len() * std::mem::size_of::<f32>(),
9747            )
9748        };
9749        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
9750        let mut state = state
9751            .lock()
9752            .map_err(|_| "MoE input trace writer lock is poisoned")?;
9753        if state.is_none() {
9754            let dir = std::path::PathBuf::from(&dir);
9755            std::fs::create_dir_all(&dir)?;
9756            let index = std::fs::OpenOptions::new()
9757                .create(true)
9758                .append(true)
9759                .open(dir.join("index.jsonl"))?;
9760            *state = Some(MoeInputTraceWriter {
9761                dir,
9762                index,
9763                payloads: std::collections::HashMap::new(),
9764            });
9765        }
9766        let writer = state.as_mut().unwrap();
9767        if writer.dir != std::path::Path::new(&dir) {
9768            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
9769        }
9770        let file_name = format!("layer-{il:03}.f32");
9771        if !writer.payloads.contains_key(&il) {
9772            let payload = std::fs::OpenOptions::new()
9773                .create(true)
9774                .append(true)
9775                .open(writer.dir.join(&file_name))?;
9776            let offset = payload.metadata()?.len();
9777            writer.payloads.insert(il, (payload, offset));
9778        }
9779        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
9780        let row_offset = *offset;
9781        payload.write_all(bytes)?;
9782        *offset += bytes.len() as u64;
9783        writeln!(
9784            writer.index,
9785            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
9786             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
9787             \"payload_bytes\":{}}}",
9788            bytes.len()
9789        )?;
9790        Ok(())
9791    }
9792
9793    #[allow(clippy::too_many_arguments)]
9794    #[allow(clippy::too_many_arguments)]
9795    // allow: the parameter list mirrors its moe_ffn_inner caller's dispatch contract
9796    pub(crate) fn moe_ffn_sequential_zq8(
9797        e: &Engine,
9798        m: &MoeWeights,
9799        z: &CudaSlice<f32>,
9800        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
9801        t: usize,
9802        cfg: &ModelConfig,
9803        il: u16,
9804        max_block: usize,
9805        vrows: bool,
9806    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9807        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
9808        let moe = cfg.moe.as_ref().unwrap();
9809        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
9810        let n_expert = moe.expert_count as usize; // 256
9811        let n_used = moe.expert_used_count as usize; // 8
9812        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
9813
9814        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
9815        debug_assert_eq!(m.gate_exps.in_f, n_embd);
9816        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
9817        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
9818        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
9819        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
9820
9821        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
9822        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
9823        let lim_exp = cfg.clamp_exp_at(il as u32);
9824        let lim_shexp = cfg.clamp_shexp_at(il as u32);
9825        let use_cache = Engine::moe_cache_enabled();
9826        let uniform_experts = m.has_uniform_expert_layout();
9827        let moe_q8 = uniform_experts && moe_q8_enabled_for_model(cfg, m);
9828        // Experimental secondary backend: complete experts already resident in the SLRU stay on
9829        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
9830        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
9831        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
9832        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
9833        // commands and CI have no llama.cpp or OpenMP dependency.
9834        let cpu_expert_requested = crate::cpu_experts::configured();
9835        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
9836            return Err(std::io::Error::other(
9837                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
9838            )
9839            .into());
9840        }
9841        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
9842        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
9843        // Those backends are each deterministic but are different numeric configurations, so a
9844        // later prefill eviction can change greedy output. Freeze after the first real prefill;
9845        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
9846        // staging below and cannot change backend assignment.
9847        let freeze_cpu_residency = cpu_expert_requested
9848            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
9849        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
9850            .ok()
9851            .and_then(|value| value.parse::<usize>().ok())
9852            .is_some_and(|tokens| tokens > 0);
9853        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
9854            e.freeze_moe_cache();
9855        }
9856        let cache_frozen = use_cache && e.moe_cache_frozen();
9857        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
9858
9859        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
9860        // cannot change logits, selected expert ids, or routing weights.
9861        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
9862        if let Some(sig) = cfg.sigmoid_router() {
9863            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
9864        }
9865
9866        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
9867        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
9868        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
9869        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
9870        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
9871        // per-token host stall that dominated the 35B decode wall after stages 1+2.
9872        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
9873        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
9874        // only difference is where sel/w/pointers are READ from (device instead of params).
9875        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
9876        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
9877        // Any non-resident layer falls through to host routing + the gdec/sequential path.
9878        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
9879        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
9880        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
9881        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
9882        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
9883        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
9884        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
9885        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
9886        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
9887        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
9888        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
9889        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
9890        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
9891        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
9892        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
9893        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
9894        // now rides the dev loop below (same kernels per token as decode); pairs serves real
9895        // prefill (t >= 16, where spec never verifies).
9896        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
9897        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
9898        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
9899        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
9900        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
9901        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
9902        // ride the macro-aware sequential/staged paths below or every expert output is off by
9903        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
9904        let no_exp_macros = m.gate_exps.macros.is_none()
9905            && m.up_exps.macros.is_none()
9906            && m.down_exps.macros.is_none();
9907        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
9908        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
9909        // so it cannot even see the per-layer limit.
9910        if cfg.sigmoid_router().is_none()
9911            && cfg.m3.is_none()
9912            && cfg.hy3.is_none()
9913            && !cfg.swiglu_clamped_at(il as u32)
9914            && no_exp_macros
9915            // > MOE_DEV_MAX_T, not >= PRIME_MIN_T: t==16 is a decode width under the
9916            // exact-16 tier and rides the dev per-token program (see MOE_DEV_MAX_T);
9917            // pairs serves real prefill from 17 up.
9918            && t > MOE_DEV_MAX_T
9919            && m.dev_exps.is_some()
9920            && moe_q8_enabled_for_model(cfg, m)
9921            && std::env::var("MEMRA_MOE_PAIRS")
9922                .map(|v| v != "0")
9923                .unwrap_or(true)
9924            && std::env::var("MEMRA_MOE_STATS").is_err()
9925        {
9926            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
9927        }
9928
9929        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
9930        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
9931        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
9932        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
9933        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
9934        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
9935        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
9936        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
9937        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
9938        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
9939        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
9940        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
9941        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
9942        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
9943        // Keyed off sigmoid_router() so arch #4 is denied by construction.
9944        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
9945        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
9946        let dev_ok = uniform_experts
9947            && cfg.sigmoid_router().is_none()
9948            && cfg.m3.is_none()
9949            && cfg.hy3.is_none()
9950            && !cfg.swiglu_clamped_at(il as u32);
9951        // Observation modes must route through the host-visible selection below. Otherwise a fully
9952        // resident layer returns through device dispatch before its trace/stats row is recorded,
9953        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
9954        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
9955            || std::env::var("MEMRA_MOE_TRACE").is_ok()
9956            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
9957            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
9958        if dev_ok
9959            && t <= MOE_DEV_MAX_T
9960            && m.dev_exps.is_some()
9961            && n_used <= 8
9962            && moe_dev_enabled()
9963            && !observe_routes
9964        {
9965            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
9966        }
9967        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
9968            let row_ok = e.with_moe_cache(max_block, |c, eng| {
9969                if moe_prewarm_enabled() {
9970                    c.prewarm_layer(il, m, eng)?;
9971                }
9972                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
9973            })?;
9974            if row_ok {
9975                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
9976            }
9977        }
9978
9979        // SLAB-LOCAL RESIDENT ARM bases, hoisted above the router (lane/glm5-moe-loc door D):
9980        // whether the layer can run the DEVICE vrows table build decides HOW it routes, and
9981        // that has to be settled before the router runs. Pure immutable pointer reads with no
9982        // side effects, so the hoist changes nothing for any other arm; the full rationale for
9983        // the arm itself is at the `slab_fused_may_fire` predicate below.
9984        let slab_local = m
9985            .dev_exps
9986            .as_ref()
9987            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
9988        let slab_bases = slab_local.map(|d| {
9989            use cudarc::driver::DevicePtr;
9990            let s = e.stream();
9991            let (pg, _g0) = d.gate.device_ptr(&s);
9992            let (pu, _g1) = d.up.device_ptr(&s);
9993            let (pd, _g2) = d.down.device_ptr(&s);
9994            (pg, pu, pd)
9995        });
9996        // DOOR D (`MEMRA_MOE_VROWS_DEV_TABLES`, default OFF): route WITHOUT the pinned sel/w
9997        // readback and build the pair's pointer/scale tables on device instead. On the serving
9998        // shape the host table build is the selection's ONLY consumer, and it costs a full
9999        // `cuStreamSynchronize` + 2 DtoH + 2 pageable HtoD + 2 host Vecs per MoE layer-call —
10000        // 42 device-wide drains, 84 DtoH and 84 HtoD per ship round (the decode-gap
10001        // attribution's "43 cuStreamSynchronize/token ... the per-layer router-admission sync
10002        // structure", and 44.6% of the unattributed 71.6 HtoD calls/token).
10003        //
10004        // The extra conjuncts beyond `vrows_fires` (asserted equal at the dispatch) are exactly
10005        // the host-visible consumers of `sel_all` between here and there, each of which would
10006        // silently read an empty selection: `moesd::record_host_routes`, `hidden_trace`,
10007        // `MEMRA_MOE_TRACE`/`MEMRA_MOE_STATS`/the other `observe_routes` modes. Plus
10008        // `sigmoid_router_enabled()`, because `MEMRA_SIG_ROUTER=0` is a full-logit HOST oracle
10009        // with no device selection to read. `promote_worker_h2d` needs no conjunct: it requires
10010        // t == 1 and this arm requires t >= 2. Any miss falls closed to the host readback.
10011        let vrows_dev = vrows
10012            && t >= 2
10013            && crate::moe_vrows_dev_tables_on()
10014            && slab_bases.is_some()
10015            && moe_q8
10016            && uniform_experts
10017            && n_used <= 8
10018            && cfg.sigmoid_router().is_some()
10019            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10020            && !cpu_hybrid
10021            && sigmoid_router_enabled()
10022            && !observe_routes
10023            && !memra_reference::hidden_trace::enabled()
10024            && !crate::moesd::capture_active();
10025        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
10026        // Door D adds a fourth arm returning the router's DEVICE sel/w with no readback.
10027        let mut sel_dev: Option<(CudaSlice<i32>, CudaSlice<f32>)> = None;
10028        let (sel_all, w_all, routed_cpu_input) = if vrows_dev {
10029            let (sf, route_norm) = cfg
10030                .sigmoid_router()
10031                .expect("vrows_dev carries cfg.sigmoid_router().is_some()");
10032            sel_dev = Some(e.moe_router_sigmoid_topk(
10033                &logits,
10034                t,
10035                n_expert,
10036                n_used,
10037                m.active_count(),
10038                &m.exp_probs_b_dev,
10039                &m.active_experts_dev,
10040                sf,
10041                route_norm,
10042            )?);
10043            crate::MOE_VROWS_ROUTER_SYNCS_AVOIDED
10044                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10045            (Vec::new(), Vec::new(), None)
10046        } else if let Some(sig) = cfg.sigmoid_router() {
10047            if cpu_hybrid {
10048                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
10049                    e,
10050                    &logits,
10051                    z,
10052                    t,
10053                    n_embd,
10054                    n_expert,
10055                    n_used,
10056                    m.exp_probs_b.as_deref(),
10057                    sig,
10058                    m.active_experts.as_deref(),
10059                )?;
10060                (sel, w, Some(input))
10061            } else {
10062                let (sel, w) =
10063                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
10064                (sel, w, None)
10065            }
10066        } else {
10067            let (sel, w) =
10068                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
10069            (sel, w, None)
10070        };
10071        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
10072        if memra_reference::hidden_trace::enabled() {
10073            memra_reference::hidden_trace::emit_last_row(
10074                "router",
10075                il as i64,
10076                t,
10077                n_expert,
10078                &e.dtoh(&logits)?,
10079            );
10080            let last = (t - 1) * n_used;
10081            let mut route = Vec::with_capacity(n_used * 2);
10082            for slot in 0..n_used {
10083                route.push(sel_all[last + slot] as f32);
10084                route.push(w_all[last + slot]);
10085            }
10086            memra_reference::hidden_trace::emit_last_row("route", il as i64, 1, n_used * 2, &route);
10087        }
10088
10089        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
10090        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
10091        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
10092        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
10093        Self::trace_moe_input(e, il, t, n_embd, z)?;
10094
10095        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
10096        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
10097        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
10098        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
10099        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
10100        // wait for each pending block, so later copies can overlap the earlier expert kernels while
10101        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
10102        // T=1; batched forwards can have token-local consumers still in flight between selections.
10103        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
10104        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
10105        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
10106        let worker_disk_prefetch =
10107            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
10108        let promote_worker_h2d =
10109            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
10110        if promote_worker_h2d {
10111            let mut selected_blocks = Vec::with_capacity(n_used * 3);
10112            for &ex in sel_all.iter().take(n_used) {
10113                let ex = ex as u16;
10114                selected_blocks.extend([
10115                    BlockId::new(il, PROJ_GATE, ex),
10116                    BlockId::new(il, PROJ_UP, ex),
10117                    BlockId::new(il, PROJ_DOWN, ex),
10118                ]);
10119            }
10120            for &ex in sel_all.iter().take(n_used) {
10121                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
10122            }
10123            e.with_moe_cache(max_block, |cache, eng| {
10124                cache.promote_worker_reads_at_safe_boundary(
10125                    &selected_blocks,
10126                    &selected_blocks,
10127                    eng,
10128                )?;
10129                Ok(())
10130            })?;
10131        }
10132
10133        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
10134        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
10135        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
10136            let mut cnt = vec![0u32; n_expert];
10137            for &s in sel_all.iter() {
10138                cnt[s as usize] += 1;
10139            }
10140            let total = sel_all.len() as f64;
10141            let mut h = 0.0f64;
10142            let mut active = 0usize;
10143            for &c in &cnt {
10144                if c > 0 {
10145                    active += 1;
10146                    let p = c as f64 / total;
10147                    h -= p * p.log2();
10148                }
10149            }
10150            let maxc = cnt.iter().copied().max().unwrap_or(0);
10151            println!(
10152                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
10153                il,
10154                t,
10155                sel_all.len(),
10156                active,
10157                n_expert,
10158                h,
10159                (n_expert as f64).log2(),
10160                total / active.max(1) as f64,
10161                maxc
10162            );
10163        }
10164
10165        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
10166        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
10167        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
10168        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
10169        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
10170        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
10171        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
10172        // zeroed-then-accumulated exactly as before (fallback).
10173        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
10174        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
10175        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
10176        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
10177        let gdec_may_fire = uniform_experts
10178            && use_cache
10179            && n_used <= 8
10180            && gdec_enabled()
10181            && !cfg.swiglu_clamped_at(il as u32);
10182        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
10183        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
10184        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
10185        // archs the slabs were uploaded but never read, and every expert went through the
10186        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
10187        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
10188        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
10189        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
10190        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
10191        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
10192        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
10193        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
10194        // strictly worse than staging); under PP-2 without the prime walker this admits
10195        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
10196        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
10197        // `slab_local` / `slab_bases` are bound ABOVE the router: door D
10198        // (`MEMRA_MOE_VROWS_DEV_TABLES`) has to pre-decide the routing arm, and this is the
10199        // predicate it needs. Pure immutable pointer reads, so the hoist is behaviour-neutral.
10200        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
10201        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
10202        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
10203        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
10204        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
10205        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
10206        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
10207        // all-resident tokens, staged loop for misses), which is a dispatch-class
10208        // comparison, not a provenance one.
10209        let slab_fused_may_fire = slab_bases.is_some()
10210            && n_used <= 8
10211            && gdec_enabled()
10212            && !cfg.swiglu_clamped_at(il as u32)
10213            && cfg.m3.is_none()
10214            && no_exp_macros
10215            && moe_q8;
10216        // FUSED MoE EPILOGUE (lane/glm53-epilogue 2026-08-28, MEMRA_MOE_FUSED_EPI default OFF).
10217        // The arm glm5_next is denied by every other predicate in this function. It runs the SAME
10218        // launch pair shape as gdec — one gate/up kernel, one down/FMA kernel per token — but
10219        // with the three things this arch actually needs, none of which any existing fused
10220        // epilogue has:
10221        //   * the SIGMOID noaux_tc router's sel/w (host-routed above; the fused softmax router
10222        //     `moe_router_topk` that pairs/dev use would pick different experts — the M3
10223        //     gate-MISMATCH 74602-vs-92 lesson);
10224        //   * the PRE-clamped SwiGLU epilogue `silu(min(g,l)) * clamp(u,±l)` — step35's POST form
10225        //     is a different, plausible-but-wrong program (`fused_post_limit`);
10226        //   * the per-expert NVFP4 `weight_scale_2` macro fold, gate/up through the kernel's
10227        //     gs/us and down through the routing weight, exactly as `ffn_act_lim` + `axpy_into`
10228        //     do it in the sequential loop.
10229        // UNLIKE gdec it does NOT require the layer to be already-resident: it ADMITS the
10230        // 3*n_used selected blocks through the same `dispatch_source` the sequential loop uses
10231        // (hit = no copy, miss = the identical H2D into a slot) and only then collects the fixed
10232        // slot addresses. The staged bytes are unchanged — the §B.3 provenance property — so the
10233        // arm engages at any miss rate instead of gdec's P(all resident). It needs the cache to
10234        // hold 3*n_used blocks at once; `moe_fused_epi_token_q8` returns false when it cannot and
10235        // the token falls through to the sequential loop below.
10236        // TWO PROVENANCES, ONE LAUNCH PATH (slab arm added 2026-08-28). The SLRU arm below keys
10237        // on `slab_local.is_none()`; the SLAB arm keys on the slab existing. They differ ONLY in
10238        // where the eight expert pointers come from and both call `moe_fused_epi_launch`, so the
10239        // macro fold, the clamp and the kernel pair cannot drift apart between them.
10240        //
10241        // The slab arm is not an optimization, it is the arm that matters. Full two-card expert
10242        // residency makes `dev_exps` present on every stage engine, which makes `slab_local`
10243        // `Some`, which under the original predicate DENIED the fused epilogue outright — the
10244        // measured A/B would have read 0 dispatches and looked like "no effect". The residency
10245        // config is the serving config now, so the slab provenance is the one the product runs.
10246        //
10247        // It is also the SIMPLER arm: a slab holds every expert by construction, so there is no
10248        // admission, no eviction, no pass-2 re-verification and no fall-through. The SLRU arm's
10249        // capacity floor and re-check exist only because admission can move a slot.
10250        let fused_epi_common = n_used <= 8
10251            && moe_q8
10252            && cfg.m3.is_none()
10253            && cfg.sigmoid_router().is_some()
10254            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10255            && moe_fused_epi_enabled();
10256        let fused_epi_may_fire = fused_epi_common
10257            && uniform_experts
10258            && use_cache
10259            && cache_dispatch
10260            && slab_local.is_none();
10261        let fused_epi_slab_may_fire = fused_epi_common && slab_bases.is_some();
10262        // VERIFY-ROWS BATCHED ROUTED-EXPERT ARM (lane/glm5-vrest, 2026-08-31; rides
10263        // `MEMRA_GLM5_VERIFY_BATCH` — only the verify walk's batched arm passes `vrows`).
10264        // ONE launch pair covers ALL t x n_used routed pairs (the fused-epilogue kernels'
10265        // verify-rows twins) instead of the per-(token,expert) loop's ~49 launches per
10266        // token-layer — the flip-reprice cell-2 vrest wall (9.46 ms/row marginal at K=3).
10267        // Bit identity per row vs the sequential chain is the bar and it is structural:
10268        // routing is the SAME host invocation above; per-pair dots are qmatvec_expert_q8's
10269        // g-strided order; the epilogue is swiglu_preclamped_mul_scaled_f32's expression
10270        // with the per-expert macro fold exactly where ffn_act_lim/axpy_into fold it; the
10271        // down accumulation is the slot-ordered __fmaf_rn chain (the gdec-gated class).
10272        // Gated by glm5_verify_batch_gpu (kernel pair vs sequential chain + swapped-pair
10273        // and dropped-macro reds) and the glm5_tparallel_verify_gpu walk battery on the
10274        // NVFP4+macro serving expert class. Fail-closed: any unqualified shape falls
10275        // through to the unchanged loop below. Same slab-only scope as the fused epilogue
10276        // (the serving config's provenance); n_used<=8 mirrors its cap.
10277        let vrows_fires = vrows
10278            && t >= 2
10279            && slab_bases.is_some()
10280            && moe_q8
10281            && uniform_experts
10282            && n_used <= 8
10283            && cfg.sigmoid_router().is_some()
10284            && matches!(lim_exp, Some(SwigluClamp::Pre(l)) if l > 1e-6)
10285            && !cpu_hybrid;
10286        // moe_out memset elision: EVERY full-row-overwrite arm (gdec, slab fused, fused epilogue,
10287        // verify-rows) allocates uninit; a token that falls through to any accumulating loop
10288        // zeroes its own row. The fused epilogue's `moe_down8_fma_q8` fully overwrites `dst[o]`,
10289        // same as gdec's; `moe_down8_fma_q8_rows` fully overwrites every row.
10290        let mut moe_out = if gdec_may_fire
10291            || slab_fused_may_fire
10292            || fused_epi_may_fire
10293            || fused_epi_slab_may_fire
10294            || vrows_fires
10295        {
10296            e.uninit(t * n_embd)?
10297        } else {
10298            e.zeros(t * n_embd)?
10299        };
10300        // The router readback above already established a host boundary. Copy each small-t hidden
10301        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
10302        let cpu_input = if cpu_hybrid {
10303            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
10304        } else {
10305            None
10306        };
10307
10308        // Door D's predicate was evaluated at the router, BEFORE `vrows_fires` existed. If the
10309        // two ever disagreed, the layer would have skipped its readback and then dispatched a
10310        // per-row arm holding an EMPTY host selection — a silent wrong-answer class. The two
10311        // predicates share every conjunct by construction; assert it rather than hope.
10312        if vrows_dev && !vrows_fires {
10313            return Err(
10314                "MEMRA_MOE_VROWS_DEV_TABLES routed device-only but the verify-rows arm did not \
10315                 fire: the door-D and vrows_fires predicates disagree"
10316                    .into(),
10317            );
10318        }
10319        if vrows_fires {
10320            let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10321                return Err(
10322                    "verify-rows MoE arm fired without a live PRE clamp: the predicate and \
10323                     the dispatch disagree"
10324                        .into(),
10325                );
10326            };
10327            let bases = slab_bases.expect("vrows_fires carries slab_bases.is_some()");
10328            let sel = match sel_dev.as_ref() {
10329                Some((si, sw)) => VrowsSel::Dev(si, sw),
10330                None => VrowsSel::Host(&sel_all, &w_all),
10331            };
10332            Self::moe_vrows_pairs_q8(
10333                e,
10334                m,
10335                z,
10336                sel,
10337                il,
10338                bases,
10339                t,
10340                n_embd,
10341                n_ff_exp,
10342                n_used,
10343                limit,
10344                &mut moe_out,
10345            )?;
10346            if memra_reference::hidden_trace::enabled() {
10347                memra_reference::hidden_trace::emit_last_row(
10348                    "routed",
10349                    il as i64,
10350                    t,
10351                    n_embd,
10352                    &e.dtoh(&moe_out)?,
10353                );
10354            }
10355            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
10356            return Ok(moe_out);
10357        }
10358
10359        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
10360        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
10361        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
10362        // measured ~123 memsets/token of the decode wall).
10363        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
10364        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
10365        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
10366        let mut scratch_g: Option<CudaSlice<u8>> = None;
10367        let mut scratch_u: Option<CudaSlice<u8>> = None;
10368        let mut scratch_d: Option<CudaSlice<u8>> = None;
10369        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
10370        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
10371
10372        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
10373        // the copy stream before launching the current expert's compute. Pending slots stay invisible
10374        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
10375        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
10376        let page_window = moe_page_prefetch_window();
10377
10378        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
10379        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
10380        for tok in 0..t {
10381            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
10382            let w = &w_all[tok * n_used..(tok + 1) * n_used];
10383            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
10384            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10385
10386            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
10387            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
10388            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
10389            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
10390            // memcpy, zero admission, so no slot can move under the collected pointers) — any
10391            // miss falls through to the sequential loop below, which admits as before. In steady
10392            // state on a fully-resident rig every token-layer takes the grouped path.
10393            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
10394            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
10395            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
10396            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
10397            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
10398            // per-expert macro-scales the fused kernels don't fold — those fall through too.
10399            let no_macros = m.gate_exps.macros.is_none()
10400                && m.up_exps.macros.is_none()
10401                && m.down_exps.macros.is_none();
10402            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
10403            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
10404            // with pointers computed from the resident slab base + ex*stride instead of
10405            // collected SLRU slot addresses. No cache lock, no residency predicate — the
10406            // slab holds every expert by construction, so this arm never falls through
10407            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
10408            // staging both die). Bit-identity class: pointer provenance only, the same
10409            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
10410            // slab exists it is strictly better (no lock, no miss).
10411            if slab_fused_may_fire {
10412                let (pg, pu, pd) = slab_bases.unwrap();
10413                let mut gp = [0u64; 8];
10414                let mut up = [0u64; 8];
10415                let mut dp = [0u64; 8];
10416                for (j, &ex) in sel.iter().enumerate() {
10417                    let ex = ex as usize;
10418                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10419                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10420                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10421                }
10422                let mut wv = [0f32; 8];
10423                wv[..n_used].copy_from_slice(w);
10424                if tok_q8.is_none() {
10425                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10426                }
10427                let (zq, zd) = tok_q8.as_ref().unwrap();
10428                let act = e.moe_gate_up_silu8_q8(
10429                    crate::WPtr8(gp),
10430                    crate::WPtr8(up),
10431                    zq,
10432                    zd,
10433                    n_embd,
10434                    n_ff_exp,
10435                    n_used,
10436                    m.gate_exps.qtype,
10437                    m.up_exps.qtype,
10438                    m.gate_exps.row_bytes,
10439                    m.up_exps.row_bytes,
10440                )?;
10441                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
10442                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10443                e.moe_down8_fma_q8(
10444                    crate::WPtr8(dp),
10445                    crate::F32x8(wv),
10446                    &aq2,
10447                    &ad2,
10448                    &mut dst,
10449                    n_ff_exp,
10450                    n_embd,
10451                    n_used,
10452                    m.down_exps.qtype,
10453                    m.down_exps.row_bytes,
10454                )?;
10455                continue;
10456            }
10457            // FUSED MoE EPILOGUE, SLAB PROVENANCE. Ordered first: when a local slab exists it is
10458            // strictly better than anything the SLRU can offer — every expert is present by
10459            // construction, so there is no admission, no eviction and no fall-through. This is
10460            // the arm the two-card residency serving config actually runs.
10461            if fused_epi_slab_may_fire {
10462                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10463                    return Err(
10464                        "fused MoE epilogue (slab) fired without a live PRE clamp: the \
10465                                predicate and the dispatch disagree"
10466                            .into(),
10467                    );
10468                };
10469                let (pg, pu, pd) = slab_bases.unwrap();
10470                let mut g = [0u64; 8];
10471                let mut u = [0u64; 8];
10472                let mut d = [0u64; 8];
10473                for (j, &ex) in sel.iter().enumerate() {
10474                    let ex = ex as usize;
10475                    g[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
10476                    u[j] = pu + (ex * m.up_exps.expert_stride) as u64;
10477                    d[j] = pd + (ex * m.down_exps.expert_stride) as u64;
10478                }
10479                if tok_q8.is_none() {
10480                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10481                }
10482                let (zq, zd) = tok_q8.as_ref().unwrap();
10483                Self::moe_fused_epi_launch(
10484                    e,
10485                    m,
10486                    zq,
10487                    zd,
10488                    sel,
10489                    w,
10490                    g,
10491                    u,
10492                    d,
10493                    &mut moe_out,
10494                    tok,
10495                    n_embd,
10496                    n_ff_exp,
10497                    n_used,
10498                    limit,
10499                )?;
10500                continue;
10501            }
10502            // FUSED MoE EPILOGUE, SLRU PROVENANCE. Ordered before gdec (which this arch never
10503            // reaches anyway: `gdec_may_fire` carries `!swiglu_clamped_at`). A `false` return
10504            // means the cache could not hold 3*n_used blocks at once — the token falls through to
10505            // the sequential loop, which zeroes its own row below.
10506            if fused_epi_may_fire {
10507                let Some(SwigluClamp::Pre(limit)) = lim_exp else {
10508                    return Err(
10509                        "fused MoE epilogue fired without a live PRE clamp: the predicate and \
10510                         the dispatch disagree"
10511                            .into(),
10512                    );
10513                };
10514                if tok_q8.is_none() {
10515                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10516                }
10517                let (zq, zd) = tok_q8.as_ref().unwrap();
10518                if Self::moe_fused_epi_token_q8(
10519                    e,
10520                    m,
10521                    il,
10522                    max_block,
10523                    zq,
10524                    zd,
10525                    sel,
10526                    w,
10527                    &mut moe_out,
10528                    tok,
10529                    n_embd,
10530                    n_ff_exp,
10531                    n_used,
10532                    limit,
10533                )? {
10534                    continue;
10535                }
10536            }
10537            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
10538                if tok_q8.is_none() {
10539                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10540                }
10541                let (zq, zd) = tok_q8.as_ref().unwrap();
10542                if Self::moe_gdec_token_q8(
10543                    e,
10544                    m,
10545                    il,
10546                    max_block,
10547                    zq,
10548                    zd,
10549                    sel,
10550                    w,
10551                    &mut moe_out,
10552                    tok,
10553                    n_embd,
10554                    n_ff_exp,
10555                    n_used,
10556                )? {
10557                    continue;
10558                }
10559            } else if gdec_may_fire
10560                && cfg.m3.is_none()
10561                && no_macros
10562                && Self::moe_gdec_token(
10563                    e,
10564                    m,
10565                    il,
10566                    max_block,
10567                    &zt,
10568                    sel,
10569                    w,
10570                    &mut moe_out,
10571                    tok,
10572                    n_embd,
10573                    n_ff_exp,
10574                    n_used,
10575                )?
10576            {
10577                continue;
10578            }
10579
10580            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
10581            // slab pair could fire. This token fell through to a sequential axpy loop, which
10582            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
10583            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
10584            // has no fallible predicate), included for the allocation invariant's symmetry.
10585            if gdec_may_fire || slab_fused_may_fire || fused_epi_may_fire || fused_epi_slab_may_fire
10586            {
10587                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10588                e.memset_zeros_view(&mut row)?;
10589            }
10590
10591            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
10592            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
10593            // stall this path exists to remove, while mixing projections would require another
10594            // activation round-trip. Weight addresses remain valid until this worker is joined at
10595            // the bottom of the token scope.
10596            let mut cpu_mask = vec![false; sel.len()];
10597            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
10598                let gpu_resident = if use_cache {
10599                    e.with_moe_cache(max_block, |cache, _| {
10600                        Ok(sel
10601                            .iter()
10602                            .map(|&expert| {
10603                                let expert = expert as u16;
10604                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
10605                                    .into_iter()
10606                                    .filter(|&projection| {
10607                                        cache
10608                                            .resident(BlockId::new(il, projection, expert))
10609                                            .is_some()
10610                                    })
10611                                    .count()
10612                            })
10613                            .collect::<Vec<_>>())
10614                    })?
10615                } else {
10616                    vec![0; sel.len()]
10617                };
10618                let mut cpu_selected = Vec::new();
10619                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
10620                    if gpu_resident[index] != 3 {
10621                        cpu_mask[index] = true;
10622                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
10623                        let expert = expert as usize;
10624                        cpu_selected.push((expert, route_weight));
10625                    }
10626                }
10627                if crate::cpu_experts::predictor_enabled() {
10628                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
10629                    // from this layer's MoE input and prefetches predicted-and-missing
10630                    // experts into the companion RAM cache. Never blocks this thread.
10631                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10632                    crate::cpu_experts::predictor_submit(il, row);
10633                }
10634                if cpu_selected.is_empty() {
10635                    None
10636                } else {
10637                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
10638                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
10639                        .map_err(std::io::Error::other)?;
10640                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
10641                }
10642            } else {
10643                None
10644            };
10645
10646            let worker_window = worker_disk_prefetch
10647                .then(worker_prefetch_window)
10648                .unwrap_or(0);
10649            for (j, &ex) in sel.iter().enumerate() {
10650                if cpu_mask[j] {
10651                    continue;
10652                }
10653                let ex = ex as usize;
10654                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
10655                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
10656                // fused form) and macro-carrying artifacts — still have their bytes in the
10657                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
10658                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
10659                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
10660                if let Some(d) = slab_local {
10661                    let gl = m.gate_exps.expert_layout(ex);
10662                    let ul = m.up_exps.expert_layout(ex);
10663                    let dl = m.down_exps.expert_layout(ex);
10664                    let (g0, u0, d0) = (
10665                        ex * m.gate_exps.expert_stride,
10666                        ex * m.up_exps.expert_stride,
10667                        ex * m.down_exps.expert_stride,
10668                    );
10669                    let (gate, up) = if moe_q8 {
10670                        if tok_q8.is_none() {
10671                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10672                        }
10673                        let (zq, zd) = tok_q8.as_ref().unwrap();
10674                        (
10675                            e.qmatvec_expert_q8(
10676                                &d.gate,
10677                                g0..g0 + gl.len,
10678                                zq,
10679                                zd,
10680                                1,
10681                                m.gate_exps.in_f,
10682                                m.gate_exps.out_f,
10683                                gl.qtype,
10684                                gl.row_bytes,
10685                            )?,
10686                            e.qmatvec_expert_q8(
10687                                &d.up,
10688                                u0..u0 + ul.len,
10689                                zq,
10690                                zd,
10691                                1,
10692                                m.up_exps.in_f,
10693                                m.up_exps.out_f,
10694                                ul.qtype,
10695                                ul.row_bytes,
10696                            )?,
10697                        )
10698                    } else {
10699                        (
10700                            m.qmatvec_view(
10701                                e,
10702                                &d.gate,
10703                                g0..g0 + gl.len,
10704                                &zt,
10705                                1,
10706                                m.gate_exps.in_f,
10707                                m.gate_exps.out_f,
10708                                gl.qtype,
10709                                gl.row_bytes,
10710                            )?,
10711                            m.qmatvec_view(
10712                                e,
10713                                &d.up,
10714                                u0..u0 + ul.len,
10715                                &zt,
10716                                1,
10717                                m.up_exps.in_f,
10718                                m.up_exps.out_f,
10719                                ul.qtype,
10720                                ul.row_bytes,
10721                            )?,
10722                        )
10723                    };
10724                    let mut act = e.uninit(n_ff_exp)?;
10725                    Self::ffn_act_lim(
10726                        e,
10727                        cfg,
10728                        &gate,
10729                        &up,
10730                        m.gate_exps.macro_scale(ex),
10731                        m.up_exps.macro_scale(ex),
10732                        lim_exp,
10733                        &mut act,
10734                        n_ff_exp,
10735                    )?;
10736                    let y = if moe_q8 {
10737                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10738                        e.qmatvec_expert_q8(
10739                            &d.down,
10740                            d0..d0 + dl.len,
10741                            &aq2,
10742                            &ad2,
10743                            1,
10744                            m.down_exps.in_f,
10745                            m.down_exps.out_f,
10746                            dl.qtype,
10747                            dl.row_bytes,
10748                        )?
10749                    } else {
10750                        let actv = act.slice(0..n_ff_exp);
10751                        m.qmatvec_view(
10752                            e,
10753                            &d.down,
10754                            d0..d0 + dl.len,
10755                            &actv,
10756                            1,
10757                            m.down_exps.in_f,
10758                            m.down_exps.out_f,
10759                            dl.qtype,
10760                            dl.row_bytes,
10761                        )?
10762                    };
10763                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10764                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10765                    continue;
10766                }
10767                for next in page_prefetch_positions(j, sel.len(), page_window) {
10768                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
10769                }
10770                let keep = [
10771                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
10772                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
10773                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
10774                ];
10775                if worker_disk_prefetch && worker_window > 0 {
10776                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
10777                        Self::moe_prefetch_disk_expert(
10778                            e,
10779                            il,
10780                            sel[next] as usize,
10781                            m,
10782                            max_block,
10783                            &keep,
10784                        )?;
10785                    }
10786                } else if cache_dispatch
10787                    && !cpu_hybrid
10788                    && moe_prefetch_enabled()
10789                    && j + 1 < sel.len()
10790                {
10791                    let next = sel[j + 1] as usize;
10792                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
10793                }
10794                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
10795                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
10796                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
10797                    // layouts stay on the metadata-aware f32 path.
10798                    if (gate_q8 || up_q8) && tok_q8.is_none() {
10799                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
10800                    }
10801                    let gate = if gate_q8 {
10802                        let (zq, zd) = tok_q8.as_ref().unwrap();
10803                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
10804                    } else {
10805                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
10806                    };
10807                    let up = if up_q8 {
10808                        let (zq, zd) = tok_q8.as_ref().unwrap();
10809                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
10810                    } else {
10811                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
10812                    };
10813                    let mut act = e.uninit(n_ff_exp)?;
10814                    Self::ffn_act_lim(
10815                        e,
10816                        cfg,
10817                        &gate,
10818                        &up,
10819                        m.gate_exps.macro_scale(ex),
10820                        m.up_exps.macro_scale(ex),
10821                        lim_exp,
10822                        &mut act,
10823                        n_ff_exp,
10824                    )?;
10825                    let y = if down_q8 {
10826                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
10827                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
10828                    } else {
10829                        let actv = act.slice(0..n_ff_exp);
10830                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
10831                    };
10832                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10833                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
10834                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10835                } else if cache_dispatch {
10836                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
10837                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
10838                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
10839                    // only difference between HIT and MISS is whether the memcpy_htod ran.
10840                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
10841                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
10842                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
10843                    Self::ffn_act_lim(
10844                        e,
10845                        cfg,
10846                        &gate,
10847                        &up,
10848                        m.gate_exps.macro_scale(ex),
10849                        m.up_exps.macro_scale(ex),
10850                        lim_exp,
10851                        &mut act,
10852                        n_ff_exp,
10853                    )?;
10854                    let actv = act.slice(0..n_ff_exp);
10855                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
10856                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10857                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
10858                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10859                } else if cache_frozen {
10860                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
10861                    // first prime. Reuse every fixed resident projection directly and stage only a
10862                    // true miss through the ordinary scratch slot. This preserves the established
10863                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
10864                    let gate = Self::moe_frozen_gemm(
10865                        e,
10866                        il,
10867                        PROJ_GATE,
10868                        ex,
10869                        m,
10870                        max_block,
10871                        &zt,
10872                        &mut scratch_g,
10873                        g_len,
10874                    )?;
10875                    let up = Self::moe_frozen_gemm(
10876                        e,
10877                        il,
10878                        PROJ_UP,
10879                        ex,
10880                        m,
10881                        max_block,
10882                        &zt,
10883                        &mut scratch_u,
10884                        u_len,
10885                    )?;
10886                    let mut act = e.uninit(n_ff_exp)?;
10887                    Self::ffn_act_lim(
10888                        e,
10889                        cfg,
10890                        &gate,
10891                        &up,
10892                        m.gate_exps.macro_scale(ex),
10893                        m.up_exps.macro_scale(ex),
10894                        lim_exp,
10895                        &mut act,
10896                        n_ff_exp,
10897                    )?;
10898                    let actv = act.slice(0..n_ff_exp);
10899                    let y = Self::moe_frozen_gemm(
10900                        e,
10901                        il,
10902                        PROJ_DOWN,
10903                        ex,
10904                        m,
10905                        max_block,
10906                        &actv,
10907                        &mut scratch_d,
10908                        d_len,
10909                    )?;
10910                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10911                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10912                } else {
10913                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
10914                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
10915                    // fully overwrites the byte range the GEMM reads).
10916                    if scratch_g.is_none() {
10917                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
10918                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
10919                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
10920                    }
10921                    let (sg, su, sd) = (
10922                        scratch_g.as_mut().unwrap(),
10923                        scratch_u.as_mut().unwrap(),
10924                        scratch_d.as_mut().unwrap(),
10925                    );
10926                    let gl = m.gate_exps.expert_layout(ex);
10927                    let ul = m.up_exps.expert_layout(ex);
10928                    let dl = m.down_exps.expert_layout(ex);
10929                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
10930                    let gate = m.qmatvec_view(
10931                        e,
10932                        sg,
10933                        0..gl.len,
10934                        &zt,
10935                        1,
10936                        m.gate_exps.in_f,
10937                        m.gate_exps.out_f,
10938                        gl.qtype,
10939                        gl.row_bytes,
10940                    )?;
10941
10942                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
10943                    let up = m.qmatvec_view(
10944                        e,
10945                        su,
10946                        0..ul.len,
10947                        &zt,
10948                        1,
10949                        m.up_exps.in_f,
10950                        m.up_exps.out_f,
10951                        ul.qtype,
10952                        ul.row_bytes,
10953                    )?;
10954
10955                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
10956                    Self::ffn_act_lim(
10957                        e,
10958                        cfg,
10959                        &gate,
10960                        &up,
10961                        m.gate_exps.macro_scale(ex),
10962                        m.up_exps.macro_scale(ex),
10963                        lim_exp,
10964                        &mut act,
10965                        n_ff_exp,
10966                    )?;
10967
10968                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
10969                    let actv = act.slice(0..n_ff_exp);
10970                    let y = m.qmatvec_view(
10971                        e,
10972                        sd,
10973                        0..dl.len,
10974                        &actv,
10975                        1,
10976                        m.down_exps.in_f,
10977                        m.down_exps.out_f,
10978                        dl.qtype,
10979                        dl.row_bytes,
10980                    )?;
10981
10982                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10983                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
10984                }
10985            }
10986            if let Some(worker) = cpu_worker {
10987                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
10988                let cpu_output = e.htod(&cpu_output)?;
10989                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
10990                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
10991            }
10992            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
10993                for (j, &ex) in sel.iter().enumerate() {
10994                    if cpu_mask[j] {
10995                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
10996                    }
10997                }
10998            }
10999        }
11000
11001        if memra_reference::hidden_trace::enabled() {
11002            memra_reference::hidden_trace::emit_last_row(
11003                "routed",
11004                il as i64,
11005                t,
11006                n_embd,
11007                &e.dtoh(&moe_out)?,
11008            );
11009        }
11010
11011        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11012
11013        Ok(moe_out)
11014    }
11015
11016    /// The glm5 TP-2 EP walk (`MEMRA_GLM5_TP`): one MoE layer's routed-expert FFN over
11017    /// whole-expert contiguous halves. The ROUTER is the unchanged root-side program
11018    /// (`moe_router_logits` + `moe_route_sigmoid_cfg` — bit-identical selection by
11019    /// construction); each rank computes its owned slots' UNWEIGHTED expert rows through
11020    /// the sequential per-expert program (gate/up qmatvec + `ffn_act_lim` + down qmatvec —
11021    /// per-expert-independent dots), the peer's rows return host-canonically, and root
11022    /// applies the slot-ordered `axpy` accumulation chain — the same rounded-operation
11023    /// sequence the plain sequential walk applies. The ROOT-owned shared expert then adds
11024    /// through the extracted `moe_shexp_add`, verbatim.
11025    ///
11026    /// Three transport arms behind ONE routing (lane/glm5-ep-diet; sel/w are shared so a
11027    /// dispatch change can never change selection):
11028    ///   * `MEMRA_GLM5_EP_GROUPED_PRIME` (prefill shapes only): per-rank grouped-GEMM prime
11029    ///     over the rank slabs — the plain grouped-prefill program split by ownership.
11030    ///     Falls closed to the arms below whenever the plain arm's conjuncts do not hold.
11031    ///   * `MEMRA_GLM5_EP_DIET`: the v1 walk's kernels and combine chain with dieted data
11032    ///     movement — one bulk fan-out, zero per-slot host round-trips, one combine launch.
11033    ///     Decode-byte-identical to v1 by construction.
11034    ///   * default: the v1 per-slot host-canonical walk, byte-for-byte.
11035    #[allow(clippy::too_many_arguments)]
11036    fn moe_ffn_glm5_ep(
11037        e: &Engine,
11038        m: &MoeWeights,
11039        ep: &crate::glm5_tp::Glm5EpExps,
11040        z: &CudaSlice<f32>,
11041        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11042        t: usize,
11043        cfg: &ModelConfig,
11044        il: u16,
11045        prefill: bool,
11046    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11047        let moe = cfg
11048            .moe
11049            .as_ref()
11050            .ok_or("glm5 EP execution requires MoE model metadata")?;
11051        let n_embd = cfg.n_embd as usize;
11052        let n_expert = moe.expert_count as usize;
11053        let n_used = moe.expert_used_count as usize;
11054        let n_ff_exp = moe.expert_ff_length as usize;
11055        let sig = cfg
11056            .sigmoid_router()
11057            .ok_or("glm5 EP execution requires the sigmoid router")?;
11058        let lim_exp = cfg.clamp_exp_at(il as u32);
11059        let lim_shexp = cfg.clamp_shexp_at(il as u32);
11060        if ep.slabs.iter().map(|s| s.n_experts).sum::<usize>() != n_expert {
11061            return Err(format!(
11062                "glm5 EP slabs cover {:?} experts, model declares {n_expert}",
11063                ep.slabs.iter().map(|s| s.n_experts).collect::<Vec<_>>()
11064            )
11065            .into());
11066        }
11067        let rt = &ep.rt;
11068        let ranks = ep.ranks();
11069
11070        // Root router, unchanged program (selection bit-identical to the sequential arm).
11071        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
11072        let (sel_all, w_all) =
11073            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
11074        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
11075        // The route trace taps ride this walk too (sel/w are already host-side here);
11076        // the EP walk must never be a blind spot for the co-activation measurement.
11077        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
11078
11079        // EP grouped prime (`MEMRA_GLM5_EP_GROUPED_PRIME`, default OFF): keyed exactly like
11080        // the plain grouped-prefill arm (`prefill` + t above the per-token tier) and honoring
11081        // the family rollback (`MEMRA_MOE_GROUPED_PREFILL=0` kills it too). The announce
11082        // prints once per process PER FLAG VALUE, in both arms, so an A/B grep distinguishes
11083        // engagement without the line being an arm-local cost.
11084        if prefill && t > MOE_DEV_MAX_T {
11085            static EPGP_ANNOUNCED: std::sync::atomic::AtomicU8 =
11086                std::sync::atomic::AtomicU8::new(0);
11087            let enabled = crate::glm5_ep_grouped_prime_on() && moe_grouped_prefill_enabled();
11088            let bit = 1u8 << u8::from(enabled);
11089            if EPGP_ANNOUNCED.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
11090                eprintln!(
11091                    "[glm5-ep-grouped-prime] flag={} t={t} il={il} (announce printed in both \
11092                     arms; engagement is the dispatch counter + per-layer execute line)",
11093                    if enabled { "on" } else { "off" },
11094                );
11095            }
11096            if enabled
11097                && let Some(mut out) =
11098                    Self::moe_ffn_glm5_ep_grouped_prime(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?
11099            {
11100                Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11101                return Ok(out);
11102            }
11103        }
11104
11105        // EP dispatch diet (`MEMRA_GLM5_EP_DIET`, default OFF): same kernels, same combine
11106        // chain, restructured movement. Read per call — `=0`/unset restores the v1 walk.
11107        if crate::glm5_ep_diet_on() {
11108            let mut out = Self::moe_ffn_glm5_ep_diet(e, m, ep, z, &sel_all, &w_all, t, cfg, il)?;
11109            Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut out)?;
11110            return Ok(out);
11111        }
11112
11113        let mut moe_out = e.zeros(t * n_embd)?;
11114        use crate::glm5_tp_transport::Glm5TpTransport as TpXport;
11115        let hop = ep.rt.hop(e);
11116        // Peer replicas of `z`, one arm each (lane/glm5-tp-transport):
11117        //   host-canonical — v1's EXACT pattern: one draining `dtoh` of the whole block here,
11118        //     then one row `htod` per token PER PEER RANK inside the loop (at two ranks that
11119        //     is byte- and hop-identical to v1). Preserved hop-for-hop so
11120        //     `MEMRA_GLM5_TP_TRANSPORT=0` reproduces the banked v1 walk, not a faster cousin.
11121        //   peer-pull — ONE device copy of the whole `[t, n_embd]` block per peer rank; rows
11122        //     are sliced out of it. Same bytes on the peers, the host uploads and drains
11123        //     removed.
11124        let z_host = match hop.transport {
11125            TpXport::HostCanonical => Some(crate::glm5_tp_transport::host_stage_block(
11126                &hop,
11127                0,
11128                z,
11129                t * n_embd,
11130            )?),
11131            TpXport::PeerPull => None,
11132        };
11133        let z_peer_bulks = match hop.transport {
11134            TpXport::PeerPull => Some(crate::glm5_tp_transport::fanout_f32(&hop, z, t * n_embd)?),
11135            TpXport::HostCanonical => None,
11136        };
11137        for tok in 0..t {
11138            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11139            let w = &w_all[tok * n_used..(tok + 1) * n_used];
11140            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11141            // The host-canonical arm materializes this token's peer rows here (v1's
11142            // per-token `htod`, one per peer rank); the peer-pull arm slices them out of the
11143            // bulk blocks. The holders must outlive the views, hence the two-step.
11144            let z_peer_row_holders: Option<Vec<CudaSlice<f32>>> = match &z_host {
11145                Some(h) => {
11146                    let mut rows = Vec::with_capacity(ranks - 1);
11147                    for r in 1..ranks {
11148                        rows.push(crate::glm5_tp_transport::host_row_to(
11149                            &hop,
11150                            r,
11151                            &h[tok * n_embd..(tok + 1) * n_embd],
11152                        )?);
11153                    }
11154                    Some(rows)
11155                }
11156                None => None,
11157            };
11158            // Per slot, in ROUTER SLOT ORDER: compute the UNWEIGHTED expert row on its
11159            // owner, then fmaf-accumulate on root — the plain walk's exact chain.
11160            for (j, &ex) in sel.iter().enumerate() {
11161                let ex = ex as usize;
11162                let owner = ep.owner(ex);
11163                if owner != 0 {
11164                    // Engagement counter FIRST (a red skip still counts as ROUTED).
11165                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES
11166                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11167                    // Gate red arm: dropping the peers' slots MUST diverge — the
11168                    // non-vacuity proof that the peer ranks contribute real expert work.
11169                    if matches!(
11170                        crate::glm5_tp::gate_red(),
11171                        Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11172                    ) {
11173                        continue;
11174                    }
11175                }
11176                let zin_holder;
11177                let (dev, slab, zin) = if owner == 0 {
11178                    (e, &ep.slabs[0], &zt)
11179                } else {
11180                    zin_holder = match (&z_peer_row_holders, &z_peer_bulks) {
11181                        (Some(rows), _) => rows[owner - 1].slice(0..n_embd),
11182                        (None, Some(bulks)) => {
11183                            bulks[owner - 1].slice(tok * n_embd..(tok + 1) * n_embd)
11184                        }
11185                        (None, None) => {
11186                            return Err(
11187                                "glm5 EP: neither transport arm staged the peer activation".into(),
11188                            );
11189                        }
11190                    };
11191                    (
11192                        crate::glm5_tp::rank_engine(e, rt, owner),
11193                        &ep.slabs[owner],
11194                        &zin_holder,
11195                    )
11196                };
11197                // Placement indirection: the owner's slab packs its experts in
11198                // ascending-id order; `local_of` is the slot (identical to
11199                // `ex - first_expert` under the even split).
11200                let local = ep.local_of[ex] as usize;
11201                let gl = m.gate_exps.expert_stride;
11202                let ul = m.up_exps.expert_stride;
11203                let dl = m.down_exps.expert_stride;
11204                let gate = dev.qmatvec_view(
11205                    &slab.gate,
11206                    local * gl..(local + 1) * gl,
11207                    zin,
11208                    1,
11209                    m.gate_exps.in_f,
11210                    m.gate_exps.out_f,
11211                    m.gate_exps.qtype,
11212                    m.gate_exps.row_bytes,
11213                )?;
11214                let up = dev.qmatvec_view(
11215                    &slab.up,
11216                    local * ul..(local + 1) * ul,
11217                    zin,
11218                    1,
11219                    m.up_exps.in_f,
11220                    m.up_exps.out_f,
11221                    m.up_exps.qtype,
11222                    m.up_exps.row_bytes,
11223                )?;
11224                let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
11225                Self::ffn_act_lim(
11226                    dev,
11227                    cfg,
11228                    &gate,
11229                    &up,
11230                    m.gate_exps.macro_scale(ex),
11231                    m.up_exps.macro_scale(ex),
11232                    lim_exp,
11233                    &mut act,
11234                    n_ff_exp,
11235                )?;
11236                let actv = act.slice(0..n_ff_exp);
11237                let y = dev.qmatvec_view(
11238                    &slab.down,
11239                    local * dl..(local + 1) * dl,
11240                    &actv,
11241                    1,
11242                    m.down_exps.in_f,
11243                    m.down_exps.out_f,
11244                    m.down_exps.qtype,
11245                    m.down_exps.row_bytes,
11246                )?;
11247                // The owner's row returns through the armed transport; the root row stays
11248                // put. The slot-ordered axpy below is the ONE cross-rank arithmetic site and
11249                // it reproduces the sequential walk's accumulate chain operation for
11250                // operation — unchanged by which transport delivered the row.
11251                let y_root = if owner == 0 {
11252                    y
11253                } else {
11254                    crate::glm5_tp_transport::return_row_to_root(&hop, owner, &y, n_embd)?
11255                };
11256                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
11257                e.axpy_into(
11258                    &y_root,
11259                    w[j] * m.down_exps.macro_scale(ex),
11260                    &mut dst,
11261                    n_embd,
11262                )?;
11263            }
11264        }
11265
11266        Self::moe_shexp_add(e, m, z, zq8, t, cfg, lim_shexp, &mut moe_out)?;
11267        Ok(moe_out)
11268    }
11269
11270    /// The DIETED glm5 EP walk (`MEMRA_GLM5_EP_DIET`, lane/glm5-ep-diet): the v1 walk's
11271    /// per-slot expert kernels and its exact slot-ordered combine chain, with the data
11272    /// movement restructured in whole groups (the tp2-battery's measured 13-18 ms/token v1
11273    /// join+dispatch tax, attributed to per-token host fan-out x42 layers + ~4-5 sync
11274    /// peer-slot round-trips/layer + interleaved host-blocked issue):
11275    ///
11276    ///   1. ONE bulk peer z fan-out per layer-call ([t, n_embd] in one upload; SKIPPED
11277    ///      entirely when the call routed no peer-owned expert — the placement-map
11278    ///      multiplier: a single-rank layer-call moves zero activation bytes off root).
11279    ///   2. Peer-owned rows compute back-to-back on the peer stream into a compact block
11280    ///      (issue order cannot change bytes: every row is an independent per-expert
11281    ///      program; the combine order below is fixed by the id table, not by issue).
11282    ///   3. Root-owned rows compute on the root stream, un-blocked by peer returns.
11283    ///   4. ONE bulk peer return (peer DtoH + root HtoD of the compact block) replaces the
11284    ///      per-slot round-trip dribble.
11285    ///   5. ONE `moe_pairs_scatter` launch applies the per-token slot-ordered fmaf chain —
11286    ///      the kernel header carries the byte-identity contract vs the zeros +
11287    ///      sequential-`axpy_f32` chain this replaces, and the weights are the SAME host
11288    ///      fold (`w * macro_scale(ex)`) v1 passed per launch.
11289    ///
11290    /// BYTE-IDENTICAL to the v1 walk (and to plain, wherever v1 is) by construction: same
11291    /// kernels over the same bytes; copies (dtod / bulk DtoH+HtoD) preserve bits; the one
11292    /// arithmetic site keeps its exact chain. Transport stays HOST-CANONICAL: the two bulk
11293    /// hops of steps 1 and 4 are the named native-P2P swap points for the box arc (the
11294    /// `MEMRA_STEP_TP_BULK_P2P` precedent: peer copies 61,452 -> ~21/layer on step; the
11295    /// glm5 seam inherits `configure_native_p2p` but does NOT wire it on the rig — the
11296    /// same-device dual-context emulation has no real peer transport to qualify).
11297    #[allow(clippy::too_many_arguments)]
11298    fn moe_ffn_glm5_ep_diet(
11299        e: &Engine,
11300        m: &MoeWeights,
11301        ep: &crate::glm5_tp::Glm5EpExps,
11302        z: &CudaSlice<f32>,
11303        sel_all: &[u32],
11304        w_all: &[f32],
11305        t: usize,
11306        cfg: &ModelConfig,
11307        il: u16,
11308    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11309        use std::sync::atomic::Ordering;
11310        let moe = cfg
11311            .moe
11312            .as_ref()
11313            .ok_or("glm5 EP execution requires MoE model metadata")?;
11314        let n_embd = cfg.n_embd as usize;
11315        let n_used = moe.expert_used_count as usize;
11316        let n_ff_exp = moe.expert_ff_length as usize;
11317        let lim_exp = cfg.clamp_exp_at(il as u32);
11318        let rt = &ep.rt;
11319        let n_pairs = t * n_used;
11320        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11321            return Err("glm5 EP diet geometry".into());
11322        }
11323        let red_skip_peer = matches!(
11324            crate::glm5_tp::gate_red(),
11325            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11326        );
11327
11328        // Slab-position table: root-owned pairs pack the slab head in pair order; each peer
11329        // rank's pairs pack a contiguous tail segment (so every rank's bulk return is ONE
11330        // contiguous upload). `ids[p]` is pair p's slab row; the scatter walks ids in slot
11331        // order per token, which is what pins the combine chain to v1's regardless of
11332        // packing. At two ranks this is byte-for-byte the original head/tail split.
11333        let ranks = ep.ranks();
11334        let mut per_rank = vec![0usize; ranks];
11335        for &s in sel_all.iter().take(n_pairs) {
11336            let ex = s as usize;
11337            if ex >= ep.owner_of.len() {
11338                return Err(format!("glm5 EP diet: selection {ex} outside the bank").into());
11339            }
11340            per_rank[ep.owner(ex)] += 1;
11341        }
11342        let mut base = vec![0usize; ranks];
11343        for r in 1..ranks {
11344            base[r] = base[r - 1] + per_rank[r - 1];
11345        }
11346        let mut ids = vec![0i32; n_pairs];
11347        {
11348            let mut k = vec![0usize; ranks];
11349            for (p, id) in ids.iter_mut().enumerate() {
11350                let r = ep.owner(sel_all[p] as usize);
11351                *id = (base[r] + k[r]) as i32;
11352                k[r] += 1;
11353            }
11354        }
11355
11356        crate::glm5_tp::GLM5_EP_DIET_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11357        for r in 1..ranks {
11358            crate::glm5_tp::GLM5_EP_DIET_FANOUT_UPLOADS_AVOIDED.fetch_add(
11359                if per_rank[r] > 0 {
11360                    (t - 1) as u64
11361                } else {
11362                    t as u64
11363                },
11364                Ordering::Relaxed,
11365            );
11366        }
11367        static EP_DIET_MARKED: std::sync::atomic::AtomicBool =
11368            std::sync::atomic::AtomicBool::new(false);
11369        if !EP_DIET_MARKED.swap(true, Ordering::Relaxed) {
11370            eprintln!(
11371                "[glm5-ep-diet] engaged: bulk fan-out + compact peer staging + single \
11372                 slot-ordered scatter combine; per-slot host round-trips removed \
11373                 transport={} performance_claim=false",
11374                ep.rt.transport.name(),
11375            );
11376        }
11377
11378        // Per-slot expert program, shared verbatim with the v1 walk (same kernels, same
11379        // argument order): gate/up qmatvec + ffn_act_lim + down qmatvec on the OWNING rank.
11380        let expert_row = |dev: &Engine,
11381                          slab: &crate::glm5_tp::EpRankSlab,
11382                          zin: &cudarc::driver::CudaView<f32>,
11383                          ex: usize|
11384         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11385            let local = ep.local_of[ex] as usize;
11386            let gl = m.gate_exps.expert_stride;
11387            let ul = m.up_exps.expert_stride;
11388            let dl = m.down_exps.expert_stride;
11389            let gate = dev.qmatvec_view(
11390                &slab.gate,
11391                local * gl..(local + 1) * gl,
11392                zin,
11393                1,
11394                m.gate_exps.in_f,
11395                m.gate_exps.out_f,
11396                m.gate_exps.qtype,
11397                m.gate_exps.row_bytes,
11398            )?;
11399            let up = dev.qmatvec_view(
11400                &slab.up,
11401                local * ul..(local + 1) * ul,
11402                zin,
11403                1,
11404                m.up_exps.in_f,
11405                m.up_exps.out_f,
11406                m.up_exps.qtype,
11407                m.up_exps.row_bytes,
11408            )?;
11409            let mut act = dev.uninit(n_ff_exp)?; // activation fully overwrites
11410            Self::ffn_act_lim(
11411                dev,
11412                cfg,
11413                &gate,
11414                &up,
11415                m.gate_exps.macro_scale(ex),
11416                m.up_exps.macro_scale(ex),
11417                lim_exp,
11418                &mut act,
11419                n_ff_exp,
11420            )?;
11421            let actv = act.slice(0..n_ff_exp);
11422            dev.qmatvec_view(
11423                &slab.down,
11424                local * dl..(local + 1) * dl,
11425                &actv,
11426                1,
11427                m.down_exps.in_f,
11428                m.down_exps.out_f,
11429                m.down_exps.qtype,
11430                m.down_exps.row_bytes,
11431            )
11432        };
11433
11434        // Pass 1 — PEERS: one bulk fan-out per pair-owning rank, then every owned row
11435        // back-to-back on that rank's stream into its compact block. No host boundary until
11436        // the bulk returns.
11437        let hop = ep.rt.hop(e);
11438        let mut y_peer_blks: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11439        for r in 1..ranks {
11440            if per_rank[r] == 0 {
11441                continue;
11442            }
11443            let dev = crate::glm5_tp::rank_engine(e, rt, r);
11444            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only
11445            // (a rank with zero owned pairs moves zero activation bytes off root).
11446            let z_r = crate::glm5_tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11447            // Under the skip-peer-combine red the block stays ZERO for skipped rows (a red
11448            // must drop the peer contribution loudly, never multiply garbage into the chain).
11449            let mut blk = if red_skip_peer {
11450                dev.zeros(per_rank[r] * n_embd)?
11451            } else {
11452                dev.uninit(per_rank[r] * n_embd)?
11453            };
11454            let mut k = 0usize;
11455            for tok in 0..t {
11456                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11457                for &ex in sel.iter() {
11458                    let ex = ex as usize;
11459                    if ep.owner(ex) != r {
11460                        continue;
11461                    }
11462                    // Engagement counters FIRST (a red skip still counts as ROUTED).
11463                    crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11464                    crate::glm5_tp::GLM5_EP_DIET_PEER_ROUNDTRIPS_AVOIDED
11465                        .fetch_add(1, Ordering::Relaxed);
11466                    if red_skip_peer {
11467                        k += 1;
11468                        continue;
11469                    }
11470                    let zt_r = z_r.slice(tok * n_embd..(tok + 1) * n_embd);
11471                    let y = expert_row(dev, &ep.slabs[r], &zt_r, ex)?;
11472                    dev.copy_into(&mut blk, k * n_embd, &y, n_embd)?;
11473                    k += 1;
11474                }
11475            }
11476            y_peer_blks[r] = Some(blk);
11477        }
11478
11479        // Pass 2 — ROOT: every root-owned row into the slab head, never blocked on a peer
11480        // return (the v1 walk interleaved root issue behind per-slot peer syncs).
11481        let mut y_all = e.uninit(n_pairs * n_embd)?;
11482        {
11483            let mut k = 0usize;
11484            for tok in 0..t {
11485                let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
11486                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
11487                for &ex in sel.iter() {
11488                    let ex = ex as usize;
11489                    if ep.owner(ex) != 0 {
11490                        continue;
11491                    }
11492                    let y = expert_row(e, &ep.slabs[0], &zt, ex)?;
11493                    e.copy_into(&mut y_all, k * n_embd, &y, n_embd)?;
11494                    k += 1;
11495                }
11496            }
11497        }
11498
11499        // SWAP POINT 2 (bulk returns) — Pass 3: ONE rank->root block move into each rank's
11500        // tail segment. On host-canonical each is the ONE draining peer sync of that rank's
11501        // layer-call share, exactly as before; on peer-pull each is one event-ordered device
11502        // copy and no host boundary at all.
11503        for r in 1..ranks {
11504            if let Some(blk) = &y_peer_blks[r] {
11505                crate::glm5_tp_transport::return_block_to_root(
11506                    &hop,
11507                    r,
11508                    blk,
11509                    &mut y_all,
11510                    base[r] * n_embd,
11511                    per_rank[r] * n_embd,
11512                )?;
11513                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11514            }
11515        }
11516
11517        // Pass 4 — ONE combine launch. Weights are v1's exact host fold, placed at slab
11518        // positions; the scatter walks each token's pairs in SLOT order (ids[p], p pair-major),
11519        // reproducing zeros + n_used sequential axpy_f32 per the kernel's bit contract.
11520        let mut wd = vec![0f32; n_pairs];
11521        for ((&id, &w), &s) in ids.iter().zip(w_all.iter()).zip(sel_all.iter()) {
11522            wd[id as usize] = w * m.down_exps.macro_scale(s as usize);
11523        }
11524        let pw = e.htod(&wd)?;
11525        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
11526        let toff_d = e.htod_i32(&toff)?;
11527        let ids_d = e.htod_i32(&ids)?;
11528        let mut moe_out = e.uninit(t * n_embd)?; // the scatter fully overwrites
11529        e.moe_pairs_scatter(&y_all, &pw, &toff_d, &ids_d, &mut moe_out, t, n_embd)?;
11530        Ok(moe_out)
11531    }
11532
11533    /// The EP GROUPED PRIME (`MEMRA_GLM5_EP_GROUPED_PRIME`, lane/glm5-ep-diet): the plain
11534    /// walk's grouped-prefill program (`moe_ffn_grouped_prefill_sigmoid`, default ON on the
11535    /// serving artifact — 85 -> 616-639 tok/s prefill in its box A/B) split by expert
11536    /// ownership. Per rank: expert-major CSR over the rank's OWNED (token, expert) pairs,
11537    /// one grouped f16 GEMM per projection over the rank's resident EP slab (pointer tables
11538    /// minted at arm time), the PRE-clamped SwiGLU epilogue, the per-expert macro folds,
11539    /// and the slot-ordered per-token scatter — all composed from the SAME Engine calls the
11540    /// plain arm makes, so per-expert GEMM bytes match the plain grouped arm's (grouping is
11541    /// per expert, and an expert's token rows all live on its owner). The ONE new
11542    /// reassociation is the per-token partial add (root chain + peer chain instead of one
11543    /// 8-term chain) — band-gated on minted NVFP4 slabs (`glm5_ep_diet_doors_gpu`), never
11544    /// claimed byte.
11545    ///
11546    /// Returns `Ok(None)` — fall closed to the sequential EP walk — whenever the plain
11547    /// grouped arm's own admission would (f16g-ineligible qtypes, bank/top-k shape, no
11548    /// sigmoid clamp form). The rig fixture's Q8_0 bank therefore ALWAYS falls closed;
11549    /// `glm5-tp-gate`'s grouped arm proves exactly that (dispatch counter pinned 0, walk
11550    /// bytes unchanged).
11551    #[allow(clippy::too_many_arguments)]
11552    fn moe_ffn_glm5_ep_grouped_prime(
11553        e: &Engine,
11554        m: &MoeWeights,
11555        ep: &crate::glm5_tp::Glm5EpExps,
11556        z: &CudaSlice<f32>,
11557        sel_all: &[u32],
11558        w_all: &[f32],
11559        t: usize,
11560        cfg: &ModelConfig,
11561        il: u16,
11562    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11563        use std::sync::atomic::Ordering;
11564        let moe = cfg
11565            .moe
11566            .as_ref()
11567            .ok_or("glm5 EP grouped prime requires MoE model metadata")?;
11568        let n_embd = cfg.n_embd as usize;
11569        let n_expert = moe.expert_count as usize;
11570        let n_used = moe.expert_used_count as usize;
11571        let n_ff_exp = moe.expert_ff_length as usize;
11572        // The plain grouped arm's admission, mirrored term for term (fall closed, never a
11573        // new admission class). MEMRA_MOE_GATE is the sequential byte-identity oracle; this
11574        // arm is a band class and must not shadow that comparison.
11575        if crate::moe_f16g_mode() == 0 || std::env::var("MEMRA_MOE_GATE").is_ok() {
11576            return Ok(None);
11577        }
11578        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
11579            && f16g_proj_ok(m.up_exps.qtype, n_embd)
11580            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
11581        {
11582            return Ok(None);
11583        }
11584        if n_expert > 512 || n_used == 0 || n_used > 8 {
11585            return Ok(None);
11586        }
11587        let lim_exp = cfg.clamp_exp_at(il as u32);
11588        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
11589            return Err(
11590                "EP grouped prime is qualified for the PRE-clamped SwiGLU form only; \
11591                 a POST-clamp layer must ride the sequential arm"
11592                    .into(),
11593            );
11594        }
11595        let n_pairs = t * n_used;
11596        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
11597            return Err("EP grouped prime geometry".into());
11598        }
11599        let rt = &ep.rt;
11600        let red_skip_peer = matches!(
11601            crate::glm5_tp::gate_red(),
11602            Ok(Some(crate::glm5_tp::GateRed::SkipPeerCombine))
11603        );
11604
11605        // One rank's whole grouped program: CSR over OWNED pairs -> grouped gate/up GEMMs ->
11606        // macro folds -> PRE-clamped epilogue -> grouped down GEMM -> CSR->local permute ->
11607        // slot-ordered scatter into the rank partial [t, n_embd] (empty token windows write
11608        // 0.0 — the scatter fully overwrites, so partials add cleanly on root).
11609        let rank_pass = |dev: &Engine,
11610                         rank: u8,
11611                         ptr_row: &CudaSlice<u64>,
11612                         z_dev: &CudaSlice<f32>|
11613         -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
11614            // Expert-major CSR restricted to this rank, local pair index l in ascending
11615            // global-pair order (so per-token slot order == ascending l).
11616            let mut buckets_l: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
11617            let mut local_tok = Vec::new(); // token of local pair l
11618            let mut local_ex = Vec::new(); // expert of local pair l (macro folds)
11619            let mut local_wd = Vec::new(); // v1's exact weight fold, at local positions
11620            let mut local_count_per_tok = vec![0i32; t];
11621            for p in 0..n_pairs {
11622                let ex = sel_all[p] as usize;
11623                if ex >= n_expert {
11624                    return Err(format!("EP grouped prime selection {ex} >= {n_expert}").into());
11625                }
11626                if ep.owner(ex) != rank as usize {
11627                    continue;
11628                }
11629                let l = local_tok.len() as i32;
11630                buckets_l[ex].push(l);
11631                let tok = p / n_used;
11632                local_tok.push(tok as i32);
11633                local_ex.push(ex);
11634                local_wd.push(w_all[p] * m.down_exps.macro_scale(ex));
11635                local_count_per_tok[tok] += 1;
11636            }
11637            let n_owned = local_tok.len();
11638            if n_owned == 0 {
11639                return Ok(None);
11640            }
11641            let mut ex_ids: Vec<i32> = Vec::new();
11642            let mut ex_off: Vec<i32> = vec![0];
11643            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_owned); // local l, CSR order
11644            let mut csr_tok: Vec<i32> = Vec::with_capacity(n_owned);
11645            for (e_id, b) in buckets_l.iter().enumerate() {
11646                if !b.is_empty() {
11647                    ex_ids.push(e_id as i32);
11648                    for &l in b {
11649                        ex_pairs.push(l);
11650                        csr_tok.push(local_tok[l as usize]);
11651                    }
11652                    ex_off.push(ex_pairs.len() as i32);
11653                }
11654            }
11655            let n_active = ex_ids.len();
11656            if n_active == 0 || n_active > 512 {
11657                return Err(format!("EP grouped prime n_active {n_active} outside 1..=512").into());
11658            }
11659
11660            let exi = dev.htod_i32(&ex_ids)?;
11661            let exo = dev.htod_i32(&ex_off)?;
11662            let exp_d = dev.htod_i32(&ex_pairs)?;
11663            let csr_tok_d = dev.htod_i32(&csr_tok)?;
11664
11665            // GATE/UP grouped GEMMs over the rank slab, CSR order end to end.
11666            let (z16, zs) = dev.moe_f16g_act(z_dev, Some(&csr_tok_d), n_embd, n_owned)?;
11667            let mut g = dev.moe_f16_grouped(
11668                ptr_row,
11669                0,
11670                n_expert,
11671                &exi,
11672                &ex_off,
11673                &exo,
11674                &z16,
11675                &zs,
11676                n_embd,
11677                n_ff_exp,
11678                n_active,
11679                n_owned,
11680                m.gate_exps.qtype,
11681                m.gate_exps.row_bytes,
11682            )?;
11683            if m.gate_exps.macros.is_some() {
11684                let mg: Vec<f32> = ex_pairs
11685                    .iter()
11686                    .map(|&l| m.gate_exps.macro_scale(local_ex[l as usize]))
11687                    .collect();
11688                let mg_d = dev.htod(&mg)?;
11689                dev.scale_rows(&mut g, &mg_d, n_ff_exp, n_owned)?;
11690            }
11691            let mut u = dev.moe_f16_grouped(
11692                ptr_row,
11693                1,
11694                n_expert,
11695                &exi,
11696                &ex_off,
11697                &exo,
11698                &z16,
11699                &zs,
11700                n_embd,
11701                n_ff_exp,
11702                n_active,
11703                n_owned,
11704                m.up_exps.qtype,
11705                m.up_exps.row_bytes,
11706            )?;
11707            if m.up_exps.macros.is_some() {
11708                let mu: Vec<f32> = ex_pairs
11709                    .iter()
11710                    .map(|&l| m.up_exps.macro_scale(local_ex[l as usize]))
11711                    .collect();
11712                let mu_d = dev.htod(&mu)?;
11713                dev.scale_rows(&mut u, &mu_d, n_ff_exp, n_owned)?;
11714            }
11715
11716            // Epilogue: PRE-clamped SwiGLU (POST refused above), plain-silu pair otherwise.
11717            let act = match lim_exp {
11718                Some(SwigluClamp::Pre(limit)) => {
11719                    let mut a = dev.uninit(n_owned * n_ff_exp)?;
11720                    dev.swiglu_preclamped_mul_scaled(
11721                        &g,
11722                        &u,
11723                        1.0,
11724                        1.0,
11725                        limit,
11726                        &mut a,
11727                        n_owned * n_ff_exp,
11728                    )?;
11729                    a
11730                }
11731                None => dev.moe_pairs_silu_mul(&g, &u, n_owned * n_ff_exp)?,
11732                Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
11733            };
11734
11735            // DOWN grouped GEMM, permute CSR -> local pair order, slot-ordered scatter.
11736            let (a16, a_s) = dev.moe_f16g_act(&act, None, n_ff_exp, n_owned)?;
11737            let d_csr = dev.moe_f16_grouped(
11738                ptr_row,
11739                2,
11740                n_expert,
11741                &exi,
11742                &ex_off,
11743                &exo,
11744                &a16,
11745                &a_s,
11746                n_ff_exp,
11747                n_embd,
11748                n_active,
11749                n_owned,
11750                m.down_exps.qtype,
11751                m.down_exps.row_bytes,
11752            )?;
11753            let y_local = dev.rows_permute(&d_csr, &exp_d, n_owned, n_embd)?;
11754            let mut toff: Vec<i32> = Vec::with_capacity(t + 1);
11755            let mut acc = 0i32;
11756            toff.push(0);
11757            for &c in &local_count_per_tok {
11758                acc += c;
11759                toff.push(acc);
11760            }
11761            let tids: Vec<i32> = (0..n_owned as i32).collect();
11762            let pw = dev.htod(&local_wd)?;
11763            let toff_d = dev.htod_i32(&toff)?;
11764            let tids_d = dev.htod_i32(&tids)?;
11765            let mut partial = dev.uninit(t * n_embd)?; // scatter fully overwrites
11766            dev.moe_pairs_scatter(&y_local, &pw, &toff_d, &tids_d, &mut partial, t, n_embd)?;
11767            Ok(Some(partial))
11768        };
11769
11770        // Peer passes first (their GEMMs overlap root's), each on its own runtime binding —
11771        // the grouped-MoE FFI follows the RUNTIME device, not cudarc's pushed context
11772        // (`bind_runtime_device`'s contract). Engagement counters count ROUTED peer pairs
11773        // before any red skip, exactly like the sequential walk.
11774        let ranks = ep.ranks();
11775        let n_peer_pairs = sel_all
11776            .iter()
11777            .take(n_pairs)
11778            .filter(|&&ex| ep.owner(ex as usize) != 0)
11779            .count() as u64;
11780        crate::glm5_tp::GLM5_EP_PEER_SLOT_DISPATCHES.fetch_add(n_peer_pairs, Ordering::Relaxed);
11781        let hop = ep.rt.hop(e);
11782        let mut peer_partials: Vec<Option<CudaSlice<f32>>> = (0..ranks).map(|_| None).collect();
11783        for r in 1..ranks {
11784            let rank_owns_pairs = sel_all
11785                .iter()
11786                .take(n_pairs)
11787                .any(|&ex| ep.owner(ex as usize) == r);
11788            if !rank_owns_pairs {
11789                continue;
11790            }
11791            let dev = crate::glm5_tp::rank_engine(e, rt, r);
11792            // SWAP POINT 1 (bulk fan-out) — the named transport shape, to this rank only.
11793            let z_r = crate::glm5_tp_transport::fanout_f32_to(&hop, r, z, t * n_embd)?;
11794            dev.bind_runtime_device(dev.ctx().ordinal() as i32)?;
11795            let res = rank_pass(dev, r as u8, &ep.ptr_rows[r], &z_r);
11796            e.bind_runtime_device(e.ctx().ordinal() as i32)?;
11797            peer_partials[r] = res?;
11798        }
11799        let root_partial = rank_pass(e, 0, &ep.ptr_rows[0], z)?;
11800
11801        // Root combine: root partial + bulk-returned peer partials (SWAP POINT 2, the named
11802        // transport shape). One partial add per contributing rank — the same reassociation
11803        // class the two-rank arm band-gated (root chain + per-rank chains instead of one
11804        // 8-term chain), never claimed byte. The skip-peer-combine red drops every peer
11805        // partial AFTER counting — the loud non-vacuity arm.
11806        let mut out = match root_partial {
11807            Some(p) => p,
11808            None => e.zeros(t * n_embd)?,
11809        };
11810        for r in 1..ranks {
11811            if let Some(pp) = &peer_partials[r]
11812                && !red_skip_peer
11813            {
11814                let pp_root =
11815                    crate::glm5_tp_transport::return_row_to_root(&hop, r, pp, t * n_embd)?;
11816                let mut dst = out.slice_mut(0..t * n_embd);
11817                e.axpy_into(&pp_root, 1.0, &mut dst, t * n_embd)?;
11818                crate::glm5_tp::GLM5_EP_DIET_BULK_RETURNS.fetch_add(1, Ordering::Relaxed);
11819            }
11820        }
11821        crate::glm5_tp::GLM5_EP_GROUPED_PRIME_DISPATCHES.fetch_add(1, Ordering::Relaxed);
11822        static EPGP_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
11823        let layer_bit = 1u64 << (il as u64 % 64);
11824        if EPGP_LOGGED.fetch_or(layer_bit, Ordering::Relaxed) & layer_bit == 0 {
11825            eprintln!(
11826                "[glm5-ep-grouped-prime] execute layer={il} tokens={t} \
11827                 provenance=ep-rank-slabs router=sigmoid-host-oracle epilogue=pre-clamped \
11828                 combine=rank-partial-add transport={} performance_claim=false \
11829                 (logged once per layer)",
11830                hop.transport.name(),
11831            );
11832        }
11833        Ok(Some(out))
11834    }
11835
11836    /// Step 3 of the MoE body — SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z —
11837    /// qwen35moe only. OLMoE and most vanilla MoE have NO shared expert (the shexp tensors
11838    /// are absent / `None`); skip it then. gate_inp_shexp is OPTIONAL: qwen35moe gates the
11839    /// shared expert (sigmoid(gate_inp) x sh); MiniMax-M3 (DeepSeek-V3 class) has NO shexp
11840    /// gate — the shared expert adds directly. (Extracted verbatim from the sequential body
11841    /// so the glm5 EP-2 walk adds the ROOT-owned shared expert through the identical
11842    /// program.)
11843    #[allow(clippy::too_many_arguments)]
11844    fn moe_shexp_add(
11845        e: &Engine,
11846        m: &MoeWeights,
11847        z: &CudaSlice<f32>,
11848        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
11849        t: usize,
11850        cfg: &ModelConfig,
11851        lim_shexp: Option<memra_gguf::config::SwigluClamp>,
11852        moe_out: &mut CudaSlice<f32>,
11853    ) -> Result<(), Box<dyn std::error::Error>> {
11854        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
11855            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
11856        {
11857            let n_embd = cfg.n_embd as usize;
11858            let n_ff_sh = gate_shexp.out_features(); // 512
11859            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
11860            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
11861            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
11862            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
11863            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
11864            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
11865            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
11866            let verify_t = t > 1 && t < PRIME_MIN_T;
11867            let (sg_gate, sg_up) = if t == 1 {
11868                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
11869            } else if verify_t {
11870                (
11871                    e.matmul_decode_exact(gate_shexp, z, t)?,
11872                    e.matmul_decode_exact(up_shexp, z, t)?,
11873                )
11874            } else {
11875                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
11876            };
11877            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
11878            Self::ffn_act_lim(
11879                e,
11880                cfg,
11881                &sg_gate,
11882                &sg_up,
11883                1.0,
11884                1.0,
11885                lim_shexp,
11886                &mut sa,
11887                t * n_ff_sh,
11888            )?;
11889            let sh = if verify_t {
11890                e.matmul_decode_exact(down_shexp, &sa, t)?
11891            } else {
11892                e.matmul(down_shexp, &sa, t)?
11893            }; // [T, n_embd]
11894
11895            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
11896            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
11897            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
11898            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
11899            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
11900            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
11901            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
11902            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
11903            // expert's contribution into every token's residual, so under cross-request
11904            // concat prefill a session's hidden state depended on its co-arrivals' token
11905            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
11906            // DOOR H (`MEMRA_GLM5_HTOD_DIET`): glm5 has no `ffn_gate_inp_shexp`, so this is the
11907            // LIVE arm on the serving artifact and it re-uploaded a constant `vec![1.0f32; t]`
11908            // on every MoE layer-call — 42 pageable HtoD per ship round (26.9% of the round's
11909            // 156). The resident ones buffer feeds the SAME `add_scaled_rows_f32` kernel the
11910            // same 1.0 values, so the arms are bit-identical.
11911            if m.gate_inp_shexp.is_none() && crate::glm5_htod_diet_on() {
11912                crate::GLM5_HTOD_DIET_AVOIDED.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11913                e.add_scaled_rows_ones(&sh, moe_out, n_embd, t)?;
11914                return Ok(());
11915            }
11916            let g = match &m.gate_inp_shexp {
11917                Some(gate_inp_shexp) => {
11918                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
11919                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
11920                    } else {
11921                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
11922                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
11923                        e.sigmoid(&gs, &mut g, t)?;
11924                        g
11925                    }
11926                }
11927                None => e.htod(&vec![1.0f32; t])?,
11928            };
11929            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
11930            e.add_scaled_rows(&sh, &g, moe_out, n_embd, t)?;
11931        }
11932        Ok(())
11933    }
11934
11935    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
11936    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
11937    pub fn stage1_h2d_per_token(&self) -> u64 {
11938        use crate::hybrid::Ffn;
11939        let n_used = self
11940            .cfg
11941            .moe
11942            .as_ref()
11943            .map(|m| m.expert_used_count as u64)
11944            .unwrap_or(0);
11945        let mut bytes = 0u64;
11946        for l in self.layers.iter() {
11947            if let Ffn::Moe(m) = &l.ffn {
11948                bytes += n_used
11949                    * (m.gate_exps.max_expert_bytes()
11950                        + m.up_exps.max_expert_bytes()
11951                        + m.down_exps.max_expert_bytes()) as u64;
11952            }
11953        }
11954        bytes
11955    }
11956
11957    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
11958    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
11959    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
11960    pub(crate) fn max_moe_block(&self) -> usize {
11961        use crate::hybrid::Ffn;
11962        let mut mx = 0usize;
11963        let mut scan = |ffn: &Ffn| {
11964            if let Ffn::Moe(m) = ffn {
11965                mx = mx
11966                    .max(m.gate_exps.max_expert_bytes())
11967                    .max(m.up_exps.max_expert_bytes())
11968                    .max(m.down_exps.max_expert_bytes());
11969            }
11970        };
11971        for l in self.layers.iter() {
11972            scan(&l.ffn);
11973        }
11974        if let Some(mtp) = self.mtp.as_ref() {
11975            scan(&mtp.ffn);
11976        }
11977        mx
11978    }
11979
11980    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
11981    /// but have no bytes and therefore consume no residency slot.
11982    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
11983        use crate::hybrid::Ffn;
11984        let mut sizes = Vec::new();
11985        let mut scan = |ffn: &Ffn| {
11986            let Ffn::Moe(m) = ffn else { return };
11987            for ex in 0..m.gate_exps.n_expert {
11988                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
11989                    continue;
11990                }
11991                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
11992                    let len = exps.expert_layout(ex).len;
11993                    if len > 0 {
11994                        sizes.push(len);
11995                    }
11996                }
11997            }
11998        };
11999        for layer in &self.layers {
12000            scan(&layer.ffn);
12001        }
12002        if let Some(mtp) = &self.mtp {
12003            scan(&mtp.ffn);
12004        }
12005        sizes
12006    }
12007
12008    /// Persist the frozen residency set so a later process can restage it directly and skip
12009    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
12010    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
12011    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
12012    /// post-freeze argmax gate still validates the serving assignment.
12013    pub fn save_cpu_expert_residency_profile(
12014        &self,
12015        e: &Engine,
12016        path: &std::path::Path,
12017    ) -> Result<(), Box<dyn std::error::Error>> {
12018        let Some(ids) = e.export_moe_residency() else {
12019            return Err("no MoE residency cache to persist".into());
12020        };
12021        let mut body = format!(
12022            "memra-freeze-profile v1 max_block={} blocks={}\n",
12023            self.max_moe_block(),
12024            ids.len()
12025        );
12026        for (layer, proj, ex) in &ids {
12027            body.push_str(&format!("{layer} {proj} {ex}\n"));
12028        }
12029        let tmp = path.with_extension("tmp");
12030        std::fs::write(&tmp, body)?;
12031        std::fs::rename(&tmp, path)?;
12032        println!(
12033            "[moe-cache] freeze profile saved: {} blocks -> {}",
12034            ids.len(),
12035            path.display()
12036        );
12037        Ok(())
12038    }
12039
12040    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
12041    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
12042    /// missing or its header does not match this model's slot geometry.
12043    pub fn restore_cpu_expert_residency_profile(
12044        &self,
12045        e: &Engine,
12046        path: &std::path::Path,
12047    ) -> Result<bool, Box<dyn std::error::Error>> {
12048        use crate::hybrid::Ffn;
12049        use crate::moe_cache::BlockId;
12050        let Ok(content) = std::fs::read_to_string(path) else {
12051            return Ok(false);
12052        };
12053        let mut lines = content.lines();
12054        let Some(header) = lines.next() else {
12055            return Ok(false);
12056        };
12057        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
12058        if !header.starts_with(&expected) {
12059            println!(
12060                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
12061                path.display()
12062            );
12063            return Ok(false);
12064        }
12065        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
12066            std::collections::HashMap::new();
12067        for line in lines {
12068            let mut fields = line.split_whitespace();
12069            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
12070            else {
12071                continue;
12072            };
12073            let (Ok(layer), Ok(proj), Ok(ex)) =
12074                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
12075            else {
12076                continue;
12077            };
12078            by_layer
12079                .entry(layer)
12080                .or_default()
12081                .push(BlockId::new(layer, proj, ex));
12082        }
12083        let requested: usize = by_layer.values().map(Vec::len).sum();
12084        if requested == 0 {
12085            return Ok(false);
12086        }
12087        let max_block = self.max_moe_block();
12088        let mut restaged = 0usize;
12089        let mut stage_layer =
12090            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
12091                let Ffn::Moe(m) = ffn else { return Ok(()) };
12092                let Some(ids) = by_layer.get(&layer_index) else {
12093                    return Ok(());
12094                };
12095                e.with_moe_cache(max_block, |cache, eng| {
12096                    for id in ids {
12097                        if cache.restage_block(*id, m, eng)? {
12098                            restaged += 1;
12099                        }
12100                    }
12101                    Ok(())
12102                })
12103            };
12104        for (index, layer) in self.layers.iter().enumerate() {
12105            stage_layer(index as u16, &layer.ffn)?;
12106        }
12107        if let Some(mtp) = self.mtp.as_ref() {
12108            stage_layer(u16::MAX, &mtp.ffn)?;
12109        }
12110        e.freeze_moe_cache();
12111        println!(
12112            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
12113            path.display()
12114        );
12115        Ok(true)
12116    }
12117
12118    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
12119    pub fn freeze_cpu_expert_residency(
12120        &self,
12121        e: &Engine,
12122    ) -> Result<(), Box<dyn std::error::Error>> {
12123        e.freeze_moe_cache();
12124        Ok(())
12125    }
12126
12127    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
12128    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
12129    /// the model's activation exactly.
12130    ///
12131    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
12132    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
12133    /// form for anything that can land on a clamped layer.
12134    pub fn ffn_act(
12135        e: &Engine,
12136        cfg: &ModelConfig,
12137        gate: &CudaSlice<f32>,
12138        up: &CudaSlice<f32>,
12139        act: &mut CudaSlice<f32>,
12140        n: usize,
12141    ) -> Result<(), Box<dyn std::error::Error>> {
12142        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
12143    }
12144
12145    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
12146    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
12147    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
12148    #[allow(clippy::too_many_arguments)]
12149    pub(crate) fn ffn_act_scaled(
12150        e: &Engine,
12151        cfg: &ModelConfig,
12152        gate: &CudaSlice<f32>,
12153        up: &CudaSlice<f32>,
12154        gs: f32,
12155        us: f32,
12156        act: &mut CudaSlice<f32>,
12157        n: usize,
12158    ) -> Result<(), Box<dyn std::error::Error>> {
12159        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
12160    }
12161
12162    /// ffn_act_scaled + a PER-LAYER clamped SwiGLU. `limit`:
12163    ///   * `None`          -> the unclamped dispatch (every arch with no live clamp).
12164    ///   * `Some(Post(l))` -> step35: `min(silu(gate*gs), l) * clamp(up*us, +-l)`
12165    ///     (llama-graph.cpp:2146/1751, non-DEEPSEEK4 branch).
12166    ///   * `Some(Pre(l))`  -> glm5_next: `silu(min(gate*gs, l)) * clamp(up*us, +-l)`.
12167    ///     Callers source it from `cfg.clamp_exp_at(il)` (routed experts) or `cfg.clamp_shexp_at(il)`
12168    ///     (shared expert / dense MLP) — on step35 the two arrays are SEPARATE and a layer can have
12169    ///     one without the other. The `> 1e-6` eps gate lives in the accessors, so a `Some` here is
12170    ///     already known live. The match is exhaustive so a new clamp form cannot default to either
12171    ///     existing one.
12172    #[allow(clippy::too_many_arguments)]
12173    pub(crate) fn ffn_act_lim(
12174        e: &Engine,
12175        cfg: &ModelConfig,
12176        gate: &CudaSlice<f32>,
12177        up: &CudaSlice<f32>,
12178        gs: f32,
12179        us: f32,
12180        limit: Option<SwigluClamp>,
12181        act: &mut CudaSlice<f32>,
12182        n: usize,
12183    ) -> Result<(), Box<dyn std::error::Error>> {
12184        if let Some(m3) = cfg.m3.as_ref() {
12185            debug_assert!(
12186                limit.is_none(),
12187                "m3 swigluoai and the step35/glm5_next clamps are different archs"
12188            );
12189            return e.swigluoai_mul_scaled(
12190                gate,
12191                up,
12192                gs,
12193                us,
12194                m3.swiglu_alpha,
12195                m3.swiglu_limit,
12196                act,
12197                n,
12198            );
12199        }
12200        match limit {
12201            Some(SwigluClamp::Post(l)) => {
12202                return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
12203            }
12204            Some(SwigluClamp::Pre(l)) => {
12205                return e.swiglu_preclamped_mul_scaled(gate, up, gs, us, l, act, n);
12206            }
12207            None => {}
12208        }
12209        if gs == 1.0 && us == 1.0 {
12210            return e.silu_mul(gate, up, act, n);
12211        }
12212        e.silu_mul_scaled(gate, up, gs, us, act, n)
12213    }
12214
12215    /// The bare POST limit for the fused kernels whose epilogue HARDCODES step35's form
12216    /// (`matvec_bf16_dual_silu` / `_rows`, qmatvec.cu:10676). `Ok` = the kernel may run;
12217    /// `Err(())` = glm5_next's PRE form, which has no fused twin, and the caller MUST return its
12218    /// not-handled value so the layer falls through to the unfused `ffn_act_lim` seam. Feeding a
12219    /// PRE limit to a POST epilogue compiles, runs, and returns plausible-but-wrong logits.
12220    fn fused_post_limit(lim: Option<SwigluClamp>) -> Result<Option<f32>, ()> {
12221        match lim {
12222            None => Ok(None),
12223            Some(SwigluClamp::Post(l)) => Ok(Some(l)),
12224            Some(SwigluClamp::Pre(_)) => Err(()),
12225        }
12226    }
12227
12228    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
12229    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
12230    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
12231    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
12232    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
12233    fn moe_route(
12234        e: &Engine,
12235        logits: &CudaSlice<f32>,
12236        t: usize,
12237        n_expert: usize,
12238        n_used: usize,
12239    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12240        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
12241    }
12242
12243    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
12244    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
12245    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
12246    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
12247    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
12248    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
12249    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
12250    #[allow(clippy::too_many_arguments)]
12251    fn moe_route_sigmoid_cfg(
12252        e: &Engine,
12253        logits: &CudaSlice<f32>,
12254        t: usize,
12255        n_expert: usize,
12256        n_used: usize,
12257        m: &MoeWeights,
12258        (sf, route_norm): (f32, bool),
12259    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12260        if sigmoid_router_enabled() {
12261            return e.moe_router_sigmoid_topk_host(
12262                logits,
12263                t,
12264                n_expert,
12265                n_used,
12266                m.active_count(),
12267                &m.exp_probs_b_dev,
12268                &m.active_experts_dev,
12269                sf,
12270                route_norm,
12271            );
12272        }
12273        let lg = e.dtoh(logits)?;
12274        Self::moe_route_sigmoid_host(
12275            &lg,
12276            t,
12277            n_expert,
12278            n_used,
12279            m.exp_probs_b.as_deref(),
12280            sf,
12281            route_norm,
12282            m.active_experts.as_deref(),
12283        )
12284    }
12285
12286    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
12287    /// the existing softmax device kernel has no mask input.
12288    #[allow(clippy::excessive_precision)] // allow: literal kept verbatim from the reference/measured value
12289    fn moe_route_cfg(
12290        e: &Engine,
12291        logits: &CudaSlice<f32>,
12292        t: usize,
12293        n_expert: usize,
12294        n_used: usize,
12295        active: Option<&[bool]>,
12296    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12297        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
12298        // rollback) via the single-sync pinned readback — softmax arch only.
12299        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
12300            return e.moe_router_topk_host(logits, t, n_expert, n_used);
12301        }
12302        // Host oracle (the §D bit-identity reference).
12303        let lg = e.dtoh(logits)?; // [T*n_expert] host
12304        let mut sel = vec![0u32; t * n_used];
12305        let mut w_out = vec![0f32; t * n_used];
12306        for tok in 0..t {
12307            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12308            // softmax over ALL n_expert (stable: subtract max)
12309            let maxl = row
12310                .iter()
12311                .enumerate()
12312                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
12313                .map(|(_, &x)| x)
12314                .fold(f32::NEG_INFINITY, f32::max);
12315            let mut probs = vec![0f32; n_expert];
12316            let mut den = 0f32;
12317            for i in 0..n_expert {
12318                if active.is_some_and(|mask| !mask[i]) {
12319                    continue;
12320                }
12321                let x = (row[i] - maxl).exp();
12322                probs[i] = x;
12323                den += x;
12324            }
12325            for p in probs.iter_mut() {
12326                *p /= den;
12327            }
12328            // stable DESC sort: prob DESC, ascending-index tiebreak.
12329            let mut idx: Vec<usize> = (0..n_expert)
12330                .filter(|&i| active.is_none_or(|mask| mask[i]))
12331                .collect();
12332            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
12333            let sl = &idx[..n_used];
12334            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
12335            let mut ws: f32 = wv.iter().sum();
12336            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
12337            for x in wv.iter_mut() {
12338                *x /= ws;
12339            }
12340            for j in 0..n_used {
12341                sel[tok * n_used + j] = sl[j] as u32;
12342                w_out[tok * n_used + j] = wv[j];
12343            }
12344        }
12345        Ok((sel, w_out))
12346    }
12347
12348    #[allow(clippy::too_many_arguments)]
12349    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
12350    fn moe_route_sigmoid_with_input(
12351        e: &Engine,
12352        logits: &CudaSlice<f32>,
12353        input: &CudaSlice<f32>,
12354        t: usize,
12355        in_features: usize,
12356        n_expert: usize,
12357        n_used: usize,
12358        bias: Option<&[f32]>,
12359        (sf, route_norm): (f32, bool),
12360        active: Option<&[bool]>,
12361    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
12362        let logit_values =
12363            active_matrix_values(logits.len(), t, n_expert, "sigmoid router logits")?;
12364        let input_values =
12365            active_matrix_values(input.len(), t, in_features, "sigmoid router input")?;
12366        let (lg, input) = e.dtoh_pair_views(
12367            &logits.slice(0..logit_values),
12368            &input.slice(0..input_values),
12369        )?;
12370        let (sel, w) =
12371            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
12372        Ok((sel, w, input))
12373    }
12374
12375    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
12376    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
12377    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
12378    /// active mask, prebuilt projection descriptors) so no model reference escapes.
12379    pub fn start_moe_prefetch_predictor(
12380        &self,
12381        e: &Engine,
12382        cfg: &ModelConfig,
12383    ) -> Result<(), Box<dyn std::error::Error>> {
12384        use crate::hybrid::Ffn;
12385        let Some(sig) = cfg.sigmoid_router() else {
12386            return Err("prefetch predictor requires a sigmoid-router arch".into());
12387        };
12388        let resident: std::collections::HashSet<(u16, u8, u16)> = e
12389            .export_moe_residency()
12390            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
12391            .into_iter()
12392            .collect();
12393        let mut layers = Vec::new();
12394        for (index, layer) in self.layers.iter().enumerate() {
12395            let Ffn::Moe(m) = &layer.ffn else { continue };
12396            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
12397                continue;
12398            };
12399            let router = e.dtoh(data)?;
12400            let n_expert = m.gate_exps.n_expert;
12401            let n_embd = m.gate_exps.in_f;
12402            if router.len() != n_embd * n_expert {
12403                continue;
12404            }
12405            let build = |exps: &crate::model::HostExps| {
12406                (0..n_expert)
12407                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
12408                    .collect::<Vec<_>>()
12409            };
12410            layers.push((
12411                index as u16,
12412                crate::cpu_experts::PredictLayerInit {
12413                    router,
12414                    bias: m.exp_probs_b.clone(),
12415                    active: m.active_experts.clone(),
12416                    n_embd,
12417                    n_used: cfg
12418                        .moe
12419                        .as_ref()
12420                        .map(|moe| moe.expert_used_count as usize)
12421                        .ok_or("prefetch predictor requires MoE config")?,
12422                    sig,
12423                    weights_n_expert: n_expert,
12424                    gate: build(&m.gate_exps),
12425                    up: build(&m.up_exps),
12426                    down: build(&m.down_exps),
12427                },
12428            ));
12429        }
12430        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
12431    }
12432
12433    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
12434    /// selection math to the rollback runtime, applied to host-computed logits.
12435    #[allow(clippy::too_many_arguments)]
12436    pub fn moe_route_sigmoid_host_public(
12437        logits: &[f32],
12438        t: usize,
12439        n_expert: usize,
12440        n_used: usize,
12441        bias: Option<&[f32]>,
12442        sf: f32,
12443        route_norm: bool,
12444        active: Option<&[bool]>,
12445    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12446        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
12447    }
12448
12449    #[allow(clippy::too_many_arguments)]
12450    fn moe_route_sigmoid_host(
12451        lg: &[f32],
12452        t: usize,
12453        n_expert: usize,
12454        n_used: usize,
12455        bias: Option<&[f32]>,
12456        sf: f32,
12457        route_norm: bool,
12458        active: Option<&[bool]>,
12459    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
12460        let active_count = active
12461            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
12462            .unwrap_or(n_expert);
12463        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
12464        if lg.len() != t * n_expert {
12465            return Err(format!(
12466                "sigmoid router logits length mismatch: got {}, expected {}",
12467                lg.len(),
12468                t * n_expert,
12469            )
12470            .into());
12471        }
12472        let mut sel = vec![0u32; t * n_used];
12473        let mut w_out = vec![0f32; t * n_used];
12474        for tok in 0..t {
12475            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
12476            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
12477            // selection score = sigmoid + bias; weight = plain sigmoid.
12478            let selsc: Vec<f32> = match bias {
12479                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
12480                None => scores.clone(),
12481            };
12482            let mut idx: Vec<usize> = (0..n_expert)
12483                .filter(|&i| active.is_none_or(|mask| mask[i]))
12484                .collect();
12485            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
12486            let sl = &idx[..n_used];
12487            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
12488            if route_norm {
12489                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
12490                for x in wv.iter_mut() {
12491                    *x = *x / ws * sf;
12492                }
12493            } else {
12494                for x in wv.iter_mut() {
12495                    *x *= sf;
12496                }
12497            }
12498            for j in 0..n_used {
12499                sel[tok * n_used + j] = sl[j] as u32;
12500                w_out[tok * n_used + j] = wv[j];
12501            }
12502        }
12503        Ok((sel, w_out))
12504    }
12505
12506    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
12507    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
12508    /// macro-scaled experts, and observation modes are denied by the caller.
12509    #[allow(clippy::too_many_arguments)]
12510    fn moe_ffn_sigmoid_dev(
12511        e: &Engine,
12512        m: &MoeWeights,
12513        z: &CudaSlice<f32>,
12514        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
12515        logits: &CudaSlice<f32>,
12516        t: usize,
12517        cfg: &ModelConfig,
12518        il: u16,
12519        (scaling_factor, route_norm): (f32, bool),
12520    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12521        let moe = cfg.moe.as_ref().unwrap();
12522        let n_embd = cfg.n_embd as usize;
12523        let n_expert = moe.expert_count as usize;
12524        let n_used = moe.expert_used_count as usize;
12525        let n_ff_exp = moe.expert_ff_length as usize;
12526        let dev = m.dev_exps.as_ref().unwrap();
12527        debug_assert_eq!(dev.dev, e.ctx().ordinal());
12528        debug_assert!(m.has_uniform_expert_layout());
12529        debug_assert!(!m.has_macros);
12530
12531        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
12532            logits,
12533            t,
12534            n_expert,
12535            n_used,
12536            m.active_count(),
12537            &m.exp_probs_b_dev,
12538            &m.active_experts_dev,
12539            scaling_factor,
12540            route_norm,
12541        )?;
12542        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
12543        if let Some(fp8) = dev.fp8_blk.as_ref() {
12544            debug_assert_eq!(m.gate_exps.qtype, crate::QT_F8_E4M3_BLK);
12545            debug_assert_eq!(m.up_exps.qtype, crate::QT_F8_E4M3_BLK);
12546            debug_assert_eq!(m.down_exps.qtype, crate::QT_F8_E4M3_BLK);
12547            debug_assert_eq!(fp8.gate.rows, m.gate_exps.out_f.div_ceil(128));
12548            debug_assert_eq!(fp8.up.rows, m.up_exps.out_f.div_ceil(128));
12549            debug_assert_eq!(fp8.down.rows, m.down_exps.out_f.div_ceil(128));
12550
12551            // Official Step-3.7 FP8 uses dynamic per-token/per-128 E4M3
12552            // activations with block-128 E4M3 weights. This deliberately
12553            // simple resident reference is the correctness oracle for later
12554            // grouped and TP/EP execution. MEMRA_ST_E4M3=0 chooses the
12555            // load-time Q8 diagnostic representation, so one process never
12556            // crosses between numerical programs.
12557            let selected = e.dtoh_i32(&sel_d)?;
12558            let route_weights = e.dtoh(&w_d)?;
12559            let mut moe_out = e.zeros(t * n_embd)?;
12560            for tok in 0..t {
12561                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
12562                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
12563                for j in 0..n_used {
12564                    let pair = tok * n_used + j;
12565                    let expert = selected[pair] as usize;
12566                    let gate = Self::moe_resident_fp8_e4m3(
12567                        e,
12568                        &m.gate_exps,
12569                        &dev.gate,
12570                        &fp8.gate,
12571                        expert,
12572                        &zt,
12573                        1,
12574                    )?;
12575                    let up = Self::moe_resident_fp8_e4m3(
12576                        e, &m.up_exps, &dev.up, &fp8.up, expert, &zt, 1,
12577                    )?;
12578                    let mut act = e.uninit(n_ff_exp)?;
12579                    Self::ffn_act_lim(
12580                        e,
12581                        cfg,
12582                        &gate,
12583                        &up,
12584                        1.0,
12585                        1.0,
12586                        cfg.clamp_exp_at(il as u32),
12587                        &mut act,
12588                        n_ff_exp,
12589                    )?;
12590                    let act = act.slice(0..n_ff_exp);
12591                    let down = Self::moe_resident_fp8_e4m3(
12592                        e,
12593                        &m.down_exps,
12594                        &dev.down,
12595                        &fp8.down,
12596                        expert,
12597                        &act,
12598                        1,
12599                    )?;
12600                    e.axpy_into(&down, route_weights[pair], &mut dst, n_embd)?;
12601                }
12602            }
12603            if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12604                eprintln!(
12605                    "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} \
12606                     native=fp8blk-w8a8-e4m3-reference clamp={}",
12607                    cfg.clamp_exp_at(il as u32).is_some(),
12608                );
12609            }
12610            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12611            return Ok(moe_out);
12612        }
12613        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
12614            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12615            (combined, combined)
12616        } else {
12617            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12618        };
12619        let (zq, zd) = match (t, zq8) {
12620            (1, Some((q, d))) => (q.clone(), d.clone()),
12621            _ => e.quantize_q8_1(z, t, n_embd)?,
12622        };
12623        let n_pairs = t * n_used;
12624        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
12625            // The final Step layers retain the established separate gate/up -> clamp -> down
12626            // arithmetic. Pair rows are derived from token position; selected expert ids and
12627            // routing weights remain the device router's buffers throughout.
12628            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
12629            let pair_tok_d = e.htod_i32(&pair_tok)?;
12630            let gate = e.moe_pairs_matvec_q8(
12631                &dev.ptr_row,
12632                0,
12633                &pair_tok_d,
12634                &sel_d,
12635                &zq,
12636                &zd,
12637                n_embd,
12638                n_ff_exp,
12639                n_expert,
12640                n_pairs,
12641                m.gate_exps.qtype,
12642                gate_row_bytes,
12643            )?;
12644            let up = e.moe_pairs_matvec_q8(
12645                &dev.ptr_row,
12646                1,
12647                &pair_tok_d,
12648                &sel_d,
12649                &zq,
12650                &zd,
12651                n_embd,
12652                n_ff_exp,
12653                n_expert,
12654                n_pairs,
12655                m.up_exps.qtype,
12656                up_row_bytes,
12657            )?;
12658            let mut act = e.uninit(n_pairs * n_ff_exp)?;
12659            Self::ffn_act_lim(
12660                e,
12661                cfg,
12662                &gate,
12663                &up,
12664                1.0,
12665                1.0,
12666                cfg.clamp_exp_at(il as u32),
12667                &mut act,
12668                n_pairs * n_ff_exp,
12669            )?;
12670            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12671            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
12672            let pair_self_d = e.htod_i32(&pair_self)?;
12673            let down = e.moe_pairs_matvec_q8(
12674                &dev.ptr_row,
12675                2,
12676                &pair_self_d,
12677                &sel_d,
12678                &aq2,
12679                &ad2,
12680                n_ff_exp,
12681                n_embd,
12682                n_expert,
12683                n_pairs,
12684                m.down_exps.qtype,
12685                m.down_exps.row_bytes,
12686            )?;
12687            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12688            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12689            let tok_off_d = e.htod_i32(&tok_off)?;
12690            let tok_ids_d = e.htod_i32(&tok_ids)?;
12691            let mut output = e.uninit(t * n_embd)?;
12692            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
12693            output
12694        } else {
12695            let act = e.moe_gate_up_silu8_dev_q8_rows(
12696                &dev.ptr_row,
12697                &sel_d,
12698                &zq,
12699                &zd,
12700                t,
12701                n_embd,
12702                n_ff_exp,
12703                n_used,
12704                n_expert,
12705                m.gate_exps.qtype,
12706                m.up_exps.qtype,
12707                gate_row_bytes,
12708                up_row_bytes,
12709                &m.dev_macros,
12710            )?;
12711            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
12712            let mut output = e.uninit(t * n_embd)?;
12713            e.moe_down8_fma_dev_q8_rows_g(
12714                &dev.ptr_row,
12715                &sel_d,
12716                &w_d,
12717                &aq2,
12718                &ad2,
12719                &mut output,
12720                t,
12721                n_ff_exp,
12722                n_embd,
12723                n_used,
12724                n_expert,
12725                m.down_exps.qtype,
12726                m.down_exps.row_bytes,
12727            )?;
12728            output
12729        };
12730
12731        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
12732            eprintln!(
12733                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
12734                cfg.clamp_exp_at(il as u32).is_some(),
12735                dev.gu_il,
12736            );
12737        }
12738        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
12739        Ok(moe_out)
12740    }
12741
12742    #[allow(clippy::too_many_arguments)]
12743    fn moe_resident_fp8_e4m3(
12744        e: &Engine,
12745        exps: &crate::model::HostExps,
12746        bytes: &CudaSlice<u8>,
12747        scales: &crate::hybrid::DevExpertFp8ProjectionScales,
12748        expert: usize,
12749        x: &cudarc::driver::CudaView<f32>,
12750        m: usize,
12751    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12752        let layout = exps.expert_layout(expert);
12753        debug_assert_eq!(layout.qtype, crate::QT_F8_E4M3_BLK);
12754        debug_assert_eq!(scales.rows * scales.cols, scales.expert_stride);
12755        let byte_start = expert * exps.expert_stride;
12756        let scale_start = expert * scales.expert_stride;
12757        let weight = bytes.slice(byte_start..byte_start + layout.len);
12758        let scale = scales
12759            .scales
12760            .slice(scale_start..scale_start + scales.expert_stride);
12761        e.qmatvec_mmq_fp8_blk_view(&weight, &scale, x, m, exps.in_f, exps.out_f)
12762    }
12763
12764    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
12765    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
12766    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
12767    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
12768    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
12769    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
12770    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
12771    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
12772    fn moe_ffn_pairs(
12773        e: &Engine,
12774        m: &MoeWeights,
12775        z: &CudaSlice<f32>,
12776        logits: &CudaSlice<f32>,
12777        t: usize,
12778        cfg: &ModelConfig,
12779    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12780        let moe = cfg.moe.as_ref().unwrap();
12781        let n_embd = cfg.n_embd as usize;
12782        let n_expert = moe.expert_count as usize;
12783        let n_used = moe.expert_used_count as usize;
12784        let n_ff_exp = moe.expert_ff_length as usize;
12785        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
12786        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
12787        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
12788        // that forgets the gate fails loudly in debug instead of returning wrong logits.
12789        debug_assert!(
12790            !cfg.swiglu_clamped_anywhere(),
12791            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
12792        );
12793        let dev = m.dev_exps.as_ref().unwrap();
12794        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
12795        let (rbg_d, rbu_d) = if dev.gu_il {
12796            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
12797            (sxx, sxx)
12798        } else {
12799            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
12800        };
12801
12802        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
12803        let n_pairs = t * n_used;
12804        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
12805        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
12806        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
12807        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
12808        let pair_w: Vec<f32> = w_all.clone();
12809        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
12810        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
12811        let pt = e.htod_i32(&pair_tok)?;
12812        let px = e.htod_i32(&pair_ex)?;
12813        let pw = e.htod(&pair_w)?;
12814        let toff = e.htod_i32(&tok_off)?;
12815        let tids = e.htod_i32(&tok_ids)?;
12816
12817        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
12818        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
12819        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
12820        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
12821        for p in 0..n_pairs {
12822            by_ex[pair_ex[p] as usize].push(p as i32);
12823        }
12824        let mut ex_ids: Vec<i32> = Vec::new();
12825        let mut ex_off: Vec<i32> = vec![0];
12826        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
12827        for (ex, list) in by_ex.iter().enumerate() {
12828            if list.is_empty() {
12829                continue;
12830            }
12831            ex_ids.push(ex as i32);
12832            ex_pairs.extend_from_slice(list);
12833            ex_off.push(ex_pairs.len() as i32);
12834        }
12835        let n_active = ex_ids.len();
12836        let exi = e.htod_i32(&ex_ids)?;
12837        let exo = e.htod_i32(&ex_off)?;
12838        let exp_d = e.htod_i32(&ex_pairs)?;
12839        let _ = &px; // pair-major twin keeps it; em path uses CSR
12840
12841        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
12842        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
12843        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
12844        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
12845        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
12846        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
12847        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
12848        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
12849        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
12850        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
12851        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
12852        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
12853        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
12854        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
12855        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
12856        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
12857        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
12858        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
12859        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
12860        let mma_t = *MMA_T.get_or_init(|| {
12861            std::env::var("MEMRA_MOE_MMA_T")
12862                .ok()
12863                .and_then(|v| v.parse().ok())
12864                .unwrap_or(16)
12865        });
12866        let use_mma = std::env::var("MEMRA_MOE_MMA")
12867            .map(|v| v != "0")
12868            .unwrap_or(true)
12869            && t >= mma_t
12870            && q8_expert_dec_supported(m.gate_exps.qtype)
12871            && q8_expert_dec_supported(m.up_exps.qtype)
12872            && q8_expert_dec_supported(m.down_exps.qtype)
12873            && n_embd.is_multiple_of(256)
12874            && n_ff_exp.is_multiple_of(256);
12875        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
12876        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
12877        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
12878        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
12879        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
12880        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
12881        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
12882        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
12883        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
12884        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
12885        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
12886        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
12887        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
12888        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
12889        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
12890        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
12891            && q8_expert_dec_supported(m.up_exps.qtype)
12892            && q8_expert_dec_supported(m.down_exps.qtype)
12893            && n_embd.is_multiple_of(256)
12894            && n_ff_exp.is_multiple_of(256);
12895        let f16g_mode = crate::moe_f16g_mode();
12896        let f16g = f16g_mode != 0
12897            && t >= mma_t
12898            && (f16g_mode != 3 || !mma_capable)
12899            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
12900            && f16g_proj_ok(m.up_exps.qtype, n_embd)
12901            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
12902        if use_mma || f16g {
12903            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
12904            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
12905            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
12906            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
12907            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
12908            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
12909            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
12910            let y_down = if f16g {
12911                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
12912                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
12913                // permute at the very end back to pair-id order for the scatter.
12914                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
12915                let csr_tok_d = e.htod_i32(&csr_tok)?;
12916                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
12917                let g_csr = e.moe_f16_grouped(
12918                    &dev.ptr_row,
12919                    0,
12920                    n_expert,
12921                    &exi,
12922                    &ex_off,
12923                    &exo,
12924                    &z_f16,
12925                    &z_s,
12926                    n_embd,
12927                    n_ff_exp,
12928                    n_active,
12929                    n_pairs,
12930                    m.gate_exps.qtype,
12931                    rbg_d,
12932                )?;
12933                let u_csr = e.moe_f16_grouped(
12934                    &dev.ptr_row,
12935                    1,
12936                    n_expert,
12937                    &exi,
12938                    &ex_off,
12939                    &exo,
12940                    &z_f16,
12941                    &z_s,
12942                    n_embd,
12943                    n_ff_exp,
12944                    n_active,
12945                    n_pairs,
12946                    m.up_exps.qtype,
12947                    rbu_d,
12948                )?;
12949                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
12950                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
12951                let d_csr = e.moe_f16_grouped(
12952                    &dev.ptr_row,
12953                    2,
12954                    n_expert,
12955                    &exi,
12956                    &ex_off,
12957                    &exo,
12958                    &a_f16,
12959                    &a_s,
12960                    n_ff_exp,
12961                    n_embd,
12962                    n_active,
12963                    n_pairs,
12964                    m.down_exps.qtype,
12965                    m.down_exps.row_bytes,
12966                )?;
12967                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
12968            } else {
12969                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
12970                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
12971                let gate = e.mmq_iq_experts(
12972                    &dev.ptr_row,
12973                    0,
12974                    n_expert,
12975                    &exi,
12976                    &exo,
12977                    &exp_d,
12978                    &pt,
12979                    &z_scr,
12980                    n_embd,
12981                    n_ff_exp,
12982                    n_active,
12983                    n_pairs,
12984                    t,
12985                    m.gate_exps.qtype,
12986                    rbg_d,
12987                )?;
12988                let up = e.mmq_iq_experts(
12989                    &dev.ptr_row,
12990                    1,
12991                    n_expert,
12992                    &exi,
12993                    &exo,
12994                    &exp_d,
12995                    &pt,
12996                    &z_scr,
12997                    n_embd,
12998                    n_ff_exp,
12999                    n_active,
13000                    n_pairs,
13001                    t,
13002                    m.up_exps.qtype,
13003                    rbu_d,
13004                )?;
13005                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
13006                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
13007                // registers and writes ONLY the quantized scratch — the two-pass chain
13008                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
13009                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
13010                let a_scr = if crate::moe_fuse_actq_on() {
13011                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
13012                } else {
13013                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13014                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
13015                };
13016                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13017                let pself = e.htod_i32(&pair_self)?;
13018                e.mmq_iq_experts(
13019                    &dev.ptr_row,
13020                    2,
13021                    n_expert,
13022                    &exi,
13023                    &exo,
13024                    &exp_d,
13025                    &pself,
13026                    &a_scr,
13027                    n_ff_exp,
13028                    n_embd,
13029                    n_active,
13030                    n_pairs,
13031                    n_pairs,
13032                    m.down_exps.qtype,
13033                    m.down_exps.row_bytes,
13034                )?
13035            };
13036            let mut moe_out = e.uninit(t * n_embd)?;
13037            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13038            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13039                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13040            {
13041                let n_ff_sh = gate_shexp.out_features();
13042                let sg_gate = e.matmul(gate_shexp, z, t)?;
13043                let sg_up = e.matmul(up_shexp, z, t)?;
13044                let mut sa = e.uninit(t * n_ff_sh)?;
13045                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13046                let sh = e.matmul(down_shexp, &sa, t)?;
13047                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13048                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
13049                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
13050                // i.e. the one real prefill actually takes on a resident-expert MoE model,
13051                // so the concat-prime isolation fix has to land here as well.
13052                let g = match &m.gate_inp_shexp {
13053                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13054                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13055                    }
13056                    Some(gate_inp_shexp) => {
13057                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13058                        let mut g = e.uninit(t)?;
13059                        e.sigmoid(&gs, &mut g, t)?;
13060                        g
13061                    }
13062                    None => e.htod(&vec![1.0f32; t])?,
13063                };
13064                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13065            }
13066            return Ok(moe_out);
13067        }
13068
13069        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
13070        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
13071        let dec = std::env::var("MEMRA_MOE_DEC")
13072            .map(|v| v != "0")
13073            .unwrap_or(true);
13074        let matvec = |proj,
13075                      exi: &_,
13076                      exo: &_,
13077                      exp_d: &_,
13078                      pt: &_,
13079                      aq: &_,
13080                      ad: &_,
13081                      inf,
13082                      outf,
13083                      qtype,
13084                      rb|
13085         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13086            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
13087            let dec = dec && q8_expert_dec_supported(qtype);
13088            if dec {
13089                e.moe_pairs_matvec_q8_dec(
13090                    &dev.ptr_row,
13091                    proj,
13092                    exi,
13093                    exo,
13094                    exp_d,
13095                    pt,
13096                    aq,
13097                    ad,
13098                    inf,
13099                    outf,
13100                    n_expert,
13101                    n_active,
13102                    n_pairs,
13103                    qtype,
13104                    rb,
13105                )
13106            } else {
13107                e.moe_pairs_matvec_q8_em(
13108                    &dev.ptr_row,
13109                    proj,
13110                    exi,
13111                    exo,
13112                    exp_d,
13113                    pt,
13114                    aq,
13115                    ad,
13116                    inf,
13117                    outf,
13118                    n_expert,
13119                    n_active,
13120                    n_pairs,
13121                    qtype,
13122                    rb,
13123                )
13124            }
13125        };
13126        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13127        let gate = matvec(
13128            0,
13129            &exi,
13130            &exo,
13131            &exp_d,
13132            &pt,
13133            &zq,
13134            &zd,
13135            n_embd,
13136            n_ff_exp,
13137            m.gate_exps.qtype,
13138            rbg_d,
13139        )?;
13140        let up = matvec(
13141            1,
13142            &exi,
13143            &exo,
13144            &exp_d,
13145            &pt,
13146            &zq,
13147            &zd,
13148            n_embd,
13149            n_ff_exp,
13150            m.up_exps.qtype,
13151            rbu_d,
13152        )?;
13153        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
13154        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13155        // down consumes PAIR-major activation rows: pair_tok = identity.
13156        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
13157        let pself = e.htod_i32(&pair_self)?;
13158        let y_down = matvec(
13159            2,
13160            &exi,
13161            &exo,
13162            &exp_d,
13163            &pself,
13164            &aq2,
13165            &ad2,
13166            n_ff_exp,
13167            n_embd,
13168            m.down_exps.qtype,
13169            m.down_exps.row_bytes,
13170        )?;
13171        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
13172        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
13173
13174        // SHARED EXPERT epilogue — same as the other paths.
13175        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
13176        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
13177        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13178            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13179        {
13180            let n_ff_sh = gate_shexp.out_features();
13181            // These decode-exact forms are required by the new Step resident arm. Keep the
13182            // established grouped shared-expert program for every other architecture: widening
13183            // this to Gemma changed its speculative acceptance despite green argmax gates.
13184            let step_exact = true;
13185            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
13186            let (sg_gate, sg_up) = if step_exact && t == 1 {
13187                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, None)?
13188            } else if verify_t {
13189                let mut fused = None;
13190                if crate::spec::spec_fused_t()
13191                    && (2..=4).contains(&t)
13192                    && e.uses_q8_1_fast(gate_shexp)
13193                    && e.uses_q8_1_fast(up_shexp)
13194                {
13195                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13196                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13197                }
13198                match fused {
13199                    Some(pair) => pair,
13200                    None => (
13201                        e.matmul_decode_exact(gate_shexp, z, t)?,
13202                        e.matmul_decode_exact(up_shexp, z, t)?,
13203                    ),
13204                }
13205            } else {
13206                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13207            };
13208            let mut sa = e.uninit(t * n_ff_sh)?;
13209            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13210            let sh = if verify_t {
13211                e.matmul_decode_exact(down_shexp, &sa, t)?
13212            } else {
13213                e.matmul(down_shexp, &sa, t)?
13214            };
13215            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13216            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
13217            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
13218            // dispatch choice cannot change bits.
13219            let g = match &m.gate_inp_shexp {
13220                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
13221                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13222                }
13223                Some(gate_inp_shexp) => {
13224                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13225                    let mut g = e.uninit(t)?;
13226                    e.sigmoid(&gs, &mut g, t)?;
13227                    g
13228                }
13229                None => e.htod(&vec![1.0f32; t])?,
13230            };
13231            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13232        }
13233        Ok(moe_out)
13234    }
13235
13236    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
13237    #[allow(clippy::too_many_arguments)]
13238    #[allow(clippy::too_many_arguments)]
13239    fn moe_ffn_dev(
13240        e: &Engine,
13241        m: &MoeWeights,
13242        z: &CudaSlice<f32>,
13243        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
13244        logits: &CudaSlice<f32>,
13245        t: usize,
13246        cfg: &ModelConfig,
13247        il: u16,
13248        max_block: usize,
13249    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13250        let moe = cfg.moe.as_ref().unwrap();
13251        let n_embd = cfg.n_embd as usize;
13252        let n_expert = moe.expert_count as usize;
13253        let n_used = moe.expert_used_count as usize;
13254        let n_ff_exp = moe.expert_ff_length as usize;
13255        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
13256        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
13257        // clamped layers; assert both so a future caller that skips the gate fails loudly.
13258        debug_assert!(
13259            cfg.sigmoid_router().is_none(),
13260            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
13261        );
13262        debug_assert!(
13263            !cfg.swiglu_clamped_at(il as u32),
13264            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
13265        );
13266
13267        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
13268        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
13269        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
13270        // skipped entirely for macro-free experts (every k-quant GGUF).
13271        if m.has_macros {
13272            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
13273        }
13274
13275        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
13276        let mut moe_out = e.uninit(t * n_embd)?;
13277
13278        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
13279        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
13280        if let Some(dev) = m.dev_exps.as_ref() {
13281            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
13282            // the combined stride; up's base is offset in the ptr table. Down unchanged.
13283            let (rbg_d, rbu_d) = if dev.gu_il {
13284                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
13285                (sxx, sxx)
13286            } else {
13287                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
13288            };
13289            let q8 = moe_q8_enabled_for_model(cfg, m);
13290            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
13291            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
13292            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
13293            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
13294            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
13295            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
13296            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
13297            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
13298            let rows_arm = q8
13299                && t > 1
13300                && crate::spec::spec_m2()
13301                && n_ff_exp == 512
13302                && n_used <= 8
13303                && std::env::var("MEMRA_MOE_DEVQ8_GU")
13304                    .map(|v| v.is_empty() || v == "v")
13305                    .unwrap_or(true)
13306                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
13307                    .map(|v| v.is_empty() || v == "w8h2v")
13308                    .unwrap_or(true);
13309            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
13310            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
13311            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
13312            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
13313            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
13314            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
13315            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
13316            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
13317            let csr_mode = std::env::var("MEMRA_MOE_CSR")
13318                .ok()
13319                .and_then(|v| v.parse::<i32>().ok())
13320                .unwrap_or(1);
13321            // NVFP4 admission REVERTED 2026-08-21 (lane/samplat, decode-batch-gate2 find):
13322            // the csr_nvfp4 kernel drifts last-ULP vs the rows program (11041/32768 ACT
13323            // elements at t=8) and the drift is BATCH-COMPOSITION-DEPENDENT — gate2 (B=8 vs
13324            // isolated) FAILED on the ornith15 artifact, the one-numeric-program law's batch
13325            // axis. Three chain-pinning attempts did not close it (receipts,
13326            // research/samplat-20260821/); a source-verbatim per-pair helper form IS
13327            // bit-identical but loses the dedup win (-3% vs rows). NVFP4 stays on the rows
13328            // twins until a cached form passes gate2 + the =2 byte-compare at t=8. The
13329            // increment-1 qualification hole: =2 ran across run-spec (solo verify shapes),
13330            // never decode-batch-gate at B=8 on the MoE model itself.
13331            // MEMRA_MOE_CSR_NVFP4=1 (lane/orndecode, DIAGNOSTIC PROBE ONLY): re-admits NVFP4
13332            // to the CSR arm and widens it to the exact-16 decode widths, so gate2 B=12/16 +
13333            // the =2 byte-compare can re-adjudicate the cached form at the widths where the
13334            // serial dev loop hurts most (B=16 tick: 1280 launches/step). The v0.100.1
13335            // de-admission verdict above stands until those gates are GREEN on the MoE
13336            // artifact; this door must never default on.
13337            let csr_nvfp4_probe = std::env::var("MEMRA_MOE_CSR_NVFP4").as_deref() == Ok("1");
13338            let csr_qt = |qt: i32| {
13339                qt == crate::QT_IQ4_XS
13340                    || qt == crate::QT_IQ3_S
13341                    || (csr_nvfp4_probe && qt == crate::QT_NVFP4)
13342            };
13343            let csr_t_max = if csr_nvfp4_probe { MOE_DEV_MAX_T } else { 10 };
13344            let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
13345            let csr_arm = rows_arm
13346                && csr_mode > 0
13347                && t <= csr_t_max
13348                && csr_uniform
13349                && csr_qt(m.gate_exps.qtype)
13350                && csr_qt(m.up_exps.qtype)
13351                && csr_qt(m.down_exps.qtype);
13352            if csr_arm {
13353                if csr_mode == 2 {
13354                    static ENGAGED: std::sync::Once = std::sync::Once::new();
13355                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
13356                }
13357                let n_pairs = t * n_used;
13358                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13359                let act = e.moe_gate_up_silu8_dev_q8_csr(
13360                    &dev.ptr_row,
13361                    &sel_d,
13362                    &zq,
13363                    &zd,
13364                    n_pairs,
13365                    n_embd,
13366                    n_ff_exp,
13367                    n_used,
13368                    n_expert,
13369                    m.gate_exps.qtype,
13370                    m.up_exps.qtype,
13371                    rbg_d,
13372                    rbu_d,
13373                )?;
13374                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
13375                // down stays on the _rows twin — BOTH CSR down variants measured negative
13376                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
13377                // 16-group rows have too little decode to amortize any dedup structure.
13378                e.moe_down8_fma_dev_q8_rows(
13379                    &dev.ptr_row,
13380                    &sel_d,
13381                    &w_d,
13382                    &aq2,
13383                    &ad2,
13384                    &mut moe_out,
13385                    t,
13386                    n_ff_exp,
13387                    n_embd,
13388                    n_used,
13389                    n_expert,
13390                    m.down_exps.qtype,
13391                    m.down_exps.row_bytes,
13392                )?;
13393                if csr_mode == 2 {
13394                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
13395                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
13396                        &dev.ptr_row,
13397                        &sel_d,
13398                        &zq,
13399                        &zd,
13400                        t,
13401                        n_embd,
13402                        n_ff_exp,
13403                        n_used,
13404                        n_expert,
13405                        m.gate_exps.qtype,
13406                        m.up_exps.qtype,
13407                        rbg_d,
13408                        rbu_d,
13409                        &m.dev_macros,
13410                    )?;
13411                    let mut out_r = e.uninit(t * n_embd)?;
13412                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
13413                    e.moe_down8_fma_dev_q8_rows(
13414                        &dev.ptr_row,
13415                        &sel_d,
13416                        &w_d,
13417                        &aq2r,
13418                        &ad2r,
13419                        &mut out_r,
13420                        t,
13421                        n_ff_exp,
13422                        n_embd,
13423                        n_used,
13424                        n_expert,
13425                        m.down_exps.qtype,
13426                        m.down_exps.row_bytes,
13427                    )?;
13428                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
13429                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
13430                    let ba = a1
13431                        .iter()
13432                        .zip(&a2)
13433                        .filter(|(x, y)| x.to_bits() != y.to_bits())
13434                        .count();
13435                    let bo = o1
13436                        .iter()
13437                        .zip(&o2)
13438                        .filter(|(x, y)| x.to_bits() != y.to_bits())
13439                        .count();
13440                    if ba + bo > 0 {
13441                        eprintln!(
13442                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
13443                            a1.len(),
13444                            o1.len()
13445                        );
13446                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
13447                        let sel_h = e.dtoh_i32(&sel_d)?;
13448                        let mut shown = 0;
13449                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
13450                            if x.to_bits() != y.to_bits() && shown < 4 {
13451                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
13452                                let ex = sel_h[p];
13453                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
13454                                eprintln!(
13455                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
13456                                );
13457                                shown += 1;
13458                            }
13459                        }
13460                        std::process::exit(3);
13461                    }
13462                }
13463            } else if rows_arm {
13464                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
13465                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
13466                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
13467                    use std::sync::atomic::{AtomicU64, Ordering};
13468                    static PAIRS: AtomicU64 = AtomicU64::new(0);
13469                    static UNIQ: AtomicU64 = AtomicU64::new(0);
13470                    static CALLS: AtomicU64 = AtomicU64::new(0);
13471                    let sel_h = e.dtoh_i32(&sel_d)?;
13472                    let mut u: Vec<i32> = sel_h.clone();
13473                    u.sort_unstable();
13474                    u.dedup();
13475                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
13476                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
13477                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
13478                    if c.is_multiple_of(480) {
13479                        let p = PAIRS.load(Ordering::Relaxed);
13480                        let q = UNIQ.load(Ordering::Relaxed);
13481                        eprintln!(
13482                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
13483                            q as f64 / p as f64
13484                        );
13485                    }
13486                }
13487                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13488                let act = e.moe_gate_up_silu8_dev_q8_rows(
13489                    &dev.ptr_row,
13490                    &sel_d,
13491                    &zq,
13492                    &zd,
13493                    t,
13494                    n_embd,
13495                    n_ff_exp,
13496                    n_used,
13497                    n_expert,
13498                    m.gate_exps.qtype,
13499                    m.up_exps.qtype,
13500                    rbg_d,
13501                    rbu_d,
13502                    &m.dev_macros,
13503                )?;
13504                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
13505                e.moe_down8_fma_dev_q8_rows(
13506                    &dev.ptr_row,
13507                    &sel_d,
13508                    &w_d,
13509                    &aq2,
13510                    &ad2,
13511                    &mut moe_out,
13512                    t,
13513                    n_ff_exp,
13514                    n_embd,
13515                    n_used,
13516                    n_expert,
13517                    m.down_exps.qtype,
13518                    m.down_exps.row_bytes,
13519                )?;
13520            } else {
13521                for tok in 0..t {
13522                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13523                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13524                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13525                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13526                    if q8 {
13527                        let (zq, zd) = match (t, zq8) {
13528                            (1, Some((q, d))) => (q.clone(), d.clone()),
13529                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
13530                        };
13531                        let act = e.moe_gate_up_silu8_dev_q8(
13532                            &dev.ptr_row,
13533                            &selt,
13534                            &zq,
13535                            &zd,
13536                            n_embd,
13537                            n_ff_exp,
13538                            n_used,
13539                            n_expert,
13540                            m.gate_exps.qtype,
13541                            m.up_exps.qtype,
13542                            rbg_d,
13543                            rbu_d,
13544                            &m.dev_macros,
13545                        )?;
13546                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13547                        e.moe_down8_fma_dev_q8(
13548                            &dev.ptr_row,
13549                            &selt,
13550                            &wt,
13551                            &aq2,
13552                            &ad2,
13553                            &mut dst,
13554                            n_ff_exp,
13555                            n_embd,
13556                            n_used,
13557                            n_expert,
13558                            m.down_exps.qtype,
13559                            m.down_exps.row_bytes,
13560                        )?;
13561                    } else {
13562                        let act = e.moe_gate_up_silu8_dev(
13563                            &dev.ptr_row,
13564                            &selt,
13565                            &zt,
13566                            n_embd,
13567                            n_ff_exp,
13568                            n_used,
13569                            n_expert,
13570                            m.gate_exps.qtype,
13571                            m.up_exps.qtype,
13572                            rbg_d,
13573                            rbu_d,
13574                            &m.dev_macros,
13575                        )?;
13576                        e.moe_down8_fma_dev(
13577                            &dev.ptr_row,
13578                            &selt,
13579                            &wt,
13580                            &act,
13581                            &mut dst,
13582                            n_ff_exp,
13583                            n_embd,
13584                            n_used,
13585                            n_expert,
13586                            m.down_exps.qtype,
13587                            m.down_exps.row_bytes,
13588                        )?;
13589                    }
13590                }
13591            }
13592        } else {
13593            // Launch under the cache lock: the row borrow lives as long as the closure, and the
13594            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
13595            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
13596            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
13597            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
13598            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
13599            let q8 = moe_q8_enabled_for_model(cfg, m);
13600            e.with_moe_cache(max_block, |c, eng| {
13601                let row = c
13602                    .layer_dev_row(il, n_expert, eng)?
13603                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
13604                for tok in 0..t {
13605                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
13606                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
13607                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
13608                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13609                    if q8 {
13610                        let (zq, zd) = match (t, zq8) {
13611                            (1, Some((q, d))) => (q.clone(), d.clone()),
13612                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
13613                        };
13614                        let act = eng.moe_gate_up_silu8_dev_q8(
13615                            row,
13616                            &selt,
13617                            &zq,
13618                            &zd,
13619                            n_embd,
13620                            n_ff_exp,
13621                            n_used,
13622                            n_expert,
13623                            m.gate_exps.qtype,
13624                            m.up_exps.qtype,
13625                            m.gate_exps.row_bytes,
13626                            m.up_exps.row_bytes,
13627                            &m.dev_macros,
13628                        )?;
13629                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
13630                        eng.moe_down8_fma_dev_q8(
13631                            row,
13632                            &selt,
13633                            &wt,
13634                            &aq2,
13635                            &ad2,
13636                            &mut dst,
13637                            n_ff_exp,
13638                            n_embd,
13639                            n_used,
13640                            n_expert,
13641                            m.down_exps.qtype,
13642                            m.down_exps.row_bytes,
13643                        )?;
13644                    } else {
13645                        let act = eng.moe_gate_up_silu8_dev(
13646                            row,
13647                            &selt,
13648                            &zt,
13649                            n_embd,
13650                            n_ff_exp,
13651                            n_used,
13652                            n_expert,
13653                            m.gate_exps.qtype,
13654                            m.up_exps.qtype,
13655                            m.gate_exps.row_bytes,
13656                            m.up_exps.row_bytes,
13657                            &m.dev_macros,
13658                        )?;
13659                        eng.moe_down8_fma_dev(
13660                            row,
13661                            &selt,
13662                            &wt,
13663                            &act,
13664                            &mut dst,
13665                            n_ff_exp,
13666                            n_embd,
13667                            n_used,
13668                            n_expert,
13669                            m.down_exps.qtype,
13670                            m.down_exps.row_bytes,
13671                        )?;
13672                    }
13673                }
13674                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
13675                c.hits += (t * 3 * n_used) as u64;
13676                Ok(())
13677            })?;
13678        }
13679
13680        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
13681        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
13682        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
13683        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
13684        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
13685            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
13686        {
13687            let n_ff_sh = gate_shexp.out_features();
13688            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
13689            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
13690            let verify_t = t > 1 && t < PRIME_MIN_T;
13691            let (sg_gate, sg_up) = if t == 1 {
13692                shexp_gate_up_t1(e, gate_shexp, up_shexp, z, zq8)?
13693            } else if verify_t {
13694                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
13695                // rides one shared quantize + one fused2 batched launch instead of two
13696                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
13697                let mut fused = None;
13698                if crate::spec::spec_fused_t()
13699                    && (2..=4).contains(&t)
13700                    && e.uses_q8_1_fast(gate_shexp)
13701                    && e.uses_q8_1_fast(up_shexp)
13702                {
13703                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
13704                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
13705                }
13706                match fused {
13707                    Some(pair) => pair,
13708                    None => (
13709                        e.matmul_decode_exact(gate_shexp, z, t)?,
13710                        e.matmul_decode_exact(up_shexp, z, t)?,
13711                    ),
13712                }
13713            } else {
13714                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
13715            };
13716            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
13717            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
13718            let sh = if verify_t {
13719                e.matmul_decode_exact(down_shexp, &sa, t)?
13720            } else {
13721                e.matmul(down_shexp, &sa, t)?
13722            };
13723            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
13724            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
13725            // between the two arms; prefill keeps the batched cuBLASLt linear).
13726            let g = match &m.gate_inp_shexp {
13727                Some(gate_inp_shexp) => {
13728                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
13729                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
13730                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
13731                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
13732                    } else {
13733                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
13734                        let mut g = e.uninit(t)?;
13735                        e.sigmoid(&gs, &mut g, t)?;
13736                        g
13737                    }
13738                }
13739                None => e.htod(&vec![1.0f32; t])?,
13740            };
13741            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
13742        }
13743
13744        Ok(moe_out)
13745    }
13746
13747    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
13748    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
13749    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
13750    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
13751    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
13752    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
13753    /// the collected raw pointers cannot move between collection and launch (single-threaded
13754    /// decode; the lock is held only for collection, launches are stream-ordered after any
13755    /// prior same-stream staging writes).
13756    #[allow(clippy::too_many_arguments)]
13757    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
13758    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
13759    #[allow(clippy::too_many_arguments)]
13760    fn moe_gdec_token_q8(
13761        e: &Engine,
13762        m: &MoeWeights,
13763        il: u16,
13764        max_block: usize,
13765        zq: &CudaSlice<i8>,
13766        zd: &CudaSlice<f32>,
13767        sel: &[u32],
13768        w: &[f32],
13769        moe_out: &mut CudaSlice<f32>,
13770        tok: usize,
13771        n_embd: usize,
13772        n_ff_exp: usize,
13773        n_used: usize,
13774    ) -> Result<bool, Box<dyn std::error::Error>> {
13775        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
13776        use cudarc::driver::DevicePtr;
13777        let ptrs = e.with_moe_cache(max_block, |c, eng| {
13778            let mut g = [0u64; 8];
13779            let mut u = [0u64; 8];
13780            let mut d = [0u64; 8];
13781            for (j, &ex) in sel.iter().enumerate() {
13782                let ex = ex as u16;
13783                let (Some(sg), Some(su), Some(sd)) = (
13784                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
13785                    c.resident(BlockId::new(il, PROJ_UP, ex)),
13786                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
13787                ) else {
13788                    return Ok(None);
13789                };
13790                let __s = eng.stream();
13791                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
13792                let (pu, _e1) = c.slot(su).device_ptr(&__s);
13793                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
13794                g[j] = pg;
13795                u[j] = pu;
13796                d[j] = pd;
13797            }
13798            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
13799                for &ex in sel {
13800                    let ex = ex as u16;
13801                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
13802                        c.note_profile_hit(BlockId::new(il, proj, ex));
13803                    }
13804                }
13805            }
13806            c.hits += (3 * n_used) as u64;
13807            Ok(Some((g, u, d)))
13808        })?;
13809        let Some((g, u, d)) = ptrs else {
13810            return Ok(false);
13811        };
13812        let mut wv = [0f32; 8];
13813        wv[..n_used].copy_from_slice(w);
13814        let act = e.moe_gate_up_silu8_q8(
13815            crate::WPtr8(g),
13816            crate::WPtr8(u),
13817            zq,
13818            zd,
13819            n_embd,
13820            n_ff_exp,
13821            n_used,
13822            m.gate_exps.qtype,
13823            m.up_exps.qtype,
13824            m.gate_exps.row_bytes,
13825            m.up_exps.row_bytes,
13826        )?;
13827        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
13828        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
13829        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
13830        e.moe_down8_fma_q8(
13831            crate::WPtr8(d),
13832            crate::F32x8(wv),
13833            &aq2,
13834            &ad2,
13835            &mut dst,
13836            n_ff_exp,
13837            n_embd,
13838            n_used,
13839            m.down_exps.qtype,
13840            m.down_exps.row_bytes,
13841        )?;
13842        Ok(true)
13843    }
13844
13845    /// glm5_next's fused MoE epilogue for ONE token-layer: the sigmoid router's already-selected
13846    /// `(sel, w)`, the PRE-clamped SwiGLU and the per-expert NVFP4 macro fold, in one launch
13847    /// pair. Returns `false` when the cache cannot hold this token's `3*n_used` blocks at once,
13848    /// in which case the caller must run the sequential loop (which zeroes its own row).
13849    ///
13850    /// WHY IT DOES NOT NEED A RESIDENT LAYER, unlike `moe_gdec_token_q8`. gdec collects pointers
13851    /// from blocks that are ALREADY resident and bails on the first miss, because a miss would
13852    /// mean an admission that could move a slot under the pointers it has already taken. This arm
13853    /// inverts the order: it ADMITS all `3*n_used` blocks first, through the same
13854    /// `dispatch_source` the sequential loop calls per projection (a hit copies nothing, a miss
13855    /// runs the identical `memcpy_htod` into a slot), and only then takes the addresses, in a
13856    /// second pass, with the cache lock still held. Nothing can move between the last admission
13857    /// and the pointer read, and the kernels are issued on the compute stream immediately after —
13858    /// the same in-order guarantee the sequential loop already relies on when it dispatches
13859    /// expert j+1 after launching expert j's kernels.
13860    ///
13861    /// The slot-capacity check is the fail-closed seam: `n_slots()` is a whole-cache bound, and
13862    /// with fewer than `3*n_used` slots an admission is guaranteed to evict one of this token's
13863    /// own blocks. The second pass re-reads `resident()` for every block rather than trusting the
13864    /// dispatch's return, so an eviction the capacity check did not predict falls through loudly
13865    /// to the sequential loop instead of dereferencing a reused slot.
13866    ///
13867    /// BIT-IDENTITY CLASS. Against the sequential loop this arm is a DISPATCH-class change, not a
13868    /// provenance one: the same block bytes and the same macro scales, but the gate/up dots are
13869    /// the fused kernel's warp reduction rather than `qmatvec_expert_q8`'s per-projection one, and
13870    /// the down accumulation is `moe_down8_fma_q8`'s slot-ordered `__fmaf_rn` chain rather than 8
13871    /// separate `axpy_into` calls. Those chains are the ones the gdec family documents as
13872    /// reproducing the sequential chain exactly; `tests/glm5_moe_epilogue_gpu.rs::the_two_arms_agree`
13873    /// measures the actual bit disagreement rather than asserting the claim.
13874    #[allow(clippy::too_many_arguments)]
13875    fn moe_fused_epi_token_q8(
13876        e: &Engine,
13877        m: &MoeWeights,
13878        il: u16,
13879        max_block: usize,
13880        zq: &CudaSlice<i8>,
13881        zd: &CudaSlice<f32>,
13882        sel: &[u32],
13883        w: &[f32],
13884        moe_out: &mut CudaSlice<f32>,
13885        tok: usize,
13886        n_embd: usize,
13887        n_ff_exp: usize,
13888        n_used: usize,
13889        limit: f32,
13890    ) -> Result<bool, Box<dyn std::error::Error>> {
13891        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_DOWN, PROJ_GATE, PROJ_UP};
13892        use cudarc::driver::DevicePtr;
13893        debug_assert!(
13894            limit > 1e-6,
13895            "the fused epilogue's kernel collapses every gate to silu(0) at limit 0"
13896        );
13897        debug_assert_eq!(sel.len(), n_used);
13898        debug_assert_eq!(w.len(), n_used);
13899
13900        let ptrs = e.with_moe_cache(max_block, |c, eng| {
13901            // Fail closed: below this bound an admission MUST evict one of this token's own
13902            // blocks, so there is no pointer set that stays valid.
13903            if c.n_slots() < 3 * n_used {
13904                return Ok(None);
13905            }
13906            // PASS 1 — admit. Identical dispatch to the sequential loop's `moe_cached_gemm_q8`,
13907            // projection for projection; only the GEMM is deferred.
13908            for &ex in sel.iter() {
13909                let ex_usize = ex as usize;
13910                for (proj, exps) in [
13911                    (PROJ_GATE, &m.gate_exps),
13912                    (PROJ_UP, &m.up_exps),
13913                    (PROJ_DOWN, &m.down_exps),
13914                ] {
13915                    let id = BlockId::new(il, proj, ex as u16);
13916                    let DispatchSlot::Resident(_) =
13917                        c.dispatch_source(id, exps.expert_source(ex_usize), eng)?;
13918                }
13919            }
13920            // PASS 2 — take the fixed slot addresses, with nothing left to admit.
13921            let mut g = [0u64; 8];
13922            let mut u = [0u64; 8];
13923            let mut d = [0u64; 8];
13924            for (j, &ex) in sel.iter().enumerate() {
13925                let ex = ex as u16;
13926                let (Some(sg), Some(su), Some(sd)) = (
13927                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
13928                    c.resident(BlockId::new(il, PROJ_UP, ex)),
13929                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
13930                ) else {
13931                    return Ok(None);
13932                };
13933                let __s = eng.stream();
13934                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
13935                let (pu, _e1) = c.slot(su).device_ptr(&__s);
13936                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
13937                g[j] = pg;
13938                u[j] = pu;
13939                d[j] = pd;
13940            }
13941            Ok(Some((g, u, d)))
13942        })?;
13943        let Some((g, u, d)) = ptrs else {
13944            return Ok(false);
13945        };
13946
13947        Self::moe_fused_epi_launch(
13948            e, m, zq, zd, sel, w, g, u, d, moe_out, tok, n_embd, n_ff_exp, n_used, limit,
13949        )?;
13950        Ok(true)
13951    }
13952
13953    /// The fused epilogue's ONLY launch path, shared by both provenances (SLRU slot addresses and
13954    /// device-resident slab base+stride). Everything that could differ semantically between them
13955    /// — the per-expert macro fold, the clamp, the kernel pair, the dispatch counter — lives here
13956    /// exactly once, so the two arms cannot drift into being different programs. The callers
13957    /// differ only in how they filled `g`/`u`/`d`.
13958    ///
13959    /// Per-expert macro scales in router slot order: gate/up ride the kernel's epilogue exactly
13960    /// where `ffn_act_lim`'s gs/us ride the unfused loop, and down folds into the routing weight
13961    /// exactly where `axpy_into`'s `w[j] * macro_scale(ex)` folds it. `macro_scale` answers 1.0
13962    /// for a macro-free bank, so a k-quant GGUF takes this path with no fold and no branch.
13963    #[allow(clippy::too_many_arguments)]
13964    fn moe_fused_epi_launch(
13965        e: &Engine,
13966        m: &MoeWeights,
13967        zq: &CudaSlice<i8>,
13968        zd: &CudaSlice<f32>,
13969        sel: &[u32],
13970        w: &[f32],
13971        g: [u64; 8],
13972        u: [u64; 8],
13973        d: [u64; 8],
13974        moe_out: &mut CudaSlice<f32>,
13975        tok: usize,
13976        n_embd: usize,
13977        n_ff_exp: usize,
13978        n_used: usize,
13979        limit: f32,
13980    ) -> Result<(), Box<dyn std::error::Error>> {
13981        let mut gs = [0f32; 8];
13982        let mut us = [0f32; 8];
13983        let mut wv = [0f32; 8];
13984        for (j, &ex) in sel.iter().enumerate() {
13985            let ex = ex as usize;
13986            gs[j] = m.gate_exps.macro_scale(ex);
13987            us[j] = m.up_exps.macro_scale(ex);
13988            wv[j] = w[j] * m.down_exps.macro_scale(ex);
13989        }
13990        let act = e.moe_gate_up_preclamp8_q8(
13991            crate::WPtr8(g),
13992            crate::WPtr8(u),
13993            zq,
13994            zd,
13995            crate::F32x8(gs),
13996            crate::F32x8(us),
13997            limit,
13998            n_embd,
13999            n_ff_exp,
14000            n_used,
14001            m.gate_exps.qtype,
14002            m.up_exps.qtype,
14003            m.gate_exps.row_bytes,
14004            m.up_exps.row_bytes,
14005        )?;
14006        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
14007        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
14008        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14009        e.moe_down8_fma_q8(
14010            crate::WPtr8(d),
14011            crate::F32x8(wv),
14012            &aq2,
14013            &ad2,
14014            &mut dst,
14015            n_ff_exp,
14016            n_embd,
14017            n_used,
14018            m.down_exps.qtype,
14019            m.down_exps.row_bytes,
14020        )?;
14021        crate::MOE_FUSED_EPI_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14022        Ok(())
14023    }
14024
14025    /// The verify-rows batched routed-expert program (lane/glm5-vrest): one layer-call's
14026    /// WHOLE t x n_used pair union through the fused-epilogue kernels' rows twins —
14027    /// one gate/up+preclamp launch, one pair-major activation quantize, one down+FMA
14028    /// launch — with pointers computed from the resident slab base + ex*stride (the
14029    /// sequential slab arm's exact arithmetic) and the macro folds landing exactly where
14030    /// `ffn_act_lim` / `axpy_into` land them. `moe_out` rows are FULLY overwritten.
14031    #[allow(clippy::too_many_arguments)]
14032    // allow: the parameter list mirrors its dispatch-arm caller's contract
14033    fn moe_vrows_pairs_q8(
14034        e: &Engine,
14035        m: &MoeWeights,
14036        z: &CudaSlice<f32>,
14037        sel: VrowsSel<'_>,
14038        il: u16,
14039        (pg, pu, pd): (u64, u64, u64),
14040        t: usize,
14041        n_embd: usize,
14042        n_ff_exp: usize,
14043        n_used: usize,
14044        limit: f32,
14045        moe_out: &mut CudaSlice<f32>,
14046    ) -> Result<(), Box<dyn std::error::Error>> {
14047        let n_pairs = t * n_used;
14048        // Door E (MEMRA_MOE_VROWS_DEDUP_ORDER, default OFF): the gate/up launch walks the pair
14049        // union EXPERT-MAJOR, reading the visit order from a FOURTH plane appended to the pointer
14050        // table (planes: gate | up | down | order). Carrying it in the existing table is what
14051        // makes the door free on the host arm — the order plane rides the single `htod_u64_into`
14052        // that was already uploading the pointers, so no new transfer and no new pool appear.
14053        // Door M (`MEMRA_MOE_VROWS_PACK`) refuses it in the launcher, so do not build the plane
14054        // when the refuted pack door is armed.
14055        let order_on = crate::moe_vrows_dedup_order_on() && !crate::moe_vrows_pack_on();
14056        let n_planes = if order_on { 4 } else { 3 };
14057        // Door W (MEMRA_GLM5_VERIFY_WS): the whole staging set — tables, token quantize,
14058        // act, pair quantize — draws from the verify workspace and recycles at the end of
14059        // the call (vws_* are alloc_uninit/plain-drop with the door off, so the OFF arm is
14060        // byte-for-byte the shipped program). Every buffer is fully overwritten before any
14061        // read by the SAME kernels (the sites' standing uninit contract).
14062        let mut ptrs_d = e.vws_uninit_u64(n_planes * n_pairs)?;
14063        let mut scl_d = e.vws_uninit(3 * n_pairs)?;
14064        // ONE launch path, TWO table provenances (the fused-epilogue arm's own discipline):
14065        // only the plane-major (gate | up | down) pointer/scale tables are built differently,
14066        // and door D's kernel evaluates the SAME terms as the host loop, so nothing downstream
14067        // can tell the arms apart. See the `moe_vrows_tables_from_sel` kernel comment for the
14068        // term-by-term bit-identity argument.
14069        match sel {
14070            VrowsSel::Host(sel_all, w_all) => {
14071                debug_assert_eq!(sel_all.len(), n_pairs);
14072                debug_assert_eq!(w_all.len(), n_pairs);
14073                let mut ptrs = vec![0u64; n_planes * n_pairs];
14074                let mut scl = vec![0f32; 3 * n_pairs];
14075                for (p, (&ex, &w)) in sel_all.iter().zip(w_all).enumerate() {
14076                    let ex = ex as usize;
14077                    ptrs[p] = pg + (ex * m.gate_exps.expert_stride) as u64;
14078                    ptrs[n_pairs + p] = pu + (ex * m.up_exps.expert_stride) as u64;
14079                    ptrs[2 * n_pairs + p] = pd + (ex * m.down_exps.expert_stride) as u64;
14080                    scl[p] = m.gate_exps.macro_scale(ex);
14081                    scl[n_pairs + p] = m.up_exps.macro_scale(ex);
14082                    // down-proj macro folds into the accumulate weight (1.0 for non-macro
14083                    // banks) — the axpy_into fold, verbatim.
14084                    scl[2 * n_pairs + p] = w * m.down_exps.macro_scale(ex);
14085                }
14086                if order_on {
14087                    // The order plane rides the SAME upload — the door adds no HtoD on this arm.
14088                    ptrs[3 * n_pairs..].copy_from_slice(&crate::vrows_expert_major_order(sel_all));
14089                    // The box receipt: the slab reads whose repeat visit this schedule places
14090                    // inside the reuse window (host arm only — see MOE_VROWS_SLAB_READS_AVOIDED).
14091                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14092                    crate::MOE_VROWS_SLAB_READS_AVOIDED
14093                        .fetch_add(visits - distinct, std::sync::atomic::Ordering::Relaxed);
14094                }
14095                e.htod_u64_into(&ptrs, &mut ptrs_d)?;
14096                e.htod_f32_into(&scl, &mut scl_d)?;
14097                // MEMRA_MOE_VROWS_DEDUP_STAT: size the ONLY remaining byte lever on this pair
14098                // (LANE.md §1 — it already runs at ~90% of theoretical DRAM peak, so the sole
14099                // way to cut it further is reading a shared expert slab once for the rows that
14100                // share it). `1 - distinct/visits` IS that lever; measuring it costs a bitset.
14101                if crate::moe_vrows_dedup_stat_on() {
14102                    let (visits, distinct) = crate::vrows_overlap_counts(sel_all);
14103                    debug_assert_eq!(visits, n_pairs as u64);
14104                    crate::MOE_VROWS_PAIR_VISITS
14105                        .fetch_add(visits, std::sync::atomic::Ordering::Relaxed);
14106                    crate::MOE_VROWS_PAIR_DISTINCT
14107                        .fetch_add(distinct, std::sync::atomic::Ordering::Relaxed);
14108                    crate::moe_vrows_dedup_report();
14109                }
14110            }
14111            VrowsSel::Dev(sel_d, selw_d) => {
14112                let macros = match (
14113                    m.gate_exps.macros.as_deref(),
14114                    m.up_exps.macros.as_deref(),
14115                    m.down_exps.macros.as_deref(),
14116                ) {
14117                    (Some(g), Some(u), Some(d)) => Some((g, u, d)),
14118                    // A partially-macro bank would need per-plane 1.0 defaults the kernel does
14119                    // not carry; the serving artifact's three planes are all present or all
14120                    // absent, so refuse rather than guess.
14121                    (None, None, None) => None,
14122                    _ => {
14123                        return Err("vrows device tables: expert macro planes are not uniform \
14124                                    across gate/up/down"
14125                            .into());
14126                    }
14127                };
14128                e.moe_vrows_tables_from_sel(
14129                    sel_d,
14130                    selw_d,
14131                    il,
14132                    macros,
14133                    (pg, pu, pd),
14134                    (
14135                        m.gate_exps.expert_stride,
14136                        m.up_exps.expert_stride,
14137                        m.down_exps.expert_stride,
14138                    ),
14139                    n_pairs,
14140                    &mut ptrs_d,
14141                    &mut scl_d,
14142                )?;
14143                if order_on {
14144                    // Door E on the device arm: one extra launch (the host arm gets the plane for
14145                    // free inside its existing upload). Bit-identical to the host's stable sort by
14146                    // (expert id, pair index) — gated directly against it in glm5_dedup_sched_gpu.
14147                    e.moe_vrows_order_from_sel(sel_d, n_pairs, &mut ptrs_d)?;
14148                }
14149                if crate::MOE_VROWS_DEV_TABLES_DISPATCHES
14150                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
14151                    == 0
14152                {
14153                    eprintln!(
14154                        "[moe-vrows-dev-tables] engaged: pointer/scale tables built on device \
14155                         from the router's own sel/w; the per-layer pinned readback and its \
14156                         cuStreamSynchronize are skipped (MEMRA_MOE_VROWS_DEV_TABLES=1)"
14157                    );
14158                }
14159            }
14160        }
14161        // Token rows quantized in one launch; per-row q8_1 bytes are position-independent
14162        // (the batched-MMVQ class), bit-gated against the per-token quantize_q8_1_view.
14163        let (mut zq, mut zd) = (
14164            e.vws_uninit_i8(t * n_embd)?,
14165            e.vws_uninit(t * (n_embd / 32))?,
14166        );
14167        e.quantize_q8_1_into(z, t, n_embd, &mut zq, &mut zd)?;
14168        let act = e.moe_gate_up_preclamp8_q8_rows(
14169            &ptrs_d,
14170            &scl_d,
14171            &zq,
14172            &zd,
14173            limit,
14174            n_embd,
14175            n_ff_exp,
14176            n_used,
14177            n_pairs,
14178            m.gate_exps.qtype,
14179            m.up_exps.qtype,
14180            m.gate_exps.row_bytes,
14181            m.up_exps.row_bytes,
14182        )?;
14183        // Pair-major activation quantize: [n_pairs, n_ff] rows in one launch.
14184        let (mut aq2, mut ad2) = (
14185            e.vws_uninit_i8(n_pairs * n_ff_exp)?,
14186            e.vws_uninit(n_pairs * (n_ff_exp / 32))?,
14187        );
14188        e.quantize_q8_1_into(&act, n_pairs, n_ff_exp, &mut aq2, &mut ad2)?;
14189        e.moe_down8_fma_q8_rows(
14190            &ptrs_d,
14191            &scl_d,
14192            &aq2,
14193            &ad2,
14194            moe_out,
14195            n_ff_exp,
14196            n_embd,
14197            n_used,
14198            n_pairs,
14199            m.down_exps.qtype,
14200            m.down_exps.row_bytes,
14201        )?;
14202        // Everything above is dead after the down launch (stream-ordered reuse is safe on
14203        // this engine's stream, the same guarantee the async free relies on).
14204        e.vws_recycle_u64(ptrs_d);
14205        e.vws_recycle(scl_d);
14206        e.vws_recycle_i8(zq);
14207        e.vws_recycle(zd);
14208        e.vws_recycle(act);
14209        e.vws_recycle_i8(aq2);
14210        e.vws_recycle(ad2);
14211        if crate::MOE_VROWS_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) == 0 {
14212            eprintln!(
14213                "[glm5-vrows] verify MoE batched across rows: pairs={n_pairs} (t={t} x \
14214                 {n_used}), one gate/up+preclamp launch + one down/FMA launch per layer-call \
14215                 (rides MEMRA_GLM5_VERIFY_BATCH)"
14216            );
14217        }
14218        Ok(())
14219    }
14220
14221    /// GROUPED MoE PREFILL for the sigmoid-router glm5_next class (`MEMRA_MOE_GROUPED_PREFILL`,
14222    /// default ON since 2026-08-29, `=0` rollback). One call covers a whole prefill chunk's
14223    /// routed-expert FFN for one layer:
14224    ///
14225    ///   1. ROUTER: the SAME m-invariant `moe_router_logits` + `moe_route_sigmoid_cfg` host
14226    ///      oracle invocation the sequential arm makes, so selected experts and routing weights
14227    ///      are BIT-identical to the sequential arm by construction. Only the GEMM accumulation
14228    ///      order may move (the grouped GEMM is measured non-bit-stable,
14229    ///      `run_tensor_parallel_routes_nvfp4_prime_grouped`'s MEMRA_MOE_DETERM note), which is
14230    ///      why the acceptance gate is reference-band + routing-exactness, not byte identity.
14231    ///   2. TOKEN-SORT BY EXPERT: host counting sort of the (token, expert) pairs into an
14232    ///      expert-major CSR (vLLM's `moe_align_block_size` shape; same O(pairs) build the
14233    ///      softmax `moe_ffn_pairs` arm and the step37 grouped prime use).
14234    ///   3. ONE GROUPED TENSOR-CORE GEMM PER PROJECTION over the resident NVFP4 slab
14235    ///      (`moe_f16_grouped`, the sk single-kernel visitor with the NVFP4 direct tile
14236    ///      loaders; the step37 grouped-prime kernel class, 170-270 TFLOP/s on its lane's
14237    ///      sizing rows, generalized off the TP runtime to the single-device `dev_exps`
14238    ///      pointer-table provenance). Each expert's weights stream through tensor cores ONCE
14239    ///      per layer per chunk instead of once per (token, expert): at t=4096 that replaces
14240    ///      the sequential loop's 49 launches x 4096 tokens (~200k launches and ~113 MB x 4096
14241    ///      of expert VRAM re-reads per layer) with a ~15-launch chunk-wide program.
14242    ///   4. EPILOGUE: glm5_next's PRE-clamped SwiGLU `silu(min(g,l)) * clamp(u,±l)` with the
14243    ///      per-expert `weight_scale_2` macro fold: gate/up macros land BEFORE the nonlinearity
14244    ///      (`scale_rows` per CSR row; silu is nonlinear, so the fold cannot commute past it),
14245    ///      down macros fold into the scatter weight, exactly where the fused epilogue and the
14246    ///      sequential loop's `ffn_act_lim`/`axpy_into` put them.
14247    ///   5. SCATTER: permute CSR rows back to pair order, then the slot-ordered weighted
14248    ///      per-token accumulation (`moe_pairs_scatter`, the sequential-axpy accumulation
14249    ///      class). Shared expert rides the canonical clamp-aware grouped add.
14250    ///
14251    /// Returns `Ok(None)` (fail closed to the sequential arm) for every unqualified shape:
14252    /// no local resident slab, f16g door off, a projection the grouped kernel cannot walk, or
14253    /// an expert count past the sk visitor's group cap. Numeric class: f16-mirror activations
14254    /// (`moe_f16g_act` row-normalized f16), the class the softmax pairs f16g arm and the
14255    /// step37 grouped prime already serve prefill with; gated by
14256    /// `tests/glm5_moe_grouped_prefill_gpu.rs` against `memra_reference` at the fused-epilogue
14257    /// gate's tolerance class, plus the run-gen first-token argmax gate on real prompts.
14258    fn moe_ffn_grouped_prefill_sigmoid(
14259        e: &Engine,
14260        m: &MoeWeights,
14261        z: &CudaSlice<f32>,
14262        t: usize,
14263        cfg: &ModelConfig,
14264        il: u16,
14265    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14266        // Placement: the LOCAL resident slab only. The SLRU cannot serve a 4096-token chunk's
14267        // expert working set (a glm5 layer is 288 x 3 blocks against ~285 slots/layer on the
14268        // serving recipe), and a remote slab must never be dereferenced (m=1 peer reads are the
14269        // measured 34-150x class). Fail closed: the sequential arm stages as before.
14270        let Some(dev) = m
14271            .dev_exps
14272            .as_ref()
14273            .filter(|d| moe_slab_enabled() && d.dev == e.ctx().ordinal())
14274        else {
14275            return Ok(None);
14276        };
14277        if crate::moe_f16g_mode() == 0 {
14278            return Ok(None);
14279        }
14280        // MEMRA_MOE_GATE is the BYTE-identity oracle between sequential-class dispatches; this
14281        // arm is a different numeric class with its own reference-band gate, so it must not
14282        // shadow that comparison.
14283        if std::env::var("MEMRA_MOE_GATE").is_ok() {
14284            return Ok(None);
14285        }
14286        let moe = cfg
14287            .moe
14288            .as_ref()
14289            .ok_or("grouped sigmoid prefill requires MoE model metadata")?;
14290        let n_embd = cfg.n_embd as usize;
14291        let n_expert = moe.expert_count as usize;
14292        let n_used = moe.expert_used_count as usize;
14293        let n_ff_exp = moe.expert_ff_length as usize;
14294        if !(f16g_proj_ok(m.gate_exps.qtype, n_embd)
14295            && f16g_proj_ok(m.up_exps.qtype, n_embd)
14296            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp))
14297        {
14298            return Ok(None);
14299        }
14300        // The sk visitor's direct-lane group cap (mirrors the grouped prime's guard). glm5's
14301        // 288 experts fit; a wider bank falls closed rather than erring mid-forward.
14302        if n_expert > 512 || n_used == 0 || n_used > 8 {
14303            return Ok(None);
14304        }
14305        let sigmoid = cfg
14306            .sigmoid_router()
14307            .ok_or("grouped sigmoid prefill requires the sigmoid router")?;
14308        // glm5_next carries the PRE form on every clamped layer (`clamp_exp_at`); a POST-form
14309        // arch reaching this arm is an unqualified semantic program, and the clamp-form law
14310        // says no dispatch site may pick a form by default. Err, not assert.
14311        let lim_exp = cfg.clamp_exp_at(il as u32);
14312        if matches!(lim_exp, Some(SwigluClamp::Post(_))) {
14313            return Err(
14314                "grouped sigmoid prefill is qualified for the PRE-clamped SwiGLU form only; \
14315                 a POST-clamp layer must ride the sequential arm"
14316                    .into(),
14317            );
14318        }
14319
14320        // 1. ROUTER. One selector shared with the sequential arm so changing dispatch cannot
14321        // change logits, selected expert ids, or routing weights (the routing-exactness half of
14322        // the acceptance gate holds by construction). The host readback here is the same one
14323        // the sequential arm performs; killing it is the L4 host-sync diet, not this arm.
14324        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
14325        Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
14326        let (sel_all, w_all) =
14327            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sigmoid)?;
14328        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
14329        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
14330        Self::trace_moe_input(e, il, t, n_embd, z)?;
14331
14332        let mprof = std::env::var("MEMRA_PRIME_PROF").as_deref() == Ok("1");
14333        let mut mt = std::time::Instant::now();
14334        let mut phase = |on: bool| -> f64 {
14335            if on {
14336                let _ = e.stream().synchronize();
14337                let v = mt.elapsed().as_secs_f64() * 1e3;
14338                mt = std::time::Instant::now();
14339                v
14340            } else {
14341                0.0
14342            }
14343        };
14344        let d_router = phase(mprof);
14345
14346        // 2. TOKEN-SORT BY EXPERT: expert-major CSR over the (token, expert) pairs.
14347        let n_pairs = t * n_used;
14348        if sel_all.len() < n_pairs || w_all.len() < n_pairs || z.len() < t * n_embd {
14349            return Err("grouped sigmoid prefill geometry".into());
14350        }
14351        let mut buckets: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14352        for (p, &s_id) in sel_all.iter().take(n_pairs).enumerate() {
14353            let s_id = s_id as usize;
14354            if s_id >= n_expert {
14355                return Err(format!("grouped prefill selection {s_id} >= {n_expert}").into());
14356            }
14357            buckets[s_id].push(p as i32);
14358        }
14359        let mut ex_ids: Vec<i32> = Vec::new();
14360        let mut ex_off: Vec<i32> = vec![0];
14361        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
14362        for (e_id, b) in buckets.iter().enumerate() {
14363            if !b.is_empty() {
14364                ex_ids.push(e_id as i32);
14365                ex_pairs.extend_from_slice(b);
14366                ex_off.push(ex_pairs.len() as i32);
14367            }
14368        }
14369        let n_active = ex_ids.len();
14370        if n_active == 0 || n_active > 512 {
14371            return Err(format!("grouped prefill n_active {n_active} outside 1..=512").into());
14372        }
14373        let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
14374
14375        // 3. MACRO PLANES. Gate/up per-CSR-row scales (before silu); down folds into the
14376        // scatter weight; `macro_scale` answers 1.0 for a macro-free bank, so the fold is
14377        // skipped rather than launched as a no-op.
14378        let wd: Vec<f32> = (0..n_pairs)
14379            .map(|p| w_all[p] * m.down_exps.macro_scale(sel_all[p] as usize))
14380            .collect();
14381
14382        // Interleaved gate/up slab strides (see moe_ffn_pairs / moe_ffn_dev).
14383        let (rbg_d, rbu_d) = if dev.gu_il {
14384            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14385            (sxx, sxx)
14386        } else {
14387            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14388        };
14389
14390        let exi = e.htod_i32(&ex_ids)?;
14391        let exo = e.htod_i32(&ex_off)?;
14392        let exp_d = e.htod_i32(&ex_pairs)?;
14393        let csr_tok_d = e.htod_i32(&csr_tok)?;
14394        let pw = e.htod(&wd)?;
14395
14396        // 4. GATE/UP grouped GEMMs over the bank, CSR order end-to-end.
14397        let (z16, zs) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
14398        let mut g = e.moe_f16_grouped(
14399            &dev.ptr_row,
14400            0,
14401            n_expert,
14402            &exi,
14403            &ex_off,
14404            &exo,
14405            &z16,
14406            &zs,
14407            n_embd,
14408            n_ff_exp,
14409            n_active,
14410            n_pairs,
14411            m.gate_exps.qtype,
14412            rbg_d,
14413        )?;
14414        if m.gate_exps.macros.is_some() {
14415            let mg: Vec<f32> = ex_pairs
14416                .iter()
14417                .map(|&p| m.gate_exps.macro_scale(sel_all[p as usize] as usize))
14418                .collect();
14419            let mg_d = e.htod(&mg)?;
14420            e.scale_rows(&mut g, &mg_d, n_ff_exp, n_pairs)?;
14421        }
14422        let mut u = e.moe_f16_grouped(
14423            &dev.ptr_row,
14424            1,
14425            n_expert,
14426            &exi,
14427            &ex_off,
14428            &exo,
14429            &z16,
14430            &zs,
14431            n_embd,
14432            n_ff_exp,
14433            n_active,
14434            n_pairs,
14435            m.up_exps.qtype,
14436            rbu_d,
14437        )?;
14438        if m.up_exps.macros.is_some() {
14439            let mu: Vec<f32> = ex_pairs
14440                .iter()
14441                .map(|&p| m.up_exps.macro_scale(sel_all[p as usize] as usize))
14442                .collect();
14443            let mu_d = e.htod(&mu)?;
14444            e.scale_rows(&mut u, &mu_d, n_ff_exp, n_pairs)?;
14445        }
14446
14447        // 5. EPILOGUE: glm5_next's PRE-clamped SwiGLU (the POST form was refused above); a
14448        // config with no live limit takes the plain-silu pair kernel.
14449        let act = match lim_exp {
14450            Some(SwigluClamp::Pre(limit)) => {
14451                let mut a = e.uninit(n_pairs * n_ff_exp)?;
14452                // Scales are 1.0: the per-expert macros already landed via scale_rows (an
14453                // exact *1.0 inside the kernel keeps the value chain unchanged).
14454                e.swiglu_preclamped_mul_scaled(
14455                    &g,
14456                    &u,
14457                    1.0,
14458                    1.0,
14459                    limit,
14460                    &mut a,
14461                    n_pairs * n_ff_exp,
14462                )?;
14463                a
14464            }
14465            None => e.moe_pairs_silu_mul(&g, &u, n_pairs * n_ff_exp)?,
14466            Some(SwigluClamp::Post(_)) => unreachable!("refused before any launch"),
14467        };
14468        let d_gemm_gu = phase(mprof);
14469
14470        // 6. DOWN grouped GEMM (CSR order), permute back to pair order, weighted scatter.
14471        let (a16, a_s) = e.moe_f16g_act(&act, None, n_ff_exp, n_pairs)?;
14472        let d_csr = e.moe_f16_grouped(
14473            &dev.ptr_row,
14474            2,
14475            n_expert,
14476            &exi,
14477            &ex_off,
14478            &exo,
14479            &a16,
14480            &a_s,
14481            n_ff_exp,
14482            n_embd,
14483            n_active,
14484            n_pairs,
14485            m.down_exps.qtype,
14486            m.down_exps.row_bytes,
14487        )?;
14488        let y_pair = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
14489        let toff: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14490        let tids: Vec<i32> = (0..n_pairs as i32).collect();
14491        let toff_d = e.htod_i32(&toff)?;
14492        let tids_d = e.htod_i32(&tids)?;
14493        // The scatter fully overwrites every (token, col): slot-ordered accumulation over the
14494        // token's n_used pairs, the sequential-axpy class.
14495        let mut moe_out = e.uninit(t * n_embd)?;
14496        e.moe_pairs_scatter(&y_pair, &pw, &toff_d, &tids_d, &mut moe_out, t, n_embd)?;
14497        let d_down = phase(mprof);
14498
14499        // 7. SHARED EXPERT: the canonical clamp-aware grouped add (reads clamp_shexp_at and
14500        // the optional shexp gate; glm5_next has a live shared expert on every MoE layer).
14501        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
14502        if mprof {
14503            let d_shared = phase(true);
14504            eprintln!(
14505                "[moe-grouped-prefill-prof] il={il} t={t} router={d_router:.1}ms \
14506                 gemm_gu={d_gemm_gu:.1}ms down_scatter={d_down:.1}ms shared={d_shared:.1}ms"
14507            );
14508        }
14509
14510        crate::MOE_GROUPED_PREFILL_DISPATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14511        // Once per layer per process: the engagement receipt line (the A/B greps for it; the
14512        // both-arms flag announce lives at the dispatch site).
14513        static GPF_LOGGED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
14514        let layer_bit = 1u64 << (il as u64 % 64);
14515        if GPF_LOGGED.fetch_or(layer_bit, std::sync::atomic::Ordering::Relaxed) & layer_bit == 0 {
14516            eprintln!(
14517                "[moe-grouped-prefill] execute layer={il} tokens={t} n_active={n_active} \
14518                 provenance=resident-slab router=sigmoid-host-oracle epilogue=pre-clamped \
14519                 macro_fold=gate-up-rows+down-weight performance_claim=false \
14520                 (logged once per layer)"
14521            );
14522        }
14523        Ok(Some(moe_out))
14524    }
14525
14526    #[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
14527    fn moe_gdec_token(
14528        e: &Engine,
14529        m: &MoeWeights,
14530        il: u16,
14531        max_block: usize,
14532        zt: &cudarc::driver::CudaView<f32>,
14533        sel: &[u32],
14534        w: &[f32],
14535        moe_out: &mut CudaSlice<f32>,
14536        tok: usize,
14537        n_embd: usize,
14538        n_ff_exp: usize,
14539        n_used: usize,
14540    ) -> Result<bool, Box<dyn std::error::Error>> {
14541        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14542        use cudarc::driver::DevicePtr;
14543        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
14544        let ptrs = e.with_moe_cache(max_block, |c, eng| {
14545            let mut g = [0u64; 8];
14546            let mut u = [0u64; 8];
14547            let mut d = [0u64; 8];
14548            for (j, &ex) in sel.iter().enumerate() {
14549                let ex = ex as u16;
14550                let (Some(sg), Some(su), Some(sd)) = (
14551                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
14552                    c.resident(BlockId::new(il, PROJ_UP, ex)),
14553                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
14554                ) else {
14555                    return Ok(None);
14556                };
14557                let __s = eng.stream();
14558                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
14559                let (pu, _e1) = c.slot(su).device_ptr(&__s);
14560                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
14561                g[j] = pg;
14562                u[j] = pu;
14563                d[j] = pd;
14564            }
14565            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
14566                for &ex in sel {
14567                    let ex = ex as u16;
14568                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
14569                        c.note_profile_hit(BlockId::new(il, proj, ex));
14570                    }
14571                }
14572            }
14573            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
14574            Ok(Some((g, u, d)))
14575        })?;
14576        let Some((g, u, d)) = ptrs else {
14577            return Ok(false);
14578        };
14579        let mut wv = [0f32; 8];
14580        wv[..n_used].copy_from_slice(w);
14581        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
14582        let act = e.moe_gate_up_silu8(
14583            crate::WPtr8(g),
14584            crate::WPtr8(u),
14585            zt,
14586            n_embd,
14587            n_ff_exp,
14588            n_used,
14589            m.gate_exps.qtype,
14590            m.up_exps.qtype,
14591            m.gate_exps.row_bytes,
14592            m.up_exps.row_bytes,
14593        )?;
14594        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
14595        e.moe_down8_fma_into(
14596            crate::WPtr8(d),
14597            crate::F32x8(wv),
14598            &act,
14599            &mut dst,
14600            n_ff_exp,
14601            n_embd,
14602            n_used,
14603            m.down_exps.qtype,
14604            m.down_exps.row_bytes,
14605        )?;
14606        Ok(true)
14607    }
14608
14609    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
14610    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
14611    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
14612    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
14613    #[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
14614    fn moe_cached_gemm_q8(
14615        e: &Engine,
14616        il: u16,
14617        proj: u8,
14618        ex: usize,
14619        m: &MoeWeights,
14620        max_block: usize,
14621        aq: &CudaSlice<i8>,
14622        ad: &CudaSlice<f32>,
14623    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14624        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14625        let exps = match proj {
14626            PROJ_GATE => &m.gate_exps,
14627            PROJ_UP => &m.up_exps,
14628            _ => &m.down_exps,
14629        };
14630        let layout = exps.expert_layout(ex);
14631        let id = BlockId::new(il, proj, ex as u16);
14632        let source = exps.expert_source(ex);
14633        e.with_moe_cache(max_block, |c, eng| {
14634            let slot = c.dispatch_source(id, source, eng)?;
14635            let DispatchSlot::Resident(sl) = slot;
14636            let buf = c.slot(sl);
14637            eng.qmatvec_expert_q8(
14638                buf,
14639                0..layout.len,
14640                aq,
14641                ad,
14642                1,
14643                exps.in_f,
14644                exps.out_f,
14645                layout.qtype,
14646                layout.row_bytes,
14647            )
14648        })
14649    }
14650
14651    fn moe_cached_gemm(
14652        e: &Engine,
14653        il: u16,
14654        proj: u8,
14655        ex: usize,
14656        m: &MoeWeights,
14657        max_block: usize,
14658        x: &cudarc::driver::CudaView<f32>,
14659    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14660        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
14661        let exps = match proj {
14662            PROJ_GATE => &m.gate_exps,
14663            PROJ_UP => &m.up_exps,
14664            _ => &m.down_exps,
14665        };
14666        let layout = exps.expert_layout(ex);
14667        let id = BlockId::new(il, proj, ex as u16);
14668        let source = exps.expert_source(ex);
14669        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
14670        e.with_moe_cache(max_block, |c, eng| {
14671            let slot = c.dispatch_source(id, source, eng)?;
14672            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
14673            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
14674            let DispatchSlot::Resident(sl) = slot;
14675            let buf = c.slot(sl);
14676            m.qmatvec_view(
14677                eng,
14678                buf,
14679                0..layout.len,
14680                x,
14681                1,
14682                exps.in_f,
14683                exps.out_f,
14684                layout.qtype,
14685                layout.row_bytes,
14686            )
14687        })
14688    }
14689
14690    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
14691    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
14692    /// so the current forward's backend assignment and output remain unchanged.
14693    fn moe_profile_admit_expert(
14694        e: &Engine,
14695        il: u16,
14696        ex: usize,
14697        m: &MoeWeights,
14698        max_block: usize,
14699    ) -> Result<(), Box<dyn std::error::Error>> {
14700        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14701        e.with_moe_cache(max_block, |cache, eng| {
14702            for (proj, exps) in [
14703                (PROJ_GATE, &m.gate_exps),
14704                (PROJ_UP, &m.up_exps),
14705                (PROJ_DOWN, &m.down_exps),
14706            ] {
14707                let id = BlockId::new(il, proj, ex as u16);
14708                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
14709            }
14710            Ok(())
14711        })
14712    }
14713
14714    /// Read a projection from the immutable residency set when present; otherwise use one
14715    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
14716    #[allow(clippy::too_many_arguments)]
14717    fn moe_frozen_gemm(
14718        e: &Engine,
14719        il: u16,
14720        proj: u8,
14721        ex: usize,
14722        m: &MoeWeights,
14723        max_block: usize,
14724        x: &cudarc::driver::CudaView<f32>,
14725        scratch: &mut Option<CudaSlice<u8>>,
14726        scratch_len: usize,
14727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14728        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
14729        let exps = match proj {
14730            PROJ_GATE => &m.gate_exps,
14731            PROJ_UP => &m.up_exps,
14732            _ => &m.down_exps,
14733        };
14734        let layout = exps.expert_layout(ex);
14735        let id = BlockId::new(il, proj, ex as u16);
14736        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
14737            let Some(slot) = cache.resident(id) else {
14738                return Ok(None);
14739            };
14740            let buf = cache.slot(slot);
14741            Ok(Some(m.qmatvec_view(
14742                eng,
14743                buf,
14744                0..layout.len,
14745                x,
14746                1,
14747                exps.in_f,
14748                exps.out_f,
14749                layout.qtype,
14750                layout.row_bytes,
14751            )?))
14752        })? {
14753            return Ok(output);
14754        }
14755        if scratch.is_none() {
14756            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
14757        }
14758        let scratch = scratch.as_mut().unwrap();
14759        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
14760        m.qmatvec_view(
14761            e,
14762            scratch,
14763            0..layout.len,
14764            x,
14765            1,
14766            exps.in_f,
14767            exps.out_f,
14768            layout.qtype,
14769            layout.row_bytes,
14770        )
14771    }
14772
14773    fn moe_prefetch_expert(
14774        e: &Engine,
14775        il: u16,
14776        ex: usize,
14777        m: &MoeWeights,
14778        max_block: usize,
14779        keep: &[crate::moe_cache::BlockId],
14780    ) -> Result<(), Box<dyn std::error::Error>> {
14781        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14782        e.with_moe_cache(max_block, |c, eng| {
14783            for (proj, exps) in [
14784                (PROJ_GATE, &m.gate_exps),
14785                (PROJ_UP, &m.up_exps),
14786                (PROJ_DOWN, &m.down_exps),
14787            ] {
14788                let id = BlockId::new(il, proj, ex as u16);
14789                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
14790            }
14791            Ok(())
14792        })
14793    }
14794
14795    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
14796    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
14797    fn moe_prefetch_disk_expert(
14798        e: &Engine,
14799        il: u16,
14800        ex: usize,
14801        m: &MoeWeights,
14802        max_block: usize,
14803        keep: &[crate::moe_cache::BlockId],
14804    ) -> Result<(), Box<dyn std::error::Error>> {
14805        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
14806        e.with_moe_cache(max_block, |c, eng| {
14807            for (proj, exps) in [
14808                (PROJ_GATE, &m.gate_exps),
14809                (PROJ_UP, &m.up_exps),
14810                (PROJ_DOWN, &m.down_exps),
14811            ] {
14812                let source = exps.expert_source(ex);
14813                if let crate::model::ExpertSource::Disk { .. } = &source {
14814                    let id = BlockId::new(il, proj, ex as u16);
14815                    let _ = c.prefetch_source(id, source, keep, eng)?;
14816                }
14817            }
14818            Ok(())
14819        })
14820    }
14821
14822    #[inline]
14823    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
14824        let _ = m.gate_exps.prefetch_expert_pages(ex);
14825        let _ = m.up_exps.prefetch_expert_pages(ex);
14826        let _ = m.down_exps.prefetch_expert_pages(ex);
14827    }
14828}
14829
14830// ================================================================================================
14831// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
14832//
14833// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
14834// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
14835// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
14836//
14837// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
14838// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
14839// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
14840// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
14841// identical to the per-token loop regardless of expert processing order.
14842//
14843// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
14844// ================================================================================================
14845
14846impl HybridModel {
14847    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
14848    /// sequential fused q8 program over the token axis; clamped layers use the separate
14849    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
14850    #[allow(clippy::too_many_arguments)]
14851    fn moe_ffn_grouped_resident_q8(
14852        e: &Engine,
14853        m: &MoeWeights,
14854        z: &CudaSlice<f32>,
14855        t: usize,
14856        cfg: &ModelConfig,
14857        il: u16,
14858        sel_all: &[u32],
14859        w_all: &[f32],
14860        table: &CudaSlice<u64>,
14861        gu_il: bool,
14862    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14863        let moe = cfg.moe.as_ref().unwrap();
14864        let n_embd = cfg.n_embd as usize;
14865        let n_expert = moe.expert_count as usize;
14866        let n_used = moe.expert_used_count as usize;
14867        let n_ff_exp = moe.expert_ff_length as usize;
14868        let n_pairs = t * n_used;
14869        debug_assert_eq!(sel_all.len(), n_pairs);
14870        debug_assert_eq!(w_all.len(), n_pairs);
14871        debug_assert!(
14872            m.gate_exps.macros.is_none()
14873                && m.up_exps.macros.is_none()
14874                && m.down_exps.macros.is_none(),
14875            "resident grouped q8 does not fold per-expert macro scales",
14876        );
14877
14878        // The rows twins run the resident sequential program verbatim on grid.z = token:
14879        // fused gate/up/SiLU per slot, batched activation quantization, then the original
14880        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
14881        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
14882        // never enter the softmax router.
14883        if !cfg.swiglu_clamped_at(il as u32) {
14884            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
14885            let sel_d = e.htod_i32(&sel)?;
14886            let w_d = e.htod(w_all)?;
14887            let (gate_row_bytes, up_row_bytes) = if gu_il {
14888                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14889                (combined, combined)
14890            } else {
14891                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14892            };
14893            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14894            let act = e.moe_gate_up_silu8_dev_q8_rows(
14895                table,
14896                &sel_d,
14897                &zq,
14898                &zd,
14899                t,
14900                n_embd,
14901                n_ff_exp,
14902                n_used,
14903                n_expert,
14904                m.gate_exps.qtype,
14905                m.up_exps.qtype,
14906                gate_row_bytes,
14907                up_row_bytes,
14908                &m.dev_macros,
14909            )?;
14910            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
14911            let mut moe_out = e.uninit(t * n_embd)?;
14912            e.moe_down8_fma_dev_q8_rows_g(
14913                table,
14914                &sel_d,
14915                &w_d,
14916                &aq2,
14917                &ad2,
14918                &mut moe_out,
14919                t,
14920                n_ff_exp,
14921                n_embd,
14922                n_used,
14923                n_expert,
14924                m.down_exps.qtype,
14925                m.down_exps.row_bytes,
14926            )?;
14927
14928            if std::env::var("MEMRA_MOE_STATS").is_ok() {
14929                let mut counts = vec![0usize; n_expert];
14930                for &expert in sel_all {
14931                    counts[expert as usize] += 1;
14932                }
14933                let mut sizes: Vec<usize> =
14934                    counts.into_iter().filter(|&count| count != 0).collect();
14935                sizes.sort_unstable();
14936                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
14937                println!(
14938                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
14939                     m_e: min={} median={} mean={mean:.1} max={}",
14940                    sizes.len(),
14941                    n_expert,
14942                    sizes.first().copied().unwrap_or(0),
14943                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
14944                    sizes.last().copied().unwrap_or(0),
14945                );
14946            }
14947            return Ok(moe_out);
14948        }
14949
14950        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
14951        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
14952        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
14953        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
14954        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
14955        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
14956        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
14957
14958        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
14959        for (pair, &expert) in pair_ex.iter().enumerate() {
14960            by_expert[expert as usize].push(pair as i32);
14961        }
14962
14963        let pair_tok_d = e.htod_i32(&pair_tok)?;
14964        let pair_ex_d = e.htod_i32(&pair_ex)?;
14965        let pair_w_d = e.htod(w_all)?;
14966        let tok_off_d = e.htod_i32(&tok_off)?;
14967        let tok_ids_d = e.htod_i32(&tok_ids)?;
14968
14969        let matvec = |proj: i32,
14970                      pair_rows: &CudaSlice<i32>,
14971                      aq: &CudaSlice<i8>,
14972                      ad: &CudaSlice<f32>,
14973                      in_f: usize,
14974                      out_f: usize,
14975                      qtype: i32,
14976                      row_bytes: usize|
14977         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14978            e.moe_pairs_matvec_q8(
14979                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
14980                row_bytes,
14981            )
14982        };
14983
14984        let (gate_row_bytes, up_row_bytes) = if gu_il {
14985            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
14986            (combined, combined)
14987        } else {
14988            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
14989        };
14990        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
14991        let gate = matvec(
14992            0,
14993            &pair_tok_d,
14994            &zq,
14995            &zd,
14996            n_embd,
14997            n_ff_exp,
14998            m.gate_exps.qtype,
14999            gate_row_bytes,
15000        )?;
15001        let up = matvec(
15002            1,
15003            &pair_tok_d,
15004            &zq,
15005            &zd,
15006            n_embd,
15007            n_ff_exp,
15008            m.up_exps.qtype,
15009            up_row_bytes,
15010        )?;
15011        let mut act = e.uninit(n_pairs * n_ff_exp)?;
15012        Self::ffn_act_lim(
15013            e,
15014            cfg,
15015            &gate,
15016            &up,
15017            1.0,
15018            1.0,
15019            cfg.clamp_exp_at(il as u32),
15020            &mut act,
15021            n_pairs * n_ff_exp,
15022        )?;
15023        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
15024        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
15025        let pair_self_d = e.htod_i32(&pair_self)?;
15026        let down = matvec(
15027            2,
15028            &pair_self_d,
15029            &aq2,
15030            &ad2,
15031            n_ff_exp,
15032            n_embd,
15033            m.down_exps.qtype,
15034            m.down_exps.row_bytes,
15035        )?;
15036        let mut moe_out = e.uninit(t * n_embd)?;
15037        e.moe_pairs_scatter(
15038            &down,
15039            &pair_w_d,
15040            &tok_off_d,
15041            &tok_ids_d,
15042            &mut moe_out,
15043            t,
15044            n_embd,
15045        )?;
15046
15047        if std::env::var("MEMRA_MOE_STATS").is_ok() {
15048            let mut sizes: Vec<usize> = by_expert
15049                .iter()
15050                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
15051                .collect();
15052            sizes.sort_unstable();
15053            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
15054            println!(
15055                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
15056                 m_e: min={} median={} mean={mean:.1} max={}",
15057                sizes.len(),
15058                n_expert,
15059                sizes.first().copied().unwrap_or(0),
15060                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
15061                sizes.last().copied().unwrap_or(0),
15062            );
15063        }
15064        Ok(moe_out)
15065    }
15066
15067    /// MEMRA_SHEXP_SPLIT worker: the shared expert's gate/up/down rows split across both
15068    /// devices (dev1 idles during E3), act halves exchanged both ways, downs row-split —
15069    /// per-element/per-row programs identical, so `sh` is BIT-IDENTICAL to the single-device
15070    /// arm. Process-static workspace pinned by the gate tensor pointer; rank1 holds one-time
15071    /// row-half replicas (~10MB/layer x 42). Returns None when ineligible.
15072    #[allow(clippy::too_many_arguments)]
15073    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
15074    fn shexp_split_matvec(
15075        e: &Engine,
15076        rank1: &Engine,
15077        wg: &CudaSlice<u8>,
15078        wu: &CudaSlice<u8>,
15079        wd: &CudaSlice<u8>,
15080        z: &CudaSlice<f32>,
15081        lim: Option<SwigluClamp>,
15082        cfg: &ModelConfig,
15083        il: u16,
15084        n_embd: usize,
15085        n_ff_sh: usize,
15086    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15087        use cudarc::driver::DevicePtr;
15088        if !n_ff_sh.is_multiple_of(2) || !n_embd.is_multiple_of(2) {
15089            return Ok(None);
15090        }
15091        let hf = n_ff_sh / 2;
15092        let nd = n_embd / 2;
15093        struct Rep {
15094            wg1: CudaSlice<u8>,
15095            wu1: CudaSlice<u8>,
15096            wd1: CudaSlice<u8>,
15097        }
15098        struct SplitWs {
15099            pin_dev: usize,
15100            // e side
15101            gate0: CudaSlice<f32>,
15102            up0: CudaSlice<f32>,
15103            act: CudaSlice<f32>,
15104            sh_buf: CudaSlice<f32>,
15105            ev_z: cudarc::driver::CudaEvent,
15106            ev_act0: cudarc::driver::CudaEvent,
15107            // rank1 side
15108            z1: CudaSlice<f32>,
15109            g1: CudaSlice<f32>,
15110            u1: CudaSlice<f32>,
15111            a1h: CudaSlice<f32>,
15112            act1: CudaSlice<f32>,
15113            y1: CudaSlice<f32>,
15114            ev_act1: cudarc::driver::CudaEvent,
15115            ev_y1: cudarc::driver::CudaEvent,
15116            raw_act_e: u64,
15117            raw_sh_e: u64,
15118            raw_z1: u64,
15119            raw_a1h: u64,
15120            raw_act1: u64,
15121            raw_y1: u64,
15122        }
15123        static WS: std::sync::Mutex<Option<SplitWs>> = std::sync::Mutex::new(None);
15124        static REPS: std::sync::Mutex<Option<std::collections::HashMap<u64, Rep>>> =
15125            std::sync::Mutex::new(None);
15126        let mut guard = WS.lock().map_err(|_| "shexp split lock is poisoned")?;
15127        let mut reps_guard = REPS.lock().map_err(|_| "shexp reps lock is poisoned")?;
15128        let reps = reps_guard.get_or_insert_with(std::collections::HashMap::new);
15129        let pins = e.ctx().ordinal();
15130        if guard.as_ref().is_none_or(|w| w.pin_dev != pins) {
15131            let (gate0, up0, act, sh_buf, ev_z, ev_act0) = {
15132                let _m = e.gpu.enter_main()?;
15133                (
15134                    e.htod(&vec![0.0f32; hf])?,
15135                    e.htod(&vec![0.0f32; hf])?,
15136                    e.htod(&vec![0.0f32; n_ff_sh])?,
15137                    e.htod(&vec![0.0f32; n_embd])?,
15138                    e.ctx().new_event(None)?,
15139                    e.ctx().new_event(None)?,
15140                )
15141            };
15142            let (z1, g1, u1, a1h, act1, y1, ev_act1, ev_y1) = {
15143                let _r = rank1.gpu.enter_main()?;
15144                (
15145                    rank1.htod(&vec![0.0f32; n_embd])?,
15146                    rank1.htod(&vec![0.0f32; hf])?,
15147                    rank1.htod(&vec![0.0f32; hf])?,
15148                    rank1.htod(&vec![0.0f32; hf])?,
15149                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
15150                    rank1.htod(&vec![0.0f32; nd])?,
15151                    rank1.ctx().new_event(None)?,
15152                    rank1.ctx().new_event(None)?,
15153                )
15154            };
15155            let (raw_act_e, raw_sh_e) = {
15156                let _m = e.gpu.enter_main()?;
15157                let stream = e.stream();
15158                let (a, _g0) = act.device_ptr(&stream);
15159                let (b, _g1) = sh_buf.device_ptr(&stream);
15160                (a, b)
15161            };
15162            let (raw_z1, raw_a1h, raw_act1, raw_y1) = {
15163                let _r = rank1.gpu.enter_main()?;
15164                let rs = rank1.stream();
15165                let (a, _g0) = z1.device_ptr(&rs);
15166                let (b, _g1) = a1h.device_ptr(&rs);
15167                let (c, _g2) = act1.device_ptr(&rs);
15168                let (d, _g3) = y1.device_ptr(&rs);
15169                (a, b, c, d)
15170            };
15171            *guard = Some(SplitWs {
15172                pin_dev: pins,
15173                gate0,
15174                up0,
15175                act,
15176                sh_buf,
15177                ev_z,
15178                ev_act0,
15179                z1,
15180                g1,
15181                u1,
15182                a1h,
15183                act1,
15184                y1,
15185                ev_act1,
15186                ev_y1,
15187                raw_act_e,
15188                raw_sh_e,
15189                raw_z1,
15190                raw_a1h,
15191                raw_act1,
15192                raw_y1,
15193            });
15194        }
15195        let ws = guard.as_mut().expect("armed above");
15196        let wg_pin = {
15197            let _m = e.gpu.enter_main()?;
15198            let stream = e.stream();
15199            let (p, _g) = wg.device_ptr(&stream);
15200            p
15201        };
15202        if !reps.contains_key(&wg_pin) {
15203            // One-time per layer: upload rank1's row halves (gate/up rows [hf..], down rows [nd..]).
15204            let up = |src: &CudaSlice<u8>,
15205                      off_bytes: usize,
15206                      len: usize|
15207             -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
15208                use cudarc::driver::sys;
15209                let sptr = {
15210                    let _m = e.gpu.enter_main()?;
15211                    let stream = e.stream();
15212                    let (p, _g) = src.device_ptr(&stream);
15213                    p + off_bytes as u64
15214                };
15215                let dst = {
15216                    let _r = rank1.gpu.enter_main()?;
15217                    rank1.alloc_u8_uninit(len)?
15218                };
15219                let dptr = {
15220                    let _r = rank1.gpu.enter_main()?;
15221                    let rs = rank1.stream();
15222                    let (p, _g) = dst.device_ptr(&rs);
15223                    p
15224                };
15225                let _r = rank1.gpu.enter_main()?;
15226                let r = unsafe {
15227                    sys::cuMemcpyAsync(
15228                        dptr as sys::CUdeviceptr,
15229                        sptr as sys::CUdeviceptr,
15230                        len,
15231                        rank1.stream().cu_stream() as sys::CUstream,
15232                    )
15233                };
15234                if r != sys::CUresult::CUDA_SUCCESS {
15235                    return Err(format!("shexp split replica upload: {r:?}").into());
15236                }
15237                rank1.stream().synchronize()?;
15238                Ok(dst)
15239            };
15240            let wg1 = up(wg, hf * n_embd * 2, hf * n_embd * 2)?;
15241            let wu1 = up(wu, hf * n_embd * 2, hf * n_embd * 2)?;
15242            let wd1 = up(wd, nd * n_ff_sh * 2, nd * n_ff_sh * 2)?;
15243            reps.insert(wg_pin, Rep { wg1, wu1, wd1 });
15244        }
15245        let _ = il;
15246        // Per token, evented split flow.
15247        let raw_z = {
15248            let _m = e.gpu.enter_main()?;
15249            let stream = e.stream();
15250            let (p, _g) = z.device_ptr(&stream);
15251            ws.ev_z.record(&stream)?;
15252            p
15253        };
15254        // rank1: pull z, its dual half, its act half; push act half to e; pull e's act half.
15255        {
15256            let rep = reps.get(&wg_pin).expect("uploaded above");
15257            let _r = rank1.gpu.enter_main()?;
15258            rank1.stream().wait(&ws.ev_z)?;
15259            crate::tp::raw_copy_bytes(ws.raw_z1, raw_z, n_embd * 4, rank1)?;
15260            let SplitWs {
15261                z1, g1, u1, a1h, ..
15262            } = &mut *ws;
15263            rank1.matvec_bf16_dual_into(&rep.wg1, &rep.wu1, z1, g1, u1, n_embd, hf)?;
15264            Self::ffn_act_lim(rank1, cfg, g1, u1, 1.0, 1.0, lim, a1h, hf)?;
15265            // local place into act1[hf..] + P2P push into e's act[hf..]
15266            crate::tp::raw_copy_bytes(ws.raw_act1 + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15267            crate::tp::raw_copy_bytes(ws.raw_act_e + (hf * 4) as u64, ws.raw_a1h, hf * 4, rank1)?;
15268            ws.ev_act1.record(&rank1.stream())?;
15269        }
15270        // e: its dual half + act half; publish it; then wait rank1's half and run down lo.
15271        {
15272            let _m = e.gpu.enter_main()?;
15273            let SplitWs {
15274                gate0, up0, act, ..
15275            } = &mut *ws;
15276            let wg_lo = wg.slice(0..hf * n_embd * 2);
15277            let wu_lo = wu.slice(0..hf * n_embd * 2);
15278            e.matvec_bf16_dual_view_into(&wg_lo, &wu_lo, z, gate0, up0, n_embd, hf)?;
15279            Self::ffn_act_lim(e, cfg, gate0, up0, 1.0, 1.0, lim, act, hf)?;
15280            ws.ev_act0.record(&e.stream())?;
15281        }
15282        // rank1: pull e's act half into act1[0..hf], run down hi, push y half.
15283        {
15284            let rep = reps.get(&wg_pin).expect("uploaded above");
15285            let _r = rank1.gpu.enter_main()?;
15286            rank1.stream().wait(&ws.ev_act0)?;
15287            crate::tp::raw_copy_bytes(ws.raw_act1, ws.raw_act_e, hf * 4, rank1)?;
15288            let SplitWs { act1, y1, .. } = &mut *ws;
15289            rank1.matvec_bf16_into(&rep.wd1, act1, y1, n_ff_sh, nd)?;
15290            crate::tp::raw_copy_bytes(ws.raw_sh_e + (nd * 4) as u64, ws.raw_y1, nd * 4, rank1)?;
15291            ws.ev_y1.record(&rank1.stream())?;
15292        }
15293        // e: down lo into sh_buf[0..nd]; join rank1's half; hand back an owned sh.
15294        {
15295            let _m = e.gpu.enter_main()?;
15296            e.stream().wait(&ws.ev_act1)?;
15297            let SplitWs { act, sh_buf, .. } = &mut *ws;
15298            let wd_lo = wd.slice(0..nd * n_ff_sh * 2);
15299            e.matvec_bf16_view_into(&wd_lo, act, sh_buf, n_ff_sh, nd)?;
15300            e.stream().wait(&ws.ev_y1)?;
15301            let mut sh = e.uninit(n_embd)?;
15302            {
15303                let mut dst = sh.slice_mut(0..n_embd);
15304                e.stream()
15305                    .memcpy_dtod(&ws.sh_buf.slice(0..n_embd), &mut dst)?;
15306            }
15307            Ok(Some(sh))
15308        }
15309    }
15310
15311    /// SHEXP OVERLAP issue (MEMRA_SHEXP_OVERLAP=1): the shared expert reads only `z`, so
15312    /// its kernels (dual matvec+SwiGLU, down) are issued on e's stream from the routes
15313    /// PREJOIN hook — they execute while the peer rank drains its sweep, filling dev0's
15314    /// join wait. The down lands in ITS OWN row (the pre-#2e split program — receipted
15315    /// bit-identical to the fused down+addscale) and `shexp_overlap_apply` adds it after
15316    /// the join with the exact add_scaled_rows expression: values unchanged.
15317    fn shexp_overlap_issue(
15318        e: &Engine,
15319        m: &MoeWeights,
15320        z: &CudaSlice<f32>,
15321        cfg: &ModelConfig,
15322        il: u16,
15323        n_embd: usize,
15324    ) -> Result<bool, Box<dyn std::error::Error>> {
15325        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15326            return Ok(false);
15327        }
15328        let (
15329            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15330            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15331            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15332        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15333        else {
15334            return Ok(false);
15335        };
15336        let n_ff_sh = m
15337            .gate_shexp
15338            .as_ref()
15339            .expect("matched Some above")
15340            .out_features();
15341        // The dual-silu epilogue is step35's POST form only; a PRE-clamped layer declines here
15342        // and takes the unfused seam.
15343        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15344            return Ok(false);
15345        };
15346        let mut guard = SHEXP_OV_WS
15347            .lock()
15348            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15349        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15350        if guard
15351            .as_ref()
15352            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15353        {
15354            *guard = Some((
15355                pins.0,
15356                pins.1,
15357                pins.2,
15358                e.uninit(n_ff_sh)?,
15359                e.uninit(n_embd)?,
15360            ));
15361        }
15362        let (_, _, _, act, sh) = guard.as_mut().expect("armed above");
15363        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim)?;
15364        e.matvec_bf16_into(wd, act, sh, n_ff_sh, n_embd)?;
15365        drop(guard);
15366        Ok(true)
15367    }
15368
15369    /// SHEXP ON DEV1 issue (MEMRA_SHEXP_DEV1=1): the shared expert runs on rank1 — the
15370    /// idle device — with replica weights (one-time P2P upload), the SAME kernels and the
15371    /// SAME split program as the dev0 overlap (dual matvec+SwiGLU, f32acc down), so the
15372    /// values are bit-identical. z rides one 16KB P2P pull behind an e-stream event; the
15373    /// down row lands root-resident (single P2P store pass); apply waits ev_done on e.
15374    #[allow(clippy::too_many_arguments)]
15375    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
15376    fn shexp_dev1_issue(
15377        e: &Engine,
15378        rank1: &Engine,
15379        m: &MoeWeights,
15380        z: &CudaSlice<f32>,
15381        cfg: &ModelConfig,
15382        il: u16,
15383        n_embd: usize,
15384    ) -> Result<bool, Box<dyn std::error::Error>> {
15385        use cudarc::driver::DevicePtr;
15386        if cfg.m3.is_some() || m.gate_inp_shexp.is_some() {
15387            return Ok(false);
15388        }
15389        let (
15390            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
15391            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
15392            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
15393        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15394        else {
15395            return Ok(false);
15396        };
15397        let n_ff_sh = m
15398            .gate_shexp
15399            .as_ref()
15400            .expect("matched Some above")
15401            .out_features();
15402        // POST-form epilogue only (see `fused_post_limit`): a PRE-clamped layer declines
15403        // (Ok(false) = nothing issued, caller falls back per column).
15404        let Ok(lim) = Self::fused_post_limit(cfg.clamp_shexp_at(il as u32)) else {
15405            return Ok(false);
15406        };
15407        // Shared scratch, geometry-keyed.
15408        let mut ws_guard = SHEXP_D1_WS
15409            .lock()
15410            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15411        if ws_guard
15412            .as_ref()
15413            .is_none_or(|(k, ..)| *k != (n_embd, n_ff_sh))
15414        {
15415            let (act1, z1, ev_done) = {
15416                let _r1 = rank1.gpu.enter_main()?;
15417                (
15418                    rank1.htod(&vec![0.0f32; n_ff_sh])?,
15419                    rank1.htod(&vec![0.0f32; n_embd])?,
15420                    rank1.ctx().new_event(None)?,
15421                )
15422            };
15423            let (sh_root, ev_z) = {
15424                let _main = e.gpu.enter_main()?;
15425                (e.htod(&vec![0.0f32; n_embd])?, e.ctx().new_event(None)?)
15426            };
15427            *ws_guard = Some(((n_embd, n_ff_sh), act1, z1, sh_root, ev_z, ev_done));
15428        }
15429        // Per-LAYER weight replicas (gate/up/down differ per layer): one-time P2P upload.
15430        let mut reps_guard = SHEXP_D1_REPS
15431            .lock()
15432            .map_err(|_| "shexp dev1 replica lock is poisoned")?;
15433        let reps = reps_guard.get_or_insert_with(Default::default);
15434        if !reps.contains_key(&il) {
15435            let (wg1, wu1, wd1) = {
15436                let _r1 = rank1.gpu.enter_main()?;
15437                (
15438                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15439                    rank1.alloc_u8_uninit(n_ff_sh * n_embd * 2)?,
15440                    rank1.alloc_u8_uninit(n_embd * n_ff_sh * 2)?,
15441                )
15442            };
15443            for (src, dst) in [(wg, &wg1), (wu, &wu1), (wd, &wd1)] {
15444                let s_ptr = {
15445                    let _main = e.gpu.enter_main()?;
15446                    let stream = e.stream();
15447                    let (p, _g) = src.device_ptr(&stream);
15448                    p
15449                };
15450                let d_ptr = {
15451                    let _r1 = rank1.gpu.enter_main()?;
15452                    let stream = rank1.stream();
15453                    let (p, _g) = dst.device_ptr(&stream);
15454                    p
15455                };
15456                let _r1 = rank1.gpu.enter_main()?;
15457                crate::tp::raw_copy_bytes(d_ptr, s_ptr, src.len(), rank1)?;
15458            }
15459            {
15460                let _r1 = rank1.gpu.enter_main()?;
15461                rank1.stream().synchronize()?;
15462            }
15463            reps.insert(il, (wg1, wu1, wd1));
15464        }
15465        let (wg1, wu1, wd1) = reps.get(&il).expect("armed above");
15466        let (_, act1, z1, sh_root, ev_z, ev_done) = ws_guard.as_mut().expect("armed above");
15467        // z ready on e's stream -> rank1 pulls it, runs the split shexp, pushes the down
15468        // row root-side (single store pass), rings ev_done.
15469        let (raw_z, raw_sh) = {
15470            let _main = e.gpu.enter_main()?;
15471            let stream = e.stream();
15472            let (a, _g0) = z.device_ptr(&stream);
15473            let (b, _g1) = sh_root.device_ptr(&stream);
15474            ev_z.record(&stream)?;
15475            (a, b)
15476        };
15477        {
15478            let _r1 = rank1.gpu.enter_main()?;
15479            rank1.stream().wait(ev_z)?;
15480            let raw_z1 = {
15481                let stream = rank1.stream();
15482                let (p, _g) = z1.device_ptr(&stream);
15483                p
15484            };
15485            crate::tp::raw_copy_bytes(raw_z1, raw_z, n_embd * 4, rank1)?;
15486            rank1.matvec_bf16_dual_silu_into(wg1, wu1, z1, act1, n_embd, n_ff_sh, lim)?;
15487            // down writes the ROOT-resident row over P2P via the raw-output twin of
15488            // matvec_bf16_into: reuse the view launcher with a slice view is not possible
15489            // cross-device, so launch on the raw pointer.
15490            rank1.matvec_bf16_raw_out(wd1, act1, raw_sh, n_ff_sh, n_embd)?;
15491            ev_done.record(&rank1.stream())?;
15492        }
15493        Ok(true)
15494    }
15495
15496    /// Apply the dev1 shared expert: wait ev_done on e, then the exact add_scaled_rows.
15497    fn shexp_dev1_apply(
15498        e: &Engine,
15499        output: &mut CudaSlice<f32>,
15500        n_embd: usize,
15501    ) -> Result<(), Box<dyn std::error::Error>> {
15502        let guard = SHEXP_D1_WS
15503            .lock()
15504            .map_err(|_| "shexp dev1 workspace lock is poisoned")?;
15505        let (pin, _, _, sh_root, _, ev_done) =
15506            guard.as_ref().ok_or("shexp dev1 apply without issue")?;
15507        if pin.0 != n_embd {
15508            return Err("shexp dev1 width drifted".into());
15509        }
15510        let _main = e.gpu.enter_main()?;
15511        e.stream().wait(ev_done)?;
15512        static ONES_D1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15513            std::sync::Mutex::new(None);
15514        let mut og = ONES_D1.lock().map_err(|_| "ones lock is poisoned")?;
15515        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15516            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15517        }
15518        let ones = &og.as_ref().expect("armed above").1;
15519        e.add_scaled_rows(sh_root, ones, output, n_embd, 1)?;
15520        Ok(())
15521    }
15522
15523    /// MOE TAIL FUSION M1 helper: pre-arm the overlap ws + persistent ones row and
15524    /// return their RAW pointers (None when the overlap is ineligible — the caller then
15525    /// takes the split path). Mirrors shexp_overlap_issue's eligibility exactly.
15526    fn shexp_overlap_tail_ptrs(
15527        e: &Engine,
15528        m: &MoeWeights,
15529        cfg: &ModelConfig,
15530        n_embd: usize,
15531    ) -> Result<Option<(u64, u64)>, 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(None);
15535        }
15536        let (
15537            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15538            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15539            Some(crate::model::GpuTensor::FloatBf16 { .. }),
15540        ) = (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15541        else {
15542            return Ok(None);
15543        };
15544        let n_ff_sh = m
15545            .gate_shexp
15546            .as_ref()
15547            .expect("matched Some above")
15548            .out_features();
15549        let mut guard = SHEXP_OV_WS
15550            .lock()
15551            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15552        let pins = (e.ctx().ordinal(), n_embd, n_ff_sh);
15553        if guard
15554            .as_ref()
15555            .is_none_or(|(d, ne, nf, ..)| (*d, *ne, *nf) != pins)
15556        {
15557            *guard = Some((
15558                pins.0,
15559                pins.1,
15560                pins.2,
15561                e.uninit(n_ff_sh)?,
15562                e.uninit(n_embd)?,
15563            ));
15564        }
15565        let sh_raw = {
15566            let (_, _, _, _, sh) = guard.as_ref().expect("armed above");
15567            let stream = e.stream();
15568            let (p, _g) = sh.device_ptr(&stream);
15569            p
15570        };
15571        static ONES_T3: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15572            std::sync::Mutex::new(None);
15573        let mut og = ONES_T3.lock().map_err(|_| "ones lock is poisoned")?;
15574        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15575            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15576        }
15577        let ones_raw = {
15578            let stream = e.stream();
15579            let (p, _g) = og.as_ref().expect("armed above").1.device_ptr(&stream);
15580            p
15581        };
15582        Ok(Some((sh_raw, ones_raw)))
15583    }
15584
15585    /// Apply the overlapped shared expert: output[r] += sh[r] * 1.0 — the exact
15586    /// add_scaled_rows program the split path used (persistent ones row, no htod).
15587    fn shexp_overlap_apply(
15588        e: &Engine,
15589        output: &mut CudaSlice<f32>,
15590        n_embd: usize,
15591    ) -> Result<(), Box<dyn std::error::Error>> {
15592        let guard = SHEXP_OV_WS
15593            .lock()
15594            .map_err(|_| "shexp overlap workspace lock is poisoned")?;
15595        let (_, ne, _, _, sh) = guard.as_ref().ok_or("shexp overlap apply without issue")?;
15596        if *ne != n_embd {
15597            return Err("shexp overlap width drifted".into());
15598        }
15599        static ONES_OV: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15600            std::sync::Mutex::new(None);
15601        let mut og = ONES_OV.lock().map_err(|_| "ones lock is poisoned")?;
15602        if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15603            *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15604        }
15605        let ones = &og.as_ref().expect("armed above").1;
15606        e.add_scaled_rows(sh, ones, output, n_embd, 1)?;
15607        Ok(())
15608    }
15609
15610    fn moe_ffn_grouped_add_shared(
15611        e: &Engine,
15612        m: &MoeWeights,
15613        z: &CudaSlice<f32>,
15614        t: usize,
15615        cfg: &ModelConfig,
15616        il: u16,
15617        moe_out: &mut CudaSlice<f32>,
15618    ) -> Result<(), Box<dyn std::error::Error>> {
15619        // MEMRA_STEP_TP_TIMING=1: shared-expert wall (syncs e's stream at exit to bill the
15620        // queued matmuls here rather than at the next host readback).
15621        static SHEXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15622        static SHEXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
15623        let shexp_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
15624        let shexp_started = shexp_timing.then(std::time::Instant::now);
15625        let result = Self::moe_ffn_grouped_add_shared_inner(e, m, z, t, cfg, il, moe_out);
15626        if let Some(started) = shexp_started {
15627            use std::sync::atomic::Ordering;
15628            e.stream().synchronize()?;
15629            let ns = SHEXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
15630                + started.elapsed().as_nanos() as u64;
15631            let calls = SHEXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
15632            if calls.is_multiple_of(430) {
15633                eprintln!(
15634                    "[moe-shexp-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
15635                    ns as f64 / 1.0e6,
15636                    ns as f64 / calls as f64 / 1.0e3,
15637                );
15638            }
15639        }
15640        result
15641    }
15642
15643    #[allow(clippy::too_many_arguments)]
15644    fn moe_ffn_grouped_add_shared_inner(
15645        e: &Engine,
15646        m: &MoeWeights,
15647        z: &CudaSlice<f32>,
15648        t: usize,
15649        cfg: &ModelConfig,
15650        il: u16,
15651        moe_out: &mut CudaSlice<f32>,
15652    ) -> Result<(), Box<dyn std::error::Error>> {
15653        let n_embd = cfg.n_embd as usize;
15654        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
15655            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
15656        {
15657            let n_ff_sh = gate_shexp.out_features();
15658            let lim = cfg.clamp_shexp_at(il as u32);
15659            // T=1 DECODE FUSION (2026-08-20): the ffn_swiglu_decode fast-path program, ported
15660            // — one shared quantize feeds gate+up (matmul at m=1 quantizes internally per
15661            // call with identical bytes, so sharing it is bit-identical), the dual NVFP4/Q8
15662            // launch covers both when available, and silu_mul_scaled_q8_1 emits down's
15663            // operand pre-quantized (kernel_check-proven identities). This path measured
15664            // 167us/layer as separate matmuls + 5 allocs at decode.
15665            let fused = t == 1
15666                && lim.is_none()
15667                && cfg.m3.is_none()
15668                && e.uses_q8_1_fast(gate_shexp)
15669                && e.uses_q8_1_fast(up_shexp);
15670            let canonical_w4a16_rows =
15671                t <= 32 && m.step_ep.as_ref().is_some_and(|ep| ep.nvfp4_device_routes);
15672            // MEMRA_BF16_MMV class: both projections in ONE launch (bit-identical per row to
15673            // the two matvec_bf16 launches matmul would issue). W4A16 distributed execution
15674            // uses the same row program for decode and verify; a t=1-only fusion accumulated
15675            // sub-ULP residual drift from the first MoE layer onward.
15676            let bf16_dual = if (t == 1 || canonical_w4a16_rows)
15677                && crate::Engine::bf16_mmv_on()
15678                && n_embd.is_multiple_of(8)
15679            {
15680                match (gate_shexp, up_shexp) {
15681                    (
15682                        crate::model::GpuTensor::FloatBf16 { data: wg, .. },
15683                        crate::model::GpuTensor::FloatBf16 { data: wu, .. },
15684                    ) => Some((wg, wu)),
15685                    _ => None,
15686                }
15687            } else {
15688                None
15689            };
15690            let sh = if let Some((wg, wu)) = bf16_dual {
15691                // Persistent shared-expert workspace: sizes are constant across every MoE
15692                // layer, so one process-level set pinned by (device, n_embd, n_ff_sh) removes
15693                // the four per-layer allocations. Buffers are fully overwritten each call.
15694                type SharedExpertWorkspace = (
15695                    usize,
15696                    usize,
15697                    usize,
15698                    CudaSlice<f32>,
15699                    CudaSlice<f32>,
15700                    CudaSlice<f32>,
15701                    CudaSlice<f32>,
15702                );
15703                static SHEXP_WS: std::sync::Mutex<
15704                    Option<std::collections::HashMap<usize, SharedExpertWorkspace>>,
15705                > = std::sync::Mutex::new(None);
15706                let down_bf16 = match down_shexp {
15707                    crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15708                    _ => None,
15709                };
15710                let mut guard = SHEXP_WS
15711                    .lock()
15712                    .map_err(|_| "shexp workspace lock is poisoned")?;
15713                let capacity = if canonical_w4a16_rows { 32 } else { 1 };
15714                let device = e.ctx().ordinal();
15715                let workspaces = guard.get_or_insert_with(Default::default);
15716                if workspaces
15717                    .get(&device)
15718                    .is_none_or(|(ne, nf, cap, ..)| (*ne, *nf, *cap) != (n_embd, n_ff_sh, capacity))
15719                {
15720                    workspaces.insert(
15721                        device,
15722                        (
15723                            n_embd,
15724                            n_ff_sh,
15725                            capacity,
15726                            e.uninit(capacity * n_ff_sh)?,
15727                            e.uninit(capacity * n_ff_sh)?,
15728                            e.uninit(capacity * n_ff_sh)?,
15729                            e.uninit(capacity * n_embd)?,
15730                        ),
15731                    );
15732                }
15733                // MEMRA_SHEXP_SPLIT=1: bit-identical row split across both devices; falls
15734                // through to the single-device arm when ineligible.
15735                {
15736                    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15737                    let split_on = *ON
15738                        .get_or_init(|| std::env::var("MEMRA_SHEXP_SPLIT").as_deref() == Ok("1"));
15739                    if split_on
15740                        && t == 1
15741                        && let (Some(wd), Some(rank1)) = (
15742                            match down_shexp {
15743                                crate::model::GpuTensor::FloatBf16 { data, .. } => Some(data),
15744                                _ => None,
15745                            },
15746                            m.step_tp.as_ref().and_then(|st| st.runtime.rank_engine(1)),
15747                        )
15748                        && let Some(sh) = Self::shexp_split_matvec(
15749                            e, rank1, wg, wu, wd, z, lim, cfg, il, n_embd, n_ff_sh,
15750                        )?
15751                    {
15752                        drop(guard);
15753                        let gate = match &m.gate_inp_shexp {
15754                            Some(gate_inp_shexp) => {
15755                                e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15756                            }
15757                            None => e.htod(&vec![1.0f32; t])?,
15758                        };
15759                        e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
15760                        return Ok(());
15761                    }
15762                }
15763                let (_, _, _, gate, up, act, sh_buf) = workspaces
15764                    .get_mut(&device)
15765                    .expect("shexp workspace initialized above");
15766                // FUSION #2b needs the POST form its epilogue hardcodes; m3's swigluoai and
15767                // glm5_next's PRE clamp both take the unfused dual-matmul + ffn_act_lim arm.
15768                if let (true, Ok(lim_post)) = (cfg.m3.is_none(), Self::fused_post_limit(lim)) {
15769                    // dual matvec + SwiGLU act in one launch — exact dual per-row program +
15770                    // exact silu/clamped expression, bit-identical.
15771                    if canonical_w4a16_rows {
15772                        e.matvec_bf16_dual_silu_rows_into(
15773                            wg, wu, z, act, n_embd, n_ff_sh, lim_post, t,
15774                        )?;
15775                    } else {
15776                        e.matvec_bf16_dual_silu_into(wg, wu, z, act, n_embd, n_ff_sh, lim_post)?;
15777                    }
15778                    let _ = (&gate, &up);
15779                } else {
15780                    e.matvec_bf16_dual_into(wg, wu, z, gate, up, n_embd, n_ff_sh)?;
15781                    Self::ffn_act_lim(e, cfg, gate, up, 1.0, 1.0, lim, act, n_ff_sh)?;
15782                }
15783                if let Some(down) = down_bf16 {
15784                    if canonical_w4a16_rows {
15785                        e.matvec_bf16_rows_into(down, act, sh_buf, n_ff_sh, n_embd, t)?;
15786                        let mut sh = e.uninit(t * n_embd)?;
15787                        {
15788                            let mut dst = sh.slice_mut(0..t * n_embd);
15789                            e.stream()
15790                                .memcpy_dtod(&sh_buf.slice(0..t * n_embd), &mut dst)?;
15791                        }
15792                        sh
15793                    } else {
15794                        // FUSION #2e (gate-less shexp only, MEMRA_FUSE_DOWN_ADDSCALE=0 reverts):
15795                        // down matvec + scaled accumulate straight into moe_out in ONE launch —
15796                        // exact f32acc per-row program + the exact add_scaled_rows expression
15797                        // (dst[r] += y_r * 1.0). Replaces down + ownership alloc + 16KB copy +
15798                        // add_scaled (3 launches + alloc -> 1 launch); bit-identical because the
15799                        // accumulate consumes the same f32 the split path stored and reloaded.
15800                        static FUSE_DA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15801                        let fuse_da = *FUSE_DA.get_or_init(|| {
15802                            std::env::var("MEMRA_FUSE_DOWN_ADDSCALE").as_deref() != Ok("0")
15803                        });
15804                        if fuse_da && m.gate_inp_shexp.is_none() {
15805                            static ONES1: std::sync::Mutex<Option<(usize, CudaSlice<f32>)>> =
15806                                std::sync::Mutex::new(None);
15807                            let mut og = ONES1.lock().map_err(|_| "shexp ones lock is poisoned")?;
15808                            if og.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
15809                                *og = Some((e.ctx().ordinal(), e.htod(&[1.0f32])?));
15810                            }
15811                            let ones = &og.as_ref().expect("armed above").1;
15812                            e.matvec_bf16_down_addscale_into(
15813                                down, act, ones, moe_out, n_ff_sh, n_embd,
15814                            )?;
15815                            return Ok(());
15816                        }
15817                        e.matvec_bf16_into(down, act, sh_buf, n_ff_sh, n_embd)?;
15818                        let sh = e.uninit(n_embd)?;
15819                        // One alloc keeps the ownership contract; the copy is 16KB on-stream.
15820                        let mut sh = sh;
15821                        {
15822                            let mut dst = sh.slice_mut(0..n_embd);
15823                            e.stream().memcpy_dtod(&sh_buf.slice(0..n_embd), &mut dst)?;
15824                        }
15825                        sh
15826                    }
15827                } else {
15828                    e.matmul(down_shexp, act, t)?
15829                }
15830            } else if fused {
15831                let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
15832                let pair = match e.matmul_pre_dual_noscale(gate_shexp, up_shexp, &zq, &zd, 1)? {
15833                    Some((gate, up)) => Some((gate, up)),
15834                    None => {
15835                        match (
15836                            e.matmul_pre_noscale(gate_shexp, &zq, &zd, 1)?,
15837                            e.matmul_pre_noscale(up_shexp, &zq, &zd, 1)?,
15838                        ) {
15839                            (Some(gate), Some(up)) => Some((gate, up)),
15840                            _ => None,
15841                        }
15842                    }
15843                };
15844                match pair {
15845                    Some(((gate, gs), (up, us))) => {
15846                        if e.uses_q8_1_fast(down_shexp) {
15847                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff_sh)?;
15848                            e.matmul_pre(down_shexp, &aq, &ad, &gate, 1)?
15849                        } else {
15850                            let mut act = e.uninit(n_ff_sh)?;
15851                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff_sh)?;
15852                            e.matmul(down_shexp, &act, 1)?
15853                        }
15854                    }
15855                    None => {
15856                        let gate = e.matmul_pre(gate_shexp, &zq, &zd, z, 1)?;
15857                        let up = e.matmul_pre(up_shexp, &zq, &zd, z, 1)?;
15858                        let mut act = e.uninit(n_ff_sh)?;
15859                        Self::ffn_act(e, cfg, &gate, &up, &mut act, n_ff_sh)?;
15860                        e.matmul(down_shexp, &act, 1)?
15861                    }
15862                }
15863            } else {
15864                let sg_gate = e.matmul(gate_shexp, z, t)?;
15865                let sg_up = e.matmul(up_shexp, z, t)?;
15866                let mut sa = e.uninit(t * n_ff_sh)?;
15867                Self::ffn_act_lim(
15868                    e,
15869                    cfg,
15870                    &sg_gate,
15871                    &sg_up,
15872                    1.0,
15873                    1.0,
15874                    lim,
15875                    &mut sa,
15876                    t * n_ff_sh,
15877                )?;
15878                e.matmul(down_shexp, &sa, t)?
15879            };
15880            let gate = match &m.gate_inp_shexp {
15881                Some(gate_inp_shexp) => {
15882                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
15883                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
15884                    } else {
15885                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
15886                        let mut gate = e.uninit(t)?;
15887                        e.sigmoid(&raw, &mut gate, t)?;
15888                        gate
15889                    }
15890                }
15891                // t=1 hot path: the per-layer htod of a ones row is a PAGEABLE H2D that
15892                // synchronizes the stream — measured as the biggest per-layer host gap
15893                // (44.6us x 42, eager gap table 2026-08-21). One persistent ones row per
15894                // device serves every layer; larger t (prefill) keeps the plain htod.
15895                None if t == 1 => {
15896                    static ONES: std::sync::Mutex<
15897                        Option<std::collections::HashMap<usize, CudaSlice<f32>>>,
15898                    > = std::sync::Mutex::new(None);
15899                    let mut guard = ONES.lock().map_err(|_| "shexp ones lock is poisoned")?;
15900                    let device = e.ctx().ordinal();
15901                    let rows = guard.get_or_insert_with(Default::default);
15902                    // One entry lookup, not three (contains_key + insert + get). The vacant arm
15903                    // stays fallible, which is why this is `match` and not `or_insert_with`.
15904                    use std::collections::hash_map::Entry;
15905                    let ones = match rows.entry(device) {
15906                        Entry::Occupied(occupied) => occupied.into_mut(),
15907                        Entry::Vacant(vacant) => vacant.insert(e.htod(&[1.0f32])?),
15908                    };
15909                    e.add_scaled_rows(&sh, ones, moe_out, n_embd, t)?;
15910                    return Ok(());
15911                }
15912                None => e.htod(&vec![1.0f32; t])?,
15913            };
15914            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
15915        }
15916        Ok(())
15917    }
15918
15919    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
15920    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
15921    pub(crate) fn moe_ffn_grouped(
15922        e: &Engine,
15923        m: &MoeWeights,
15924        z: &CudaSlice<f32>,
15925        t: usize,
15926        cfg: &ModelConfig,
15927        il: u16,
15928        max_block: usize,
15929    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15930        let moe = cfg.moe.as_ref().unwrap();
15931        let n_embd = cfg.n_embd as usize;
15932        let n_expert = moe.expert_count as usize;
15933        let n_used = moe.expert_used_count as usize;
15934        let n_ff_exp = moe.expert_ff_length as usize;
15935        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
15936        let lim_exp = cfg.clamp_exp_at(il as u32);
15937
15938        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
15939        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
15940        // enters the softmax-only pairs/dev router.
15941        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
15942        if let Some(sig) = cfg.sigmoid_router() {
15943            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
15944        }
15945        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
15946            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
15947        } else {
15948            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
15949        };
15950        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
15951        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
15952        Self::trace_moe_input(e, il, t, n_embd, z)?;
15953
15954        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
15955        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
15956        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
15957        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
15958        let no_exp_macros = m.gate_exps.macros.is_none()
15959            && m.up_exps.macros.is_none()
15960            && m.down_exps.macros.is_none();
15961        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
15962            m.has_uniform_expert_layout()
15963                && no_exp_macros
15964                && moe_q8_enabled_for_model(cfg, m)
15965                && moe_slab_enabled()
15966                && dev.dev == e.ctx().ordinal()
15967        });
15968        if let Some(dev) = resident_q8 {
15969            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
15970                e,
15971                m,
15972                z,
15973                t,
15974                cfg,
15975                il,
15976                &sel_all,
15977                &w_all,
15978                &dev.ptr_row,
15979                dev.gu_il,
15980            )?;
15981            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
15982            return Ok(moe_out);
15983        }
15984
15985        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
15986        // For each expert e, we need: which tokens use it, their positions in z, their top-k
15987        // slot index (for bit-identical accumulation), and their weights.
15988        struct ExpertGroup {
15989            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
15990            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
15991            weights: Vec<f32>,      // renormalized weight for that token-expert pair
15992        }
15993        let mut groups: Vec<ExpertGroup> = (0..n_expert)
15994            .map(|_| ExpertGroup {
15995                tok_indices: Vec::new(),
15996                slot_indices: Vec::new(),
15997                weights: Vec::new(),
15998            })
15999            .collect();
16000
16001        for tok in 0..t {
16002            for j in 0..n_used {
16003                let ex = sel_all[tok * n_used + j] as usize;
16004                let w = w_all[tok * n_used + j];
16005                groups[ex].tok_indices.push(tok as i32);
16006                groups[ex].slot_indices.push(j as i32);
16007                groups[ex].weights.push(w);
16008            }
16009        }
16010
16011        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
16012        // Each token's 8 expert contributions land in their respective slots.
16013        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
16014        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
16015
16016        // Expert weight dimensions (used in both cache and staging paths).
16017        let g_len = m.gate_exps.max_expert_bytes();
16018        let u_len = m.up_exps.max_expert_bytes();
16019        let d_len = m.down_exps.max_expert_bytes();
16020        let moe_q8 = moe_q8_enabled_for_model(cfg, m);
16021        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
16022        // Interleaved GU slabs require the pointer-table fast path above.
16023        let slab_local = m
16024            .dev_exps
16025            .as_ref()
16026            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
16027        let use_cache =
16028            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
16029        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
16030        // also does: a local resident slab or a live SLRU dispatch.
16031        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
16032
16033        // GPU scratch for staging (only allocated without a local slab or cache).
16034        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
16035            (
16036                Some(e.alloc_u8(g_len)?),
16037                Some(e.alloc_u8(u_len)?),
16038                Some(e.alloc_u8(d_len)?),
16039            )
16040        } else {
16041            (None, None, None)
16042        };
16043
16044        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
16045        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
16046        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
16047        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
16048        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
16049        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
16050        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
16051        // at long prompts where every expert stages regardless. Order is FREE to change without
16052        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
16053        // regardless of expert processing order (the whole point of the slots).
16054        let mut order: Vec<usize> = (0..n_expert)
16055            .filter(|&ex| !groups[ex].tok_indices.is_empty())
16056            .collect();
16057        order.sort_by(|&a, &b| {
16058            groups[b]
16059                .tok_indices
16060                .len()
16061                .cmp(&groups[a].tok_indices.len())
16062                .then(a.cmp(&b))
16063        });
16064        let mut m_dist: Vec<usize> = Vec::new(); // for stats
16065        let page_window = moe_page_prefetch_window();
16066        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
16067        if worker_disk_prefetch
16068            && let Some(first) = grouped_worker_prefetch_position(order.len(), None)
16069        {
16070            Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
16071        }
16072        for (order_pos, &ex) in order.iter().enumerate() {
16073            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
16074                Self::moe_prefetch_host_expert(order[next], m);
16075            }
16076            if worker_disk_prefetch
16077                && let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos))
16078            {
16079                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16080                let keep = [
16081                    BlockId::new(il, PROJ_GATE, ex as u16),
16082                    BlockId::new(il, PROJ_UP, ex as u16),
16083                    BlockId::new(il, PROJ_DOWN, ex as u16),
16084                ];
16085                Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
16086            }
16087            let grp = &groups[ex];
16088            let m_e = grp.tok_indices.len();
16089            m_dist.push(m_e);
16090            let gl = m.gate_exps.expert_layout(ex);
16091            let ul = m.up_exps.expert_layout(ex);
16092            let dl = m.down_exps.expert_layout(ex);
16093
16094            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
16095            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
16096            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
16097            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
16098            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
16099            let dmac = m.down_exps.macro_scale(ex);
16100            let weight_d = if dmac == 1.0 {
16101                e.htod(&grp.weights)?
16102            } else {
16103                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
16104                e.htod(&scaled)?
16105            };
16106
16107            // GATHER: collect m_e activation rows from z into a contiguous buffer.
16108            let mut gathered = e.zeros(m_e * n_embd)?;
16109            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
16110            let gv = gathered.slice(0..m_e * n_embd);
16111
16112            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
16113            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
16114            let y = if let Some(dev) = slab_local {
16115                let gate_start = ex * m.gate_exps.expert_stride;
16116                let up_start = ex * m.up_exps.expert_stride;
16117                let down_start = ex * m.down_exps.expert_stride;
16118                if grouped_q8 {
16119                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16120                    let gate = e.qmatvec_expert_q8(
16121                        &dev.gate,
16122                        gate_start..gate_start + gl.len,
16123                        &zq,
16124                        &zd,
16125                        m_e,
16126                        m.gate_exps.in_f,
16127                        m.gate_exps.out_f,
16128                        gl.qtype,
16129                        gl.row_bytes,
16130                    )?;
16131                    let up = e.qmatvec_expert_q8(
16132                        &dev.up,
16133                        up_start..up_start + ul.len,
16134                        &zq,
16135                        &zd,
16136                        m_e,
16137                        m.up_exps.in_f,
16138                        m.up_exps.out_f,
16139                        ul.qtype,
16140                        ul.row_bytes,
16141                    )?;
16142                    let mut act = e.uninit(m_e * n_ff_exp)?;
16143                    Self::ffn_act_lim(
16144                        e,
16145                        cfg,
16146                        &gate,
16147                        &up,
16148                        m.gate_exps.macro_scale(ex),
16149                        m.up_exps.macro_scale(ex),
16150                        lim_exp,
16151                        &mut act,
16152                        m_e * n_ff_exp,
16153                    )?;
16154                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16155                    e.qmatvec_expert_q8(
16156                        &dev.down,
16157                        down_start..down_start + dl.len,
16158                        &aq2,
16159                        &ad2,
16160                        m_e,
16161                        m.down_exps.in_f,
16162                        m.down_exps.out_f,
16163                        dl.qtype,
16164                        dl.row_bytes,
16165                    )?
16166                } else {
16167                    let gate = m.qmatvec_view(
16168                        e,
16169                        &dev.gate,
16170                        gate_start..gate_start + gl.len,
16171                        &gv,
16172                        m_e,
16173                        m.gate_exps.in_f,
16174                        m.gate_exps.out_f,
16175                        gl.qtype,
16176                        gl.row_bytes,
16177                    )?;
16178                    let up = m.qmatvec_view(
16179                        e,
16180                        &dev.up,
16181                        up_start..up_start + ul.len,
16182                        &gv,
16183                        m_e,
16184                        m.up_exps.in_f,
16185                        m.up_exps.out_f,
16186                        ul.qtype,
16187                        ul.row_bytes,
16188                    )?;
16189                    let mut act = e.uninit(m_e * n_ff_exp)?;
16190                    Self::ffn_act_lim(
16191                        e,
16192                        cfg,
16193                        &gate,
16194                        &up,
16195                        m.gate_exps.macro_scale(ex),
16196                        m.up_exps.macro_scale(ex),
16197                        lim_exp,
16198                        &mut act,
16199                        m_e * n_ff_exp,
16200                    )?;
16201                    let actv = act.slice(0..m_e * n_ff_exp);
16202                    m.qmatvec_view(
16203                        e,
16204                        &dev.down,
16205                        down_start..down_start + dl.len,
16206                        &actv,
16207                        m_e,
16208                        m.down_exps.in_f,
16209                        m.down_exps.out_f,
16210                        dl.qtype,
16211                        dl.row_bytes,
16212                    )?
16213                }
16214            } else if use_cache {
16215                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16216                if grouped_q8 {
16217                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16218                    let gate = e.with_moe_cache(max_block, |cache, eng| {
16219                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
16220                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16221                        eng.qmatvec_expert_q8(
16222                            cache.buf(slot),
16223                            0..gl.len,
16224                            &zq,
16225                            &zd,
16226                            m_e,
16227                            m.gate_exps.in_f,
16228                            m.gate_exps.out_f,
16229                            gl.qtype,
16230                            gl.row_bytes,
16231                        )
16232                    })?;
16233                    let up = e.with_moe_cache(max_block, |cache, eng| {
16234                        let id = BlockId::new(il, PROJ_UP, ex as u16);
16235                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16236                        eng.qmatvec_expert_q8(
16237                            cache.buf(slot),
16238                            0..ul.len,
16239                            &zq,
16240                            &zd,
16241                            m_e,
16242                            m.up_exps.in_f,
16243                            m.up_exps.out_f,
16244                            ul.qtype,
16245                            ul.row_bytes,
16246                        )
16247                    })?;
16248                    let mut act = e.uninit(m_e * n_ff_exp)?;
16249                    Self::ffn_act_lim(
16250                        e,
16251                        cfg,
16252                        &gate,
16253                        &up,
16254                        m.gate_exps.macro_scale(ex),
16255                        m.up_exps.macro_scale(ex),
16256                        lim_exp,
16257                        &mut act,
16258                        m_e * n_ff_exp,
16259                    )?;
16260                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16261                    e.with_moe_cache(max_block, |cache, eng| {
16262                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16263                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16264                        eng.qmatvec_expert_q8(
16265                            cache.buf(slot),
16266                            0..dl.len,
16267                            &aq2,
16268                            &ad2,
16269                            m_e,
16270                            m.down_exps.in_f,
16271                            m.down_exps.out_f,
16272                            dl.qtype,
16273                            dl.row_bytes,
16274                        )
16275                    })?
16276                } else {
16277                    let gate = e.with_moe_cache(max_block, |cache, eng| {
16278                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
16279                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
16280                        m.qmatvec_view(
16281                            eng,
16282                            cache.buf(slot),
16283                            0..gl.len,
16284                            &gv,
16285                            m_e,
16286                            m.gate_exps.in_f,
16287                            m.gate_exps.out_f,
16288                            gl.qtype,
16289                            gl.row_bytes,
16290                        )
16291                    })?;
16292                    let up = e.with_moe_cache(max_block, |cache, eng| {
16293                        let id = BlockId::new(il, PROJ_UP, ex as u16);
16294                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
16295                        m.qmatvec_view(
16296                            eng,
16297                            cache.buf(slot),
16298                            0..ul.len,
16299                            &gv,
16300                            m_e,
16301                            m.up_exps.in_f,
16302                            m.up_exps.out_f,
16303                            ul.qtype,
16304                            ul.row_bytes,
16305                        )
16306                    })?;
16307                    let mut act = e.uninit(m_e * n_ff_exp)?;
16308                    Self::ffn_act_lim(
16309                        e,
16310                        cfg,
16311                        &gate,
16312                        &up,
16313                        m.gate_exps.macro_scale(ex),
16314                        m.up_exps.macro_scale(ex),
16315                        lim_exp,
16316                        &mut act,
16317                        m_e * n_ff_exp,
16318                    )?;
16319                    let actv = act.slice(0..m_e * n_ff_exp);
16320                    e.with_moe_cache(max_block, |cache, eng| {
16321                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
16322                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
16323                        m.qmatvec_view(
16324                            eng,
16325                            cache.buf(slot),
16326                            0..dl.len,
16327                            &actv,
16328                            m_e,
16329                            m.down_exps.in_f,
16330                            m.down_exps.out_f,
16331                            dl.qtype,
16332                            dl.row_bytes,
16333                        )
16334                    })?
16335                }
16336            } else {
16337                let sg = scratch_g.as_mut().unwrap();
16338                let su = scratch_u.as_mut().unwrap();
16339                let sd = scratch_d.as_mut().unwrap();
16340                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
16341                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
16342                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
16343                if grouped_q8 {
16344                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
16345                    let gate = e.qmatvec_expert_q8(
16346                        sg,
16347                        0..gl.len,
16348                        &zq,
16349                        &zd,
16350                        m_e,
16351                        m.gate_exps.in_f,
16352                        m.gate_exps.out_f,
16353                        gl.qtype,
16354                        gl.row_bytes,
16355                    )?;
16356                    let up = e.qmatvec_expert_q8(
16357                        su,
16358                        0..ul.len,
16359                        &zq,
16360                        &zd,
16361                        m_e,
16362                        m.up_exps.in_f,
16363                        m.up_exps.out_f,
16364                        ul.qtype,
16365                        ul.row_bytes,
16366                    )?;
16367                    let mut act = e.uninit(m_e * n_ff_exp)?;
16368                    Self::ffn_act_lim(
16369                        e,
16370                        cfg,
16371                        &gate,
16372                        &up,
16373                        m.gate_exps.macro_scale(ex),
16374                        m.up_exps.macro_scale(ex),
16375                        lim_exp,
16376                        &mut act,
16377                        m_e * n_ff_exp,
16378                    )?;
16379                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
16380                    e.qmatvec_expert_q8(
16381                        sd,
16382                        0..dl.len,
16383                        &aq2,
16384                        &ad2,
16385                        m_e,
16386                        m.down_exps.in_f,
16387                        m.down_exps.out_f,
16388                        dl.qtype,
16389                        dl.row_bytes,
16390                    )?
16391                } else {
16392                    let gate = m.qmatvec_view(
16393                        e,
16394                        sg,
16395                        0..gl.len,
16396                        &gv,
16397                        m_e,
16398                        m.gate_exps.in_f,
16399                        m.gate_exps.out_f,
16400                        gl.qtype,
16401                        gl.row_bytes,
16402                    )?;
16403                    let up = m.qmatvec_view(
16404                        e,
16405                        su,
16406                        0..ul.len,
16407                        &gv,
16408                        m_e,
16409                        m.up_exps.in_f,
16410                        m.up_exps.out_f,
16411                        ul.qtype,
16412                        ul.row_bytes,
16413                    )?;
16414                    let mut act = e.uninit(m_e * n_ff_exp)?;
16415                    Self::ffn_act_lim(
16416                        e,
16417                        cfg,
16418                        &gate,
16419                        &up,
16420                        m.gate_exps.macro_scale(ex),
16421                        m.up_exps.macro_scale(ex),
16422                        lim_exp,
16423                        &mut act,
16424                        m_e * n_ff_exp,
16425                    )?;
16426                    let actv = act.slice(0..m_e * n_ff_exp);
16427                    m.qmatvec_view(
16428                        e,
16429                        sd,
16430                        0..dl.len,
16431                        &actv,
16432                        m_e,
16433                        m.down_exps.in_f,
16434                        m.down_exps.out_f,
16435                        dl.qtype,
16436                        dl.row_bytes,
16437                    )?
16438                }
16439            };
16440
16441            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
16442            e.scatter_slot(
16443                &y,
16444                &tok_idx_d,
16445                &slot_idx_d,
16446                &weight_d,
16447                &mut slot_buf,
16448                &mut wbuf,
16449                n_embd,
16450                n_used,
16451                m_e,
16452            )?;
16453        }
16454
16455        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
16456        let mut moe_out = e.zeros(t * n_embd)?;
16457        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
16458
16459        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
16460        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
16461            m_dist.sort_unstable();
16462            let active = m_dist.len();
16463            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
16464            let median = m_dist[active / 2];
16465            let max_m = *m_dist.last().unwrap();
16466            let min_m = m_dist[0];
16467            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
16468            println!(
16469                "moe-grouped il={il} t={t} active={active}/{n_expert} \
16470                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
16471                      above_gemm_threshold(>=16)={above16}/{active}"
16472            );
16473        }
16474
16475        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
16476        Ok(moe_out)
16477    }
16478
16479    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
16480    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
16481    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
16482    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
16483    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
16484    /// expert-sum order identical to the sequential path.
16485    pub(crate) fn moe_ffn_lockstep(
16486        &self,
16487        e: &Engine,
16488        m: &MoeWeights,
16489        zbatch: &CudaSlice<f32>,
16490        mrows: usize,
16491        il: u16,
16492        max_block: usize,
16493    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16494        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
16495        let cfg = &self.cfg;
16496        let moe = cfg.moe.as_ref().unwrap();
16497        let n_embd = cfg.n_embd as usize;
16498        let n_expert = moe.expert_count as usize;
16499        let n_used = moe.expert_used_count as usize;
16500        let n_ff_exp = moe.expert_ff_length as usize;
16501        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
16502        let lim_exp = cfg.clamp_exp_at(il as u32);
16503        let lim_shexp = cfg.clamp_shexp_at(il as u32);
16504
16505        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
16506        if let Some(sig) = cfg.sigmoid_router() {
16507            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
16508        }
16509        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
16510            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
16511        } else {
16512            Self::moe_route_cfg(
16513                e,
16514                &logits,
16515                mrows,
16516                n_expert,
16517                n_used,
16518                m.active_experts.as_deref(),
16519            )?
16520        };
16521        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
16522
16523        // Residency split at whole-expert granularity against the (frozen) cache.
16524        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
16525            Ok((0..n_expert)
16526                .map(|ex| {
16527                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
16528                        .into_iter()
16529                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
16530                })
16531                .collect())
16532        })?;
16533
16534        struct Group {
16535            rows: Vec<i32>,
16536            slots: Vec<i32>,
16537            weights: Vec<f32>,
16538        }
16539        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
16540        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
16541        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
16542            Default::default();
16543        for row in 0..mrows {
16544            for j in 0..n_used {
16545                let ex = sel_all[row * n_used + j] as usize;
16546                let w = w_all[row * n_used + j];
16547                if resident_expert[ex] {
16548                    let group = groups.entry(ex).or_insert_with(|| Group {
16549                        rows: Vec::new(),
16550                        slots: Vec::new(),
16551                        weights: Vec::new(),
16552                    });
16553                    group.rows.push(row as i32);
16554                    group.slots.push(j as i32);
16555                    group.weights.push(w);
16556                } else {
16557                    crate::cpu_experts::record_incomplete_gpu_residency(0);
16558                    cpu_rows[row].push((ex, w));
16559                    cpu_by_expert.entry(ex).or_default().push((row, w));
16560                }
16561            }
16562        }
16563
16564        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
16565        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
16566        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
16567        // order per row differs from the sequential single-call chunk — part of the
16568        // documented lockstep numeric class.
16569        let host_rows = e.dtoh(zbatch)?;
16570        let rows_ok = crate::cpu_experts::rows_supported();
16571        enum CpuPart {
16572            Single { row: usize },
16573            Rows { rows: Vec<usize> },
16574        }
16575        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
16576        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
16577        if rows_ok {
16578            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
16579                .into_iter()
16580                .filter(|(_, rows)| rows.len() >= 2)
16581                .collect();
16582            shared.sort_by_key(|(ex, _)| *ex);
16583            for (ex, mut row_weights) in shared {
16584                row_weights.sort_by_key(|(row, _)| *row);
16585                let inputs: Vec<(&[f32], f32)> = row_weights
16586                    .iter()
16587                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
16588                    .collect();
16589                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
16590                    .map_err(std::io::Error::other)?;
16591                for &(row, _) in &row_weights {
16592                    rows_served.insert((row, ex));
16593                }
16594                tickets.push((
16595                    CpuPart::Rows {
16596                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
16597                    },
16598                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
16599                ));
16600            }
16601        }
16602        for (row, selected) in cpu_rows.iter().enumerate() {
16603            let leftover: Vec<(usize, f32)> = selected
16604                .iter()
16605                .copied()
16606                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
16607                .collect();
16608            if leftover.is_empty() {
16609                continue;
16610            }
16611            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
16612            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
16613                .map_err(std::io::Error::other)?;
16614            tickets.push((
16615                CpuPart::Single { row },
16616                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
16617            ));
16618        }
16619
16620        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
16621        let mut wbuf = e.zeros(mrows * n_used)?;
16622        let mut order: Vec<usize> = groups.keys().copied().collect();
16623        order.sort_by(|&a, &b| {
16624            groups[&b]
16625                .rows
16626                .len()
16627                .cmp(&groups[&a].rows.len())
16628                .then(a.cmp(&b))
16629        });
16630        for &ex in &order {
16631            let group = &groups[&ex];
16632            let m_e = group.rows.len();
16633            let gl = m.gate_exps.expert_layout(ex);
16634            let ul = m.up_exps.expert_layout(ex);
16635            let dl = m.down_exps.expert_layout(ex);
16636            let row_idx_d = e.htod_i32(&group.rows)?;
16637            let slot_idx_d = e.htod_i32(&group.slots)?;
16638            let dmac = m.down_exps.macro_scale(ex);
16639            let weight_d = if dmac == 1.0 {
16640                e.htod(&group.weights)?
16641            } else {
16642                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
16643                e.htod(&scaled)?
16644            };
16645            let mut gathered = e.zeros(m_e * n_embd)?;
16646            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
16647            let gv = gathered.slice(0..m_e * n_embd);
16648            let gate = e.with_moe_cache(max_block, |c, eng| {
16649                let slot = c
16650                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
16651                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16652                m.qmatvec_view(
16653                    eng,
16654                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16655                    0..gl.len,
16656                    &gv,
16657                    m_e,
16658                    m.gate_exps.in_f,
16659                    m.gate_exps.out_f,
16660                    gl.qtype,
16661                    gl.row_bytes,
16662                )
16663            })?;
16664            let up = e.with_moe_cache(max_block, |c, eng| {
16665                let slot = c
16666                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
16667                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16668                m.qmatvec_view(
16669                    eng,
16670                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16671                    0..ul.len,
16672                    &gv,
16673                    m_e,
16674                    m.up_exps.in_f,
16675                    m.up_exps.out_f,
16676                    ul.qtype,
16677                    ul.row_bytes,
16678                )
16679            })?;
16680            let mut act = e.zeros(m_e * n_ff_exp)?;
16681            Self::ffn_act_lim(
16682                e,
16683                cfg,
16684                &gate,
16685                &up,
16686                m.gate_exps.macro_scale(ex),
16687                m.up_exps.macro_scale(ex),
16688                lim_exp,
16689                &mut act,
16690                m_e * n_ff_exp,
16691            )?;
16692            let actv = act.slice(0..m_e * n_ff_exp);
16693            let y = e.with_moe_cache(max_block, |c, eng| {
16694                let slot = c
16695                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
16696                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
16697                m.qmatvec_view(
16698                    eng,
16699                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
16700                    0..dl.len,
16701                    &actv,
16702                    m_e,
16703                    m.down_exps.in_f,
16704                    m.down_exps.out_f,
16705                    dl.qtype,
16706                    dl.row_bytes,
16707                )
16708            })?;
16709            e.scatter_slot(
16710                &y,
16711                &row_idx_d,
16712                &slot_idx_d,
16713                &weight_d,
16714                &mut slot_buf,
16715                &mut wbuf,
16716                n_embd,
16717                n_used,
16718                m_e,
16719            )?;
16720        }
16721        let mut moe_out = e.zeros(mrows * n_embd)?;
16722        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
16723
16724        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
16725        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
16726        for (part, ticket) in tickets {
16727            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
16728            let mut add_row = |row: usize, chunk: &[f32]| {
16729                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
16730                for (accumulator, value) in sum.iter_mut().zip(chunk) {
16731                    *accumulator += value;
16732                }
16733            };
16734            match part {
16735                CpuPart::Single { row } => add_row(row, &cpu_output),
16736                CpuPart::Rows { rows } => {
16737                    for (slot, row) in rows.into_iter().enumerate() {
16738                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
16739                    }
16740                }
16741            }
16742        }
16743        for (row, sum) in row_sums.into_iter().enumerate() {
16744            let Some(sum) = sum else { continue };
16745            let cpu_output = e.htod(&sum)?;
16746            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
16747            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
16748        }
16749
16750        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
16751            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
16752        {
16753            let n_ff_sh = gate_shexp.out_features();
16754            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
16755            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
16756            let mut sa = e.zeros(mrows * n_ff_sh)?;
16757            Self::ffn_act_lim(
16758                e,
16759                cfg,
16760                &sg_gate,
16761                &sg_up,
16762                1.0,
16763                1.0,
16764                lim_shexp,
16765                &mut sa,
16766                mrows * n_ff_sh,
16767            )?;
16768            let sh = e.matmul(down_shexp, &sa, mrows)?;
16769            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
16770            // decode matches the single-sequence decode chain bit-for-bit.
16771            let g = match &m.gate_inp_shexp {
16772                Some(gate_inp_shexp) => {
16773                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
16774                }
16775                None => e.htod(&vec![1.0f32; mrows])?,
16776            };
16777            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
16778        }
16779
16780        Ok(moe_out)
16781    }
16782}
16783
16784// ============================ gemma4 (R8 verified wiring) ==================================
16785// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
16786// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
16787// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
16788// gemma variants after the correctness gate).
16789impl HybridModel {
16790    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
16791    /// Rotary width the gemma-4 config DECLARES for layer `il`, per SWA/global class — the
16792    /// `n_rot` argument of the fused `rms_norm_qkv_rope*` kernels
16793    /// (`Engine::full_width_rope_only`, lane/graph-s-key-exactness-20260819).
16794    ///
16795    /// It exists so the fusion's full-width assumption is checked against the CONFIG instead of
16796    /// against itself: passing `hd` twice would guard nothing. On every gemma-4 artifact both
16797    /// routes set `rope_dims_{global,swa} == key_length_{global,swa}` (GGUF reads
16798    /// `rope.dimension_count(_swa)`, the safetensors route derives them from
16799    /// `global_head_dim`/`head_dim`), which is exactly WHY the fusion is legal here — gemma's
16800    /// "partial" rotary is per-dimension frequency scaling through the `rope_freqs` factors
16801    /// tensor, not an n_rot truncation. An arch that truncates instead (qwen3.5: 64 of 256;
16802    /// step35 full-attn: 64 of 128) must not reach these kernels, and now cannot silently.
16803    pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
16804        let g = self
16805            .cfg
16806            .gemma4
16807            .as_ref()
16808            .expect("gemma4_rope_dims on a non-gemma4 config");
16809        if g.swa_pattern[il] {
16810            g.rope_dims_swa as usize
16811        } else {
16812            g.rope_dims_global as usize
16813        }
16814    }
16815
16816    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
16817        let g = self.cfg.gemma4.as_ref().unwrap();
16818        let swa = g.swa_pattern[il];
16819        let hd = if swa {
16820            g.key_length_swa
16821        } else {
16822            g.key_length_global
16823        } as usize;
16824        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
16825        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
16826        // rows exact (softmax over one element) while every later position drifted).
16827        (
16828            hd,
16829            g.head_count_kv[il] as usize,
16830            self.cfg.n_head as usize,
16831            if swa {
16832                g.rope_base_swa
16833            } else {
16834                g.rope_base_global
16835            },
16836            1.0,
16837            swa,
16838        )
16839    }
16840
16841    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
16842    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
16843    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
16844    pub(crate) fn gemma4_suppress(
16845        &self,
16846        e: &Engine,
16847        ld: &mut CudaSlice<f32>,
16848        t: usize,
16849    ) -> Result<(), Box<dyn std::error::Error>> {
16850        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
16851            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
16852            // stage as primary, and this tail runs only after the last stage). The assert turns
16853            // that argued invariant into a checked one: any topology violating primary==head
16854            // trips here in debug instead of silently peer-reading a device-0 buffer.
16855            #[cfg(debug_assertions)]
16856            crate::debug_assert_tensor_stream_device(
16857                ids,
16858                &e.stream(),
16859                "gemma4_suppress.suppress_d",
16860            );
16861            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
16862        }
16863        Ok(())
16864    }
16865
16866    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
16867    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
16868    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
16869    /// only (v0): attends within `tokens` via the f32 sdpa.
16870    #[allow(clippy::too_many_arguments)]
16871    /// ONE-PROGRAM seam (lane/gemma-restore-exactness-20260819, DEFAULT OFF): route dense
16872    /// gemma SWA prefill attention through the WINDOWED arm at every prompt length instead of
16873    /// switching program at `t > sliding_window`. The door is the measured cause of the
16874    /// prefix-restore exactness violation: below it the SWA layers attend unwindowed on f32
16875    /// operands, above it windowed on bf16 operands, so rows[0..932) of a 932-row prime and of
16876    /// a 1048-row prime differ in 59 of 60 layers. `=1` collapses it to one program so the
16877    /// published prefix KV stops depending on the total prompt length. Off by default because
16878    /// collapsing changes the bytes short prompts get today; the flip is an owner call.
16879    fn gemma_fa_one_program() -> bool {
16880        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16881        *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
16882    }
16883
16884    #[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
16885    fn gemma4_attn_prime(
16886        &self,
16887        e: &Engine,
16888        fa: &crate::hybrid::FullAttnLayer,
16889        il: usize,
16890        h: &CudaSlice<f32>,
16891        pos_d: &CudaSlice<i32>,
16892        t: usize,
16893        cache: Option<&mut Cache>,
16894        island: Option<&CudaSlice<i32>>,
16895    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16896        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
16897        let eps = self.cfg.rms_eps;
16898        let aux = self.gemma4_aux.as_ref().unwrap();
16899        let ones = aux.ones(e);
16900        #[cfg(debug_assertions)]
16901        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
16902
16903        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
16904        // (h stays borrowed across the triple, so the cache key can't go stale).
16905        e.mmq_act_begin();
16906        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
16907        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
16908            let v = e.dtoh(&q0)?;
16909            let nan = v.iter().filter(|x| x.is_nan()).count();
16910            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
16911            eprintln!(
16912                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
16913                v.len()
16914            );
16915        }
16916        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
16917        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
16918        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
16919        let v0 = if swa {
16920            e.matmul(&fa.wv, h, t)?
16921        } else {
16922            e.clone_dtod(&k0)?
16923        };
16924        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
16925            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
16926                let v = e.dtoh(buf)?;
16927                let nan = v.iter().filter(|x| x.is_nan()).count();
16928                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
16929                eprintln!(
16930                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
16931                    v.len()
16932                );
16933            }
16934        }
16935
16936        let mut q = e.uninit(t * nh * hd)?;
16937        let mut k = e.uninit(t * nkv * hd)?;
16938        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
16939        let mut v = e.uninit(t * nkv * hd)?;
16940        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
16941        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
16942        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
16943        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16944        // Island primes take the mask-capable naive kernel below; keep the operands f32
16945        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
16946        let emit = island.is_none()
16947            && t >= 16
16948            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
16949            && *EMIT.get_or_init(|| {
16950                std::env::var("MEMRA_FA_EMIT")
16951                    .map(|s| s != "0")
16952                    .unwrap_or(true)
16953            });
16954        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
16955        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
16956        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
16957        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
16958        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
16959        let v_f16 = emit
16960            && crate::fa_f16pv_on()
16961            && match hd {
16962                512 => true,
16963                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
16964                _ => false,
16965            };
16966        if emit {
16967            e.rms_norm_qkv_w4b(
16968                &q0,
16969                &k0,
16970                &v0,
16971                fa.q_norm.float_data(),
16972                fa.k_norm.float_data(),
16973                ones,
16974                &mut q,
16975                &mut k,
16976                &mut v,
16977                &mut vb,
16978                hd,
16979                nh * t,
16980                nkv * t,
16981                eps,
16982                v_f16,
16983            )?;
16984        } else {
16985            e.rms_norm_qkv(
16986                &q0,
16987                &k0,
16988                &v0,
16989                fa.q_norm.float_data(),
16990                fa.k_norm.float_data(),
16991                ones,
16992                &mut q,
16993                &mut k,
16994                &mut v,
16995                hd,
16996                nh * t,
16997                nkv * t,
16998                eps,
16999            )?;
17000        }
17001
17002        let ff = if swa {
17003            None
17004        } else {
17005            Some(
17006                aux.rope_freqs(e)
17007                    .expect("gemma4 global rope needs rope_freqs.weight"),
17008            )
17009        };
17010        #[cfg(debug_assertions)]
17011        if let Some(ff) = ff {
17012            crate::debug_assert_tensor_stream_device(
17013                ff,
17014                &e.stream(),
17015                "gemma4_attn_prime.rope_freqs",
17016            );
17017        }
17018        if emit {
17019            e.rope_neox2_bf16e(
17020                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
17021            )?;
17022        } else {
17023            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
17024        }
17025
17026        if let Some(cache) = cache {
17027            let kvl = cache.kv[il].as_mut().unwrap();
17028            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
17029            e.append_kv_quantized_rows(
17030                &k,
17031                &v,
17032                &mut kvl.k,
17033                &mut kvl.v,
17034                kvl.len,
17035                t,
17036                kvl.kv_dim_k,
17037                kvl.kv_dim_v,
17038                kvl.k_tok_bytes,
17039                kvl.v_tok_bytes,
17040                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
17041            )?;
17042            kvl.len += t;
17043        }
17044        let mut attn = e.zeros(t * nh * hd)?;
17045        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
17046        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
17047        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
17048        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
17049        if let Some(span) = island {
17050            // Masked-prefill arm: every layer routes through the island-aware naive
17051            // kernel (correctness-first, same posture as the vision tower v1). The
17052            // window argument keeps the R6 shortcut: 0 while the prompt fits the
17053            // window, the real window beyond it.
17054            let w = if swa && t > win { win } else { 0 };
17055            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
17056        } else if swa && (t > win || Self::gemma_fa_one_program()) {
17057            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17058                if emit {
17059                    e.fa_prefill_w_pre(
17060                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
17061                    )?;
17062                } else {
17063                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17064                }
17065            } else {
17066                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
17067            }
17068        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
17069            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17070        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
17071            if emit {
17072                e.fa_prefill_hd512_pre(
17073                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
17074                )?;
17075            } else {
17076                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17077            }
17078        } else {
17079            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
17080        }
17081        e.matmul(&fa.wo, &attn, t)
17082    }
17083
17084    /// Back-compat wrapper (pure prefill, no cache).
17085    fn gemma4_attn(
17086        &self,
17087        e: &Engine,
17088        fa: &crate::hybrid::FullAttnLayer,
17089        il: usize,
17090        h: &CudaSlice<f32>,
17091        pos_d: &CudaSlice<i32>,
17092        t: usize,
17093    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17094        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
17095    }
17096
17097    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
17098    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
17099    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
17100    /// the q8z epilogue is quantize_q8_1 verbatim).
17101    fn gemma4_moe_q8(
17102        &self,
17103        e: &Engine,
17104        m: &crate::hybrid::MoeWeights,
17105        bits: &crate::hybrid::Gemma4MoeBits,
17106        mq: &(CudaSlice<i8>, CudaSlice<f32>),
17107        router_in: &CudaSlice<f32>,
17108        t: usize,
17109    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17110        let cfg = &self.cfg;
17111        let moe = cfg.moe.as_ref().unwrap();
17112        let n_embd = cfg.n_embd as usize;
17113        let n_expert = moe.expert_count as usize;
17114        let n_used = moe.expert_used_count as usize;
17115        let n_ff_exp = moe.expert_ff_length as usize;
17116        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
17117        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
17118        // the pair's 12us is kernel time, not launch gaps.
17119        let logits = if crate::router_kernel_on() {
17120            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17121        } else {
17122            e.matmul(&m.gate_inp, router_in, t)?
17123        };
17124        let dev = m.dev_exps.as_ref().unwrap();
17125        let (sel_d, w_d) =
17126            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17127        let (zq, zd) = mq;
17128        if t == 1 {
17129            let selv = sel_d.slice(0..n_used);
17130            let wv = w_d.slice(0..n_used);
17131            let act = e.moe_gate_up_gelu8_dev_q8(
17132                &dev.ptr_row,
17133                &selv,
17134                zq,
17135                zd,
17136                n_embd,
17137                n_ff_exp,
17138                n_used,
17139                n_expert,
17140                m.gate_exps.qtype,
17141                m.up_exps.qtype,
17142                m.gate_exps.row_bytes,
17143                m.up_exps.row_bytes,
17144            )?;
17145            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17146            let mut moe_out = e.uninit(n_embd)?;
17147            e.moe_down8_fma_dev_q8(
17148                &dev.ptr_row,
17149                &selv,
17150                &wv,
17151                &aq2,
17152                &ad2,
17153                &mut moe_out.slice_mut(0..n_embd),
17154                n_ff_exp,
17155                n_embd,
17156                n_used,
17157                n_expert,
17158                m.down_exps.qtype,
17159                m.down_exps.row_bytes,
17160            )?;
17161            return Ok(moe_out);
17162        }
17163        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17164        let act = if csr {
17165            e.moe_gate_up_gelu8_dev_q8_csr(
17166                &dev.ptr_row,
17167                &sel_d,
17168                zq,
17169                zd,
17170                t * n_used,
17171                n_embd,
17172                n_ff_exp,
17173                n_used,
17174                n_expert,
17175                m.gate_exps.qtype,
17176                m.up_exps.qtype,
17177                m.gate_exps.row_bytes,
17178                m.up_exps.row_bytes,
17179            )?
17180        } else {
17181            e.moe_gate_up_gelu8_dev_q8_rows(
17182                &dev.ptr_row,
17183                &sel_d,
17184                zq,
17185                zd,
17186                t,
17187                n_embd,
17188                n_ff_exp,
17189                n_used,
17190                n_expert,
17191                m.gate_exps.qtype,
17192                m.up_exps.qtype,
17193                m.gate_exps.row_bytes,
17194                m.up_exps.row_bytes,
17195            )?
17196        };
17197        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17198        let mut moe_out = e.uninit(t * n_embd)?;
17199        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
17200        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
17201        e.moe_down8_fma_dev_q8_rows_g(
17202            &dev.ptr_row,
17203            &sel_d,
17204            &w_d,
17205            &aq2,
17206            &ad2,
17207            &mut moe_out,
17208            t,
17209            n_ff_exp,
17210            n_embd,
17211            n_used,
17212            n_expert,
17213            m.down_exps.qtype,
17214            m.down_exps.row_bytes,
17215        )?;
17216        Ok(moe_out)
17217    }
17218
17219    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
17220    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
17221    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
17222    fn gemma4_moe(
17223        &self,
17224        e: &Engine,
17225        m: &crate::hybrid::MoeWeights,
17226        bits: &crate::hybrid::Gemma4MoeBits,
17227        moe_in: &CudaSlice<f32>,
17228        router_in: &CudaSlice<f32>,
17229        t: usize,
17230    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17231        let cfg = &self.cfg;
17232        let moe = cfg.moe.as_ref().unwrap();
17233        let n_embd = cfg.n_embd as usize;
17234        let n_expert = moe.expert_count as usize;
17235        let n_used = moe.expert_used_count as usize;
17236        let n_ff_exp = moe.expert_ff_length as usize;
17237
17238        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
17239        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
17240        // batched matmul only at real prefill.
17241        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
17242            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
17243        } else {
17244            e.matmul(&m.gate_inp, router_in, t)?
17245        };
17246
17247        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
17248        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
17249        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
17250        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
17251        if t < PRIME_MIN_T
17252            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17253            && expert_dp4a_supported(m.gate_exps.qtype)
17254            && expert_dp4a_supported(m.up_exps.qtype)
17255            && expert_dp4a_supported(m.down_exps.qtype)
17256            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
17257        {
17258            let dev = m.dev_exps.as_ref().unwrap();
17259            let (sel_d, w_d) =
17260                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
17261            if t == 1 {
17262                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
17263                let selv = sel_d.slice(0..n_used);
17264                let wv = w_d.slice(0..n_used);
17265                let act = e.moe_gate_up_gelu8_dev_q8(
17266                    &dev.ptr_row,
17267                    &selv,
17268                    &zq,
17269                    &zd,
17270                    n_embd,
17271                    n_ff_exp,
17272                    n_used,
17273                    n_expert,
17274                    m.gate_exps.qtype,
17275                    m.up_exps.qtype,
17276                    m.gate_exps.row_bytes,
17277                    m.up_exps.row_bytes,
17278                )?;
17279                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
17280                let mut moe_out = e.uninit(n_embd)?;
17281                e.moe_down8_fma_dev_q8(
17282                    &dev.ptr_row,
17283                    &selv,
17284                    &wv,
17285                    &aq2,
17286                    &ad2,
17287                    &mut moe_out.slice_mut(0..n_embd),
17288                    n_ff_exp,
17289                    n_embd,
17290                    n_used,
17291                    n_expert,
17292                    m.down_exps.qtype,
17293                    m.down_exps.row_bytes,
17294                )?;
17295                return Ok(moe_out);
17296            }
17297            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
17298            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
17299            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
17300            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
17301            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17302            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
17303            let act = if csr {
17304                e.moe_gate_up_gelu8_dev_q8_csr(
17305                    &dev.ptr_row,
17306                    &sel_d,
17307                    &zq,
17308                    &zd,
17309                    t * n_used,
17310                    n_embd,
17311                    n_ff_exp,
17312                    n_used,
17313                    n_expert,
17314                    m.gate_exps.qtype,
17315                    m.up_exps.qtype,
17316                    m.gate_exps.row_bytes,
17317                    m.up_exps.row_bytes,
17318                )?
17319            } else {
17320                e.moe_gate_up_gelu8_dev_q8_rows(
17321                    &dev.ptr_row,
17322                    &sel_d,
17323                    &zq,
17324                    &zd,
17325                    t,
17326                    n_embd,
17327                    n_ff_exp,
17328                    n_used,
17329                    n_expert,
17330                    m.gate_exps.qtype,
17331                    m.up_exps.qtype,
17332                    m.gate_exps.row_bytes,
17333                    m.up_exps.row_bytes,
17334                )?
17335            };
17336            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
17337            let mut moe_out = e.uninit(t * n_embd)?;
17338            e.moe_down8_fma_dev_q8_rows_g(
17339                &dev.ptr_row,
17340                &sel_d,
17341                &w_d,
17342                &aq2,
17343                &ad2,
17344                &mut moe_out,
17345                t,
17346                n_ff_exp,
17347                n_embd,
17348                n_used,
17349                n_expert,
17350                m.down_exps.qtype,
17351                m.down_exps.row_bytes,
17352            )?;
17353            return Ok(moe_out);
17354        }
17355
17356        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
17357        for (i, &sx) in sel_all.iter().enumerate() {
17358            w_all[i] *= bits.per_expert_scale[sx as usize];
17359        }
17360
17361        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
17362        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
17363        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
17364        if t >= PRIME_MIN_T
17365            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17366            && expert_dp4a_supported(m.gate_exps.qtype)
17367            && expert_dp4a_supported(m.up_exps.qtype)
17368            && expert_dp4a_supported(m.down_exps.qtype)
17369            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
17370        {
17371            let dev = m.dev_exps.as_ref().unwrap();
17372            let n_pairs = t * n_used;
17373            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
17374            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
17375            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
17376            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
17377            let pt = e.htod_i32(&pair_tok)?;
17378            let pw = e.htod(&w_all)?;
17379            let toff = e.htod_i32(&tok_off)?;
17380            let tids = e.htod_i32(&tok_ids)?;
17381            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
17382            for p in 0..n_pairs {
17383                by_ex[pair_ex[p] as usize].push(p as i32);
17384            }
17385            let mut ex_ids: Vec<i32> = Vec::new();
17386            let mut ex_off: Vec<i32> = vec![0];
17387            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
17388            for (ex, list) in by_ex.iter().enumerate() {
17389                if list.is_empty() {
17390                    continue;
17391                }
17392                ex_ids.push(ex as i32);
17393                ex_pairs.extend_from_slice(list);
17394                ex_off.push(ex_pairs.len() as i32);
17395            }
17396            let n_active = ex_ids.len();
17397            let exi = e.htod_i32(&ex_ids)?;
17398            let exo = e.htod_i32(&ex_off)?;
17399            let exp_d = e.htod_i32(&ex_pairs)?;
17400            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
17401            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
17402            // end-to-end (gelu is elementwise), one row permute before the scatter. The
17403            // ragged down k (704) needs no padding here — cublas takes any k.
17404            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
17405            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
17406            // Hopper default — see moe_f16g_gemma_on.
17407            if crate::moe_f16g_gemma_on()
17408                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
17409                && f16g_proj_ok(m.up_exps.qtype, n_embd)
17410                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
17411            {
17412                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
17413                let csr_tok_d = e.htod_i32(&csr_tok)?;
17414                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
17415                let g_csr = e.moe_f16_grouped(
17416                    &dev.ptr_row,
17417                    0,
17418                    n_expert,
17419                    &exi,
17420                    &ex_off,
17421                    &exo,
17422                    &z_f16,
17423                    &z_s,
17424                    n_embd,
17425                    n_ff_exp,
17426                    n_active,
17427                    n_pairs,
17428                    m.gate_exps.qtype,
17429                    m.gate_exps.row_bytes,
17430                )?;
17431                let u_csr = e.moe_f16_grouped(
17432                    &dev.ptr_row,
17433                    1,
17434                    n_expert,
17435                    &exi,
17436                    &ex_off,
17437                    &exo,
17438                    &z_f16,
17439                    &z_s,
17440                    n_embd,
17441                    n_ff_exp,
17442                    n_active,
17443                    n_pairs,
17444                    m.up_exps.qtype,
17445                    m.up_exps.row_bytes,
17446                )?;
17447                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
17448                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
17449                let d_csr = e.moe_f16_grouped(
17450                    &dev.ptr_row,
17451                    2,
17452                    n_expert,
17453                    &exi,
17454                    &ex_off,
17455                    &exo,
17456                    &a_f16,
17457                    &a_s,
17458                    n_ff_exp,
17459                    n_embd,
17460                    n_active,
17461                    n_pairs,
17462                    m.down_exps.qtype,
17463                    m.down_exps.row_bytes,
17464                )?;
17465                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
17466                let mut moe_out = e.uninit(t * n_embd)?;
17467                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17468                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
17469                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
17470                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
17471                    eprintln!(
17472                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
17473                        scan(&yd),
17474                        scan(&mo)
17475                    );
17476                }
17477                return Ok(moe_out);
17478            }
17479            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
17480            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
17481            let mma = n_embd.is_multiple_of(256)
17482                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
17483            let (gate, up) = if mma {
17484                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
17485                (
17486                    e.mmq_iq_experts(
17487                        &dev.ptr_row,
17488                        0,
17489                        n_expert,
17490                        &exi,
17491                        &exo,
17492                        &exp_d,
17493                        &pt,
17494                        &z_scr,
17495                        n_embd,
17496                        n_ff_exp,
17497                        n_active,
17498                        n_pairs,
17499                        t,
17500                        m.gate_exps.qtype,
17501                        m.gate_exps.row_bytes,
17502                    )?,
17503                    e.mmq_iq_experts(
17504                        &dev.ptr_row,
17505                        1,
17506                        n_expert,
17507                        &exi,
17508                        &exo,
17509                        &exp_d,
17510                        &pt,
17511                        &z_scr,
17512                        n_embd,
17513                        n_ff_exp,
17514                        n_active,
17515                        n_pairs,
17516                        t,
17517                        m.up_exps.qtype,
17518                        m.up_exps.row_bytes,
17519                    )?,
17520                )
17521            } else {
17522                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
17523                (
17524                    e.moe_pairs_matvec_q8_dec(
17525                        &dev.ptr_row,
17526                        0,
17527                        &exi,
17528                        &exo,
17529                        &exp_d,
17530                        &pt,
17531                        &zq,
17532                        &zd,
17533                        n_embd,
17534                        n_ff_exp,
17535                        n_expert,
17536                        n_active,
17537                        n_pairs,
17538                        m.gate_exps.qtype,
17539                        m.gate_exps.row_bytes,
17540                    )?,
17541                    e.moe_pairs_matvec_q8_dec(
17542                        &dev.ptr_row,
17543                        1,
17544                        &exi,
17545                        &exo,
17546                        &exp_d,
17547                        &pt,
17548                        &zq,
17549                        &zd,
17550                        n_embd,
17551                        n_ff_exp,
17552                        n_expert,
17553                        n_active,
17554                        n_pairs,
17555                        m.up_exps.qtype,
17556                        m.up_exps.row_bytes,
17557                    )?,
17558                )
17559            };
17560            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
17561            let pself = e.htod_i32(&pair_self)?;
17562            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
17563            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
17564            // to the 256-val superblock (768) while the act quantizer's zero padding
17565            // makes every padded-k product exactly zero (weight overread bytes multiply
17566            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
17567            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
17568            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
17569            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
17570            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
17571            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
17572            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
17573            let y_down = if mma {
17574                let in_pad = n_ff_exp.div_ceil(256) * 256;
17575                let a_scr = if crate::moe_fuse_actq_on() {
17576                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
17577                } else {
17578                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17579                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
17580                };
17581                e.mmq_iq_experts(
17582                    &dev.ptr_row,
17583                    2,
17584                    n_expert,
17585                    &exi,
17586                    &exo,
17587                    &exp_d,
17588                    &pself,
17589                    &a_scr,
17590                    in_pad,
17591                    n_embd,
17592                    n_active,
17593                    n_pairs,
17594                    n_pairs,
17595                    m.down_exps.qtype,
17596                    m.down_exps.row_bytes,
17597                )?
17598            } else {
17599                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
17600                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
17601                e.moe_pairs_matvec_q8_dec(
17602                    &dev.ptr_row,
17603                    2,
17604                    &exi,
17605                    &exo,
17606                    &exp_d,
17607                    &pself,
17608                    &aq2,
17609                    &ad2,
17610                    n_ff_exp,
17611                    n_embd,
17612                    n_expert,
17613                    n_active,
17614                    n_pairs,
17615                    m.down_exps.qtype,
17616                    m.down_exps.row_bytes,
17617                )?
17618            };
17619            let mut moe_out = e.uninit(t * n_embd)?;
17620            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
17621            return Ok(moe_out);
17622        }
17623
17624        let g_len = m.gate_exps.expert_stride;
17625        let u_len = m.up_exps.expert_stride;
17626        let d_len = m.down_exps.expert_stride;
17627        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
17628        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
17629        // the spill fallback.
17630        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
17631        let (mut sg, mut su, mut sd) = if dev.is_some() {
17632            (None, None, None)
17633        } else {
17634            (
17635                Some(e.alloc_u8_uninit(g_len)?),
17636                Some(e.alloc_u8_uninit(u_len)?),
17637                Some(e.alloc_u8_uninit(d_len)?),
17638            )
17639        };
17640        let mut moe_out = e.zeros(t * n_embd)?;
17641        for tok in 0..t {
17642            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
17643            let w = &w_all[tok * n_used..(tok + 1) * n_used];
17644            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
17645            for (j, &ex) in sel.iter().enumerate() {
17646                let ex = ex as usize;
17647                let gate = match dev {
17648                    Some(d) => m.qmatvec_view(
17649                        e,
17650                        &d.gate,
17651                        ex * g_len..(ex + 1) * g_len,
17652                        &zt,
17653                        1,
17654                        m.gate_exps.in_f,
17655                        m.gate_exps.out_f,
17656                        m.gate_exps.qtype,
17657                        m.gate_exps.row_bytes,
17658                    )?,
17659                    None => {
17660                        let sg = sg.as_mut().unwrap();
17661                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
17662                        m.qmatvec_view(
17663                            e,
17664                            sg,
17665                            0..g_len,
17666                            &zt,
17667                            1,
17668                            m.gate_exps.in_f,
17669                            m.gate_exps.out_f,
17670                            m.gate_exps.qtype,
17671                            m.gate_exps.row_bytes,
17672                        )?
17673                    }
17674                };
17675                let up = match dev {
17676                    Some(d) => m.qmatvec_view(
17677                        e,
17678                        &d.up,
17679                        ex * u_len..(ex + 1) * u_len,
17680                        &zt,
17681                        1,
17682                        m.up_exps.in_f,
17683                        m.up_exps.out_f,
17684                        m.up_exps.qtype,
17685                        m.up_exps.row_bytes,
17686                    )?,
17687                    None => {
17688                        let su = su.as_mut().unwrap();
17689                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
17690                        m.qmatvec_view(
17691                            e,
17692                            su,
17693                            0..u_len,
17694                            &zt,
17695                            1,
17696                            m.up_exps.in_f,
17697                            m.up_exps.out_f,
17698                            m.up_exps.qtype,
17699                            m.up_exps.row_bytes,
17700                        )?
17701                    }
17702                };
17703                let mut act = e.uninit(n_ff_exp)?;
17704                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
17705                let actv = act.slice(0..n_ff_exp);
17706                let y = match dev {
17707                    Some(d) => m.qmatvec_view(
17708                        e,
17709                        &d.down,
17710                        ex * d_len..(ex + 1) * d_len,
17711                        &actv,
17712                        1,
17713                        m.down_exps.in_f,
17714                        m.down_exps.out_f,
17715                        m.down_exps.qtype,
17716                        m.down_exps.row_bytes,
17717                    )?,
17718                    None => {
17719                        let sd = sd.as_mut().unwrap();
17720                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
17721                        m.qmatvec_view(
17722                            e,
17723                            sd,
17724                            0..d_len,
17725                            &actv,
17726                            1,
17727                            m.down_exps.in_f,
17728                            m.down_exps.out_f,
17729                            m.down_exps.qtype,
17730                            m.down_exps.row_bytes,
17731                        )?
17732                    }
17733                };
17734                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
17735                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
17736            }
17737        }
17738        Ok(moe_out)
17739    }
17740
17741    /// One gemma4 trunk layer (R8): x -> x_next.
17742    fn gemma4_layer(
17743        &self,
17744        e: &Engine,
17745        il: usize,
17746        layer: &crate::hybrid::HybridLayer,
17747        x: &CudaSlice<f32>,
17748        pos_d: &CudaSlice<i32>,
17749        t: usize,
17750    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17751        let n_embd = self.cfg.n_embd as usize;
17752        let eps = self.cfg.rms_eps;
17753
17754        let mut h = e.zeros(t * n_embd)?;
17755        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
17756        let Mixer::Full(fa) = &layer.mixer else {
17757            panic!("gemma4 layer {il} not full-attn")
17758        };
17759        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
17760        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
17761        let mut cur = e.zeros(t * n_embd)?;
17762        e.rms_norm(
17763            &o,
17764            layer.post_attn_norm.float_data(),
17765            &mut cur,
17766            n_embd,
17767            t,
17768            eps,
17769        )?;
17770        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
17771    }
17772
17773    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
17774    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
17775    /// layer scale — shared verbatim by the prefill, decode and verify paths.
17776    fn gemma4_layer_tail_add(
17777        &self,
17778        e: &Engine,
17779        layer: &crate::hybrid::HybridLayer,
17780        cur: &CudaSlice<f32>,
17781        x: &CudaSlice<f32>,
17782        t: usize,
17783    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17784        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
17785    }
17786
17787    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
17788    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
17789    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
17790    fn gemma4_layer_tail_add_n(
17791        &self,
17792        e: &Engine,
17793        layer: &crate::hybrid::HybridLayer,
17794        cur: &CudaSlice<f32>,
17795        x: &CudaSlice<f32>,
17796        t: usize,
17797        next_norm: Option<&CudaSlice<f32>>,
17798    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
17799        let n_embd = self.cfg.n_embd as usize;
17800        let bits = layer.gemma4.as_ref().unwrap();
17801        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
17802        let mut xn = e.uninit(t * n_embd)?;
17803        match next_norm {
17804            Some(w) => {
17805                let mut hn = e.uninit(t * n_embd)?;
17806                e.add_scale_rms_norm(
17807                    &sn,
17808                    &attn_out,
17809                    bits.layer_scale,
17810                    w,
17811                    &mut xn,
17812                    &mut hn,
17813                    n_embd,
17814                    t,
17815                    self.cfg.rms_eps,
17816                )?;
17817                Ok((xn, Some(hn)))
17818            }
17819            None => {
17820                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
17821                Ok((xn, None))
17822            }
17823        }
17824    }
17825
17826    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
17827    /// norm — returns (sn, attn_out) for the closing add+scale variants.
17828    fn gemma4_layer_tail_core(
17829        &self,
17830        e: &Engine,
17831        layer: &crate::hybrid::HybridLayer,
17832        cur: &CudaSlice<f32>,
17833        x: &CudaSlice<f32>,
17834        t: usize,
17835    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17836        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
17837    }
17838
17839    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
17840    /// means `cur` is the RAW attention output and the dense entry runs
17841    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
17842    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
17843    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
17844    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
17845    #[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
17846    fn gemma4_layer_tail_core_pn(
17847        &self,
17848        e: &Engine,
17849        layer: &crate::hybrid::HybridLayer,
17850        cur: &CudaSlice<f32>,
17851        x: &CudaSlice<f32>,
17852        t: usize,
17853        pre_norm: Option<&CudaSlice<f32>>,
17854        defer_post_norm: bool,
17855    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17856        let n_embd = self.cfg.n_embd as usize;
17857        let eps = self.cfg.rms_eps;
17858        let bits = layer.gemma4.as_ref().unwrap();
17859
17860        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
17861        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
17862        let Some(mbits) = bits.moe_bits.as_ref() else {
17863            let crate::hybrid::Ffn::Dense {
17864                ffn_gate,
17865                ffn_up,
17866                ffn_down,
17867            } = &layer.ffn
17868            else {
17869                panic!("gemma4 dense layer without Dense ffn")
17870            };
17871            let mut attn_out = e.uninit(t * n_embd)?;
17872            let mut zsh = e.uninit(t * n_embd)?;
17873            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
17874            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
17875            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
17876            match pre_norm {
17877                Some(wa) if t == 1 => {
17878                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
17879                        cur,
17880                        wa,
17881                        x,
17882                        bits.ffn_norm.float_data(),
17883                        &mut attn_out,
17884                        &mut zsh,
17885                        n_embd,
17886                        t,
17887                        eps,
17888                    )?);
17889                }
17890                Some(wa) => e.rms_pre_add_rms_norm(
17891                    cur,
17892                    wa,
17893                    x,
17894                    bits.ffn_norm.float_data(),
17895                    &mut attn_out,
17896                    &mut zsh,
17897                    n_embd,
17898                    t,
17899                    eps,
17900                )?,
17901                None => e.add_rms_norm(
17902                    cur,
17903                    x,
17904                    bits.ffn_norm.float_data(),
17905                    &mut attn_out,
17906                    &mut zsh,
17907                    n_embd,
17908                    t,
17909                    eps,
17910                )?,
17911            }
17912            let n_ff = ffn_gate.out_features();
17913            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
17914            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
17915            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
17916            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
17917            // rescue segment C — the megakernel front is closed for the dense tail.
17918            let (gate, up) = if t == 1 {
17919                let (zq, zd) = match zpair {
17920                    Some(p) => p,
17921                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
17922                };
17923                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
17924                    Some(p) => p,
17925                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
17926                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
17927                        Some(p) => p,
17928                        None => (
17929                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
17930                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
17931                        ),
17932                    },
17933                }
17934            } else {
17935                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
17936                // launch for the verify's gate+up — the up segment's blocks fill SMs as
17937                // the gate segment drains (the launch-tail mechanism behind the b-tier
17938                // plateau; first positive after six falsified in-kernel variants).
17939                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17940                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
17941                let fused = if f2b {
17942                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
17943                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
17944                } else {
17945                    None
17946                };
17947                match fused {
17948                    Some(p) => p,
17949                    None => {
17950                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
17951                        e.mmq_act_begin();
17952                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
17953                    }
17954                }
17955            };
17956            let mut act = e.uninit(t * n_ff)?;
17957            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
17958            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
17959            let f0 = if e.uses_q8_1_fast(ffn_down) {
17960                let upv = e.view(&up, t * n_ff);
17961                let up_all = upv.slice(0..t * n_ff);
17962                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
17963                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
17964            } else {
17965                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
17966                e.matmul(ffn_down, &act, t)?
17967            };
17968            if defer_post_norm {
17969                return Ok((f0, attn_out));
17970            }
17971            let mut sn = e.uninit(t * n_embd)?;
17972            e.rms_norm(
17973                &f0,
17974                bits.post_ffw_norm.float_data(),
17975                &mut sn,
17976                n_embd,
17977                t,
17978                eps,
17979            )?;
17980            return Ok((sn, attn_out));
17981        };
17982
17983        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
17984        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
17985        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
17986        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
17987        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
17988        let mut attn_out = e.uninit(t * n_embd)?;
17989        let mut router_in = e.uninit(t * n_embd)?;
17990        let fast_moe = match &layer.ffn {
17991            crate::hybrid::Ffn::Moe(m) => {
17992                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
17993                    && expert_dp4a_supported(m.gate_exps.qtype)
17994                    && expert_dp4a_supported(m.up_exps.qtype)
17995                    && expert_dp4a_supported(m.down_exps.qtype)
17996                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
17997            }
17998            _ => false,
17999        };
18000        let q8z = t < PRIME_MIN_T && fast_moe;
18001        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
18002            let (z0, m2) = e.add_rms_norm3_q8z(
18003                cur,
18004                x,
18005                bits.ffn_norm.float_data(),
18006                &mbits.router_scale_pre,
18007                mbits.pre_ffw_norm_2.float_data(),
18008                &mut attn_out,
18009                &mut router_in,
18010                n_embd,
18011                t,
18012                eps,
18013            )?;
18014            (None, Some(z0), Some(m2))
18015        } else {
18016            let mut zsh = e.uninit(t * n_embd)?;
18017            let mut moe_in = e.uninit(t * n_embd)?;
18018            e.add_rms_norm3(
18019                cur,
18020                x,
18021                bits.ffn_norm.float_data(),
18022                &mbits.router_scale_pre,
18023                mbits.pre_ffw_norm_2.float_data(),
18024                &mut attn_out,
18025                &mut zsh,
18026                &mut router_in,
18027                &mut moe_in,
18028                n_embd,
18029                t,
18030                eps,
18031            )?;
18032            (Some((zsh, moe_in)), None, None)
18033        };
18034        let attn_out2 = attn_out;
18035        #[allow(unused_variables)]
18036        let attn_out = &attn_out2;
18037        let n_ff = mbits.shared_gate.out_features();
18038        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
18039            if t == 1 {
18040                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
18041                    Some(p) => p,
18042                    None => match e.matmul_nvfp4_fused2(
18043                        &mbits.shared_gate,
18044                        &mbits.shared_up,
18045                        zq,
18046                        zd,
18047                        1,
18048                    )? {
18049                        Some(p) => p,
18050                        None => {
18051                            let h0 = e.zeros(0)?;
18052                            (
18053                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
18054                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
18055                            )
18056                        }
18057                    },
18058                }
18059            } else {
18060                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
18061                let h0 = e.zeros(0)?;
18062                (
18063                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
18064                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
18065                )
18066            }
18067        } else {
18068            let (zsh, _) = zsh_f32.as_ref().unwrap();
18069            (
18070                e.matmul(&mbits.shared_gate, zsh, t)?,
18071                e.matmul(&mbits.shared_up, zsh, t)?,
18072            )
18073        };
18074        let mut act = e.uninit(t * n_ff)?;
18075        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
18076        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
18077        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
18078            panic!("gemma4 layer not MoE")
18079        };
18080        let moe0 = match (&moe_q8, &zsh_f32) {
18081            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
18082            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
18083            _ => unreachable!(),
18084        };
18085        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
18086        let mut mlp = e.uninit(t * n_embd)?;
18087        let mut moe = e.uninit(t * n_embd)?;
18088        e.rms_norm2x(
18089            &mlp0,
18090            &moe0,
18091            mbits.post_ffw_norm_1.float_data(),
18092            mbits.post_ffw_norm_2.float_data(),
18093            &mut mlp,
18094            &mut moe,
18095            n_embd,
18096            t,
18097            eps,
18098        )?;
18099
18100        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
18101        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
18102        let mut sum = e.uninit(t * n_embd)?;
18103        let mut sn = e.uninit(t * n_embd)?;
18104        e.add_rms_norm(
18105            &mlp,
18106            &moe,
18107            bits.post_ffw_norm.float_data(),
18108            &mut sum,
18109            &mut sn,
18110            n_embd,
18111            t,
18112            eps,
18113        )?;
18114        Ok((sn, attn_out2))
18115    }
18116
18117    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
18118    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
18119    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
18120    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
18121    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
18122    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
18123    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
18124    /// decode == verify == graph parity holds by construction at either seam value.
18125    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
18126    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18127    pub(crate) fn gemma4_layer_tail_add_nq_pn(
18128        &self,
18129        e: &Engine,
18130        layer: &crate::hybrid::HybridLayer,
18131        o: &CudaSlice<f32>,
18132        x: &CudaSlice<f32>,
18133        t: usize,
18134        next_norm: Option<&CudaSlice<f32>>,
18135    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18136    {
18137        let n_embd = self.cfg.n_embd as usize;
18138        let eps = self.cfg.rms_eps;
18139        let bits = layer.gemma4.as_ref().unwrap();
18140        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
18141            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
18142                e,
18143                layer,
18144                o,
18145                x,
18146                t,
18147                Some(layer.post_attn_norm.float_data()),
18148                true,
18149            )?;
18150            let mut xn = e.uninit(t * n_embd)?;
18151            return match next_norm {
18152                Some(w) => {
18153                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
18154                        &f0,
18155                        bits.post_ffw_norm.float_data(),
18156                        &attn_out,
18157                        bits.layer_scale,
18158                        w,
18159                        &mut xn,
18160                        n_embd,
18161                        t,
18162                        eps,
18163                    )?;
18164                    Ok((xn, Some(pair)))
18165                }
18166                None => {
18167                    let mut sn = e.uninit(t * n_embd)?;
18168                    e.rms_norm(
18169                        &f0,
18170                        bits.post_ffw_norm.float_data(),
18171                        &mut sn,
18172                        n_embd,
18173                        t,
18174                        eps,
18175                    )?;
18176                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18177                    Ok((xn, None))
18178                }
18179            };
18180        }
18181        let mut cur = e.uninit(t * n_embd)?;
18182        e.rms_norm(
18183            o,
18184            layer.post_attn_norm.float_data(),
18185            &mut cur,
18186            n_embd,
18187            t,
18188            eps,
18189        )?;
18190        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
18191    }
18192
18193    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18194    pub(crate) fn gemma4_layer_tail_add_nq(
18195        &self,
18196        e: &Engine,
18197        layer: &crate::hybrid::HybridLayer,
18198        cur: &CudaSlice<f32>,
18199        x: &CudaSlice<f32>,
18200        t: usize,
18201        next_norm: Option<&CudaSlice<f32>>,
18202    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
18203    {
18204        let n_embd = self.cfg.n_embd as usize;
18205        let bits = layer.gemma4.as_ref().unwrap();
18206        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
18207        let mut xn = e.uninit(t * n_embd)?;
18208        match next_norm {
18209            Some(w) => {
18210                let pair = e.add_scale_rms_norm_q8_1(
18211                    &sn,
18212                    &attn_out,
18213                    bits.layer_scale,
18214                    w,
18215                    &mut xn,
18216                    n_embd,
18217                    t,
18218                    self.cfg.rms_eps,
18219                )?;
18220                Ok((xn, Some(pair)))
18221            }
18222            None => {
18223                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
18224                Ok((xn, None))
18225            }
18226        }
18227    }
18228
18229    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
18230    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
18231    fn gemma4_forward(
18232        &self,
18233        e: &Engine,
18234        tokens: &[u32],
18235        last_only: bool,
18236    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
18237        // E4B routes to its own forward regardless of the caller's entry point (forward /
18238        // forward_last / prime paths all funnel here for gemma4).
18239        if self.is_gemma4_e4b() {
18240            return self.gemma4_e4b_forward(e, tokens, last_only);
18241        }
18242        let n_embd = self.cfg.n_embd as usize;
18243        let t = tokens.len();
18244        let pos: Vec<i32> = (0..t as i32).collect();
18245        let pos_d = e.htod_i32(&pos)?;
18246
18247        let mut x = self.embed(e, tokens)?;
18248        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18249        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
18250        // the bring-up bisect vs llama-eval-callback node stats.
18251        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
18252        let stat =
18253            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
18254                let h = e.dtoh(x)?;
18255                let bad = h.iter().filter(|v| !v.is_finite()).count();
18256                let mx = h
18257                    .iter()
18258                    .filter(|v| v.is_finite())
18259                    .fold(0.0f32, |m, v| m.max(v.abs()));
18260                eprintln!(
18261                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
18262                    &h[..3]
18263                );
18264                Ok(())
18265            };
18266        if probe {
18267            stat(e, &x, "embed")?;
18268        }
18269        for (il, layer) in self.layers.iter().enumerate() {
18270            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
18271            if probe {
18272                stat(e, &x, &format!("L{il}"))?;
18273            }
18274        }
18275        let mut hn = e.zeros(t * n_embd)?;
18276        e.rms_norm(
18277            &x,
18278            self.output_norm.float_data(),
18279            &mut hn,
18280            n_embd,
18281            t,
18282            self.cfg.rms_eps,
18283        )?;
18284        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18285        let n_vocab = self.output.out_features();
18286        let logits = if last_only {
18287            let hv = e.view(&hn, t * n_embd);
18288            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
18289            let mut hlast = e.zeros(n_embd)?;
18290            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
18291            let mut ld = e.matmul(&self.output, &hlast, 1)?;
18292            e.softcap(&mut ld, cap, n_vocab)?;
18293            self.gemma4_suppress(e, &mut ld, 1)?;
18294            e.dtoh(&ld)?
18295        } else {
18296            let mut ld = e.matmul(&self.output, &hn, t)?;
18297            e.softcap(&mut ld, cap, t * n_vocab)?;
18298            self.gemma4_suppress(e, &mut ld, t)?;
18299            e.dtoh(&ld)?
18300        };
18301        Ok(logits)
18302    }
18303
18304    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
18305    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
18306    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
18307    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
18308    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
18309    pub(crate) fn gemma4_prime(
18310        &self,
18311        e: &Engine,
18312        tokens: &[u32],
18313        cache: &mut Cache,
18314        overlay: Option<&crate::vision::EmbedOverlay>,
18315    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18316        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
18317        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
18318        // whole worker process on this line. The worker now primes gemma4 monolithically and
18319        // routes continuation suffixes tokenwise; this is the per-request backstop.
18320        if cache.pos != 0 {
18321            return Err(
18322                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
18323                        — prime the full prompt in one call or decode tokenwise"
18324                    .into(),
18325            );
18326        }
18327        let n_embd = self.cfg.n_embd as usize;
18328        let eps = self.cfg.rms_eps;
18329        let t = tokens.len();
18330        let pos: Vec<i32> = (0..t as i32).collect();
18331        let pos_d = e.htod_i32(&pos)?;
18332        let mut x = self.embed(e, tokens)?;
18333        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
18334        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
18335        // sqrt(n_embd) text scale — the reference scales token batches only
18336        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
18337        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
18338        // bidirectional within itself, causal+SWA everywhere else, matching the
18339        // reference's llama_set_causal_attn(false) image batch exactly.
18340        let island: Option<CudaSlice<i32>> = match overlay {
18341            Some(ov) => {
18342                let mut span_id = vec![-1i32; t];
18343                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
18344                    if pos + n_rows > t {
18345                        return Err(format!(
18346                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
18347                            pos + n_rows
18348                        )
18349                        .into());
18350                    }
18351                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
18352                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
18353                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
18354                        *s = i as i32;
18355                    }
18356                }
18357                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
18358                // keep the plain causal mask. Exists only so the decisive probe can show
18359                // the island mask itself changes the answer; never on in serving.
18360                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
18361                    None
18362                } else {
18363                    Some(e.htod_i32(&span_id)?)
18364                }
18365            }
18366            None => None,
18367        };
18368        for (il, layer) in self.layers.iter().enumerate() {
18369            let mut h = e.zeros(t * n_embd)?;
18370            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
18371            let Mixer::Full(fa) = &layer.mixer else {
18372                panic!("gemma4 layer not full-attn")
18373            };
18374            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
18375            if trace {
18376                let v = e.dtoh(&h)?;
18377                let nan = v.iter().filter(|x| x.is_nan()).count();
18378                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
18379            }
18380            let o =
18381                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
18382            if trace {
18383                let v = e.dtoh(&o)?;
18384                let nan = v.iter().filter(|x| x.is_nan()).count();
18385                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
18386            }
18387            let mut cur = e.zeros(t * n_embd)?;
18388            e.rms_norm(
18389                &o,
18390                layer.post_attn_norm.float_data(),
18391                &mut cur,
18392                n_embd,
18393                t,
18394                eps,
18395            )?;
18396            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
18397            self.dflash_tap(e, cache, il, &x, t)?;
18398            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
18399            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
18400                let h = e.dtoh(&x)?;
18401                let nan = h.iter().filter(|v| v.is_nan()).count();
18402                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
18403                eprintln!(
18404                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
18405                    h.len()
18406                );
18407                if nan > 0 {
18408                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
18409                }
18410            }
18411        }
18412        cache.pos += t;
18413        let hiddens = e.clone_dtod(&x)?;
18414        let xv = e.view(&x, t * n_embd);
18415        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
18416        let mut h_seed = e.zeros(n_embd)?;
18417        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
18418        let mut hn = e.uninit(n_embd)?;
18419        e.rms_norm(
18420            &h_seed,
18421            self.output_norm.float_data(),
18422            &mut hn,
18423            n_embd,
18424            1,
18425            eps,
18426        )?;
18427        let mut ld = e.matmul(&self.output, &hn, 1)?;
18428        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
18429        e.softcap(&mut ld, cap, self.output.out_features())?;
18430        self.gemma4_suppress(e, &mut ld, 1)?;
18431        let logits = e.dtoh(&ld)?;
18432        Ok((logits, h_seed, hiddens))
18433    }
18434
18435    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
18436    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
18437    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
18438    /// fused norm emits q8 directly — the f32 h never materializes).
18439    #[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
18440    fn gemma4_decode_attn(
18441        &self,
18442        e: &Engine,
18443        fa: &crate::hybrid::FullAttnLayer,
18444        il: usize,
18445        hq: &CudaSlice<i8>,
18446        hdq: &CudaSlice<f32>,
18447        pos_d: &CudaSlice<i32>,
18448        cache: &mut Cache,
18449    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18450        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18451        let eps = self.cfg.rms_eps;
18452        let aux = self.gemma4_aux.as_ref().unwrap();
18453        let ones = aux.ones(e);
18454        #[cfg(debug_assertions)]
18455        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
18456        let (hq, hdq) = (hq, hdq);
18457        let h0 = e.zeros(0)?;
18458        let h = &h0;
18459        let (q0, k0, v0) = if swa {
18460            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
18461                Some(t3) => t3,
18462                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
18463                // match — fuse the uniform (q,k) pair and take v as its own single.
18464                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18465                    Some((q0, k0)) => {
18466                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, h, 1)?;
18467                        (q0, k0, v0)
18468                    }
18469                    None => (
18470                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18471                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18472                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
18473                    ),
18474                },
18475            }
18476        } else {
18477            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
18478                Some(p) => p,
18479                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
18480                    Some(p) => p,
18481                    None => (
18482                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
18483                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
18484                    ),
18485                },
18486            };
18487            let v0 = e.clone_dtod(&k0)?;
18488            (q0, k0, v0)
18489        };
18490        let mut q = e.uninit(nh * hd)?;
18491        let mut k = e.uninit(nkv * hd)?;
18492        let mut v = e.uninit(nkv * hd)?;
18493        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
18494        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
18495        let ff = if swa {
18496            None
18497        } else {
18498            Some(
18499                aux.rope_freqs(e)
18500                    .expect("gemma4 global rope needs rope_freqs.weight"),
18501            )
18502        };
18503        #[cfg(debug_assertions)]
18504        if let Some(ff) = ff {
18505            crate::debug_assert_tensor_stream_device(
18506                ff,
18507                &e.stream(),
18508                "gemma4_decode_attn.rope_freqs",
18509            );
18510        }
18511        let kvl = cache.kv[il].as_mut().unwrap();
18512        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
18513        if crate::Engine::qkv_append_on() {
18514            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
18515            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
18516            // twin of the dc fold — bit-identical bodies, one launch per layer.
18517            e.rms_norm_qkv_rope_append(
18518                &q0,
18519                &k0,
18520                &v0,
18521                fa.q_norm.float_data(),
18522                fa.k_norm.float_data(),
18523                ones,
18524                &mut q,
18525                &mut k,
18526                &mut v,
18527                hd,
18528                self.gemma4_rope_dims(il),
18529                nh,
18530                nkv,
18531                pos_d,
18532                nh,
18533                nkv,
18534                base,
18535                1.0,
18536                ff,
18537                eps,
18538                &mut kvl.k,
18539                &mut kvl.v,
18540                kvl.len,
18541                kvl.k_tok_bytes,
18542                kvl.v_tok_bytes,
18543                kv_fp8,
18544            )?;
18545        } else {
18546            e.rms_norm_qkv_rope(
18547                &q0,
18548                &k0,
18549                &v0,
18550                fa.q_norm.float_data(),
18551                fa.k_norm.float_data(),
18552                ones,
18553                &mut q,
18554                &mut k,
18555                &mut v,
18556                hd,
18557                self.gemma4_rope_dims(il),
18558                nh,
18559                nkv,
18560                pos_d,
18561                nh,
18562                nkv,
18563                base,
18564                1.0,
18565                ff,
18566                eps,
18567            )?;
18568            e.append_kv_quantized(
18569                &k,
18570                &v,
18571                &mut kvl.k,
18572                &mut kvl.v,
18573                kvl.len,
18574                kvl.kv_dim_k,
18575                kvl.kv_dim_v,
18576                kvl.k_tok_bytes,
18577                kvl.v_tok_bytes,
18578                kv_fp8,
18579            )?;
18580        }
18581        kvl.len += 1;
18582        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
18583        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
18584        // positional). Globals attend the full history.
18585        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
18586        let mut attn = e.uninit(nh * hd)?;
18587        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
18588        if !swa
18589            && hd == 512
18590            && kvl.len >= crate::fa512_min_tkv()
18591            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18592        {
18593            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18594            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18595            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
18596            let base = kvl.len as i32;
18597            e.i32_set_k(&mut kvl.len_d, base)?;
18598            e.fa_decode_rows(
18599                &q,
18600                &kp,
18601                &vp,
18602                &mut attn,
18603                hd,
18604                nh,
18605                nkv,
18606                kvl.len - 1,
18607                1,
18608                scale,
18609                kvl.k_tok_bytes,
18610                kvl.v_tok_bytes,
18611                Some((&kvl.len_d, -1)),
18612                false,
18613                false,
18614                None,
18615            )?;
18616            return e.matmul(&fa.wo, &attn, 1);
18617        }
18618        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
18619        if swa
18620            && kvl.len > win
18621            && hd == 256
18622            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
18623        {
18624            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
18625            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
18626            let base = kvl.len as i32;
18627            e.i32_set_k(&mut kvl.len_d, base)?;
18628            e.fa_decode_rows_w(
18629                &q,
18630                &kp,
18631                &vp,
18632                &mut attn,
18633                hd,
18634                nh,
18635                nkv,
18636                &kvl.len_d,
18637                -1,
18638                1,
18639                scale,
18640                win,
18641                kvl.k_tok_bytes,
18642                kvl.v_tok_bytes,
18643                None,
18644            )?;
18645            return e.matmul(&fa.wo, &attn, 1);
18646        }
18647        let (off_tok, t_kv) = if swa && kvl.len > win {
18648            (kvl.len - win, win)
18649        } else {
18650            (0, kvl.len)
18651        };
18652        let k_view = e.view_u8_range(
18653            &kvl.k,
18654            off_tok * kvl.k_tok_bytes,
18655            (off_tok + t_kv) * kvl.k_tok_bytes,
18656        );
18657        let v_view = e.view_u8_range(
18658            &kvl.v,
18659            off_tok * kvl.v_tok_bytes,
18660            (off_tok + t_kv) * kvl.v_tok_bytes,
18661        );
18662        e.fa_decode_kvmod(
18663            &q,
18664            &k_view,
18665            &v_view,
18666            &mut attn,
18667            hd,
18668            nh,
18669            nkv,
18670            t_kv,
18671            scale,
18672            kvl.k_tok_bytes,
18673            kvl.v_tok_bytes,
18674            swa && crate::Engine::wkv_on(),
18675        )?;
18676        e.matmul(&fa.wo, &attn, 1)
18677    }
18678
18679    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
18680    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
18681    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
18682    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
18683    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
18684    /// in-graph; the driver gates).
18685    #[allow(clippy::too_many_arguments)]
18686    pub fn gemma4_decode_step_dc(
18687        &self,
18688        e: &Engine,
18689        token_d: &CudaSlice<u32>,
18690        pos_d: &mut CudaSlice<i32>,
18691        embd_gpu: &CudaSlice<u8>,
18692        embd_qt: i32,
18693        embd_rb: usize,
18694        cache: &mut Cache,
18695        n_vocab: usize,
18696        cap_bucket_max: Option<(usize, usize)>,
18697    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
18698        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
18699        self.gemma4_decode_step_dc_into(
18700            e,
18701            token_d,
18702            pos_d,
18703            embd_gpu,
18704            embd_qt,
18705            embd_rb,
18706            cache,
18707            n_vocab,
18708            cap_bucket_max,
18709            &mut tok_out,
18710        )?;
18711        Ok(tok_out)
18712    }
18713
18714    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
18715    /// every replay; pass `token_d` itself for the self-feeding graph loop).
18716    #[allow(clippy::too_many_arguments)]
18717    pub fn gemma4_decode_step_dc_into(
18718        &self,
18719        e: &Engine,
18720        token_d: &CudaSlice<u32>,
18721        pos_d: &mut CudaSlice<i32>,
18722        embd_gpu: &CudaSlice<u8>,
18723        embd_qt: i32,
18724        embd_rb: usize,
18725        cache: &mut Cache,
18726        n_vocab: usize,
18727        cap_bucket_max: Option<(usize, usize)>,
18728        tok_out: &mut CudaSlice<u32>,
18729    ) -> Result<(), Box<dyn std::error::Error>> {
18730        let n_embd = self.cfg.n_embd as usize;
18731        let eps = self.cfg.rms_eps;
18732        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
18733        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
18734        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
18735        let n_layers = self.layers.len();
18736        for (il, layer) in self.layers.iter().enumerate() {
18737            let (hq, hdq) = match h_carry.take() {
18738                Some(p) => p,
18739                None => {
18740                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
18741                }
18742            };
18743            let Mixer::Full(fa) = &layer.mixer else {
18744                panic!("gemma4 layer {il} not full-attn")
18745            };
18746            let o =
18747                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
18748            let next_norm = if il + 1 < n_layers {
18749                Some(self.layers[il + 1].attn_norm.float_data())
18750            } else {
18751                None
18752            };
18753            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
18754            x = xn;
18755            h_carry = hn;
18756        }
18757        let mut hn = e.uninit(n_embd)?;
18758        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
18759        let mut logits = e.matmul(&self.output, &hn, 1)?;
18760        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
18761        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
18762        e.inc_seqlen(pos_d)?;
18763        if cap_bucket_max.is_none() {
18764            cache.pos += 1;
18765        }
18766        Ok(())
18767    }
18768
18769    // Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
18770    // every buffer the step produces per token lives here, allocated ONCE pre-capture, so
18771    // the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
18772    // osrt 2026-07-23). Sized for the model's max per-layer shapes.
18773
18774    /// Build the slot set (call OUTSIDE any capture).
18775    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
18776        let n_embd = self.cfg.n_embd as usize;
18777        let n_vocab = self.output.out_features();
18778        let n_layers = self.layers.len();
18779        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
18780        for il in 0..n_layers {
18781            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
18782            qmax = qmax.max(nh * hd);
18783            kvmax = kvmax.max(nkv * hd);
18784            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
18785                ffmax = ffmax.max(ffn_gate.out_features());
18786            }
18787        }
18788        Ok(G4DcSlots {
18789            x: e.uninit(n_embd)?,
18790            xn: e.uninit(n_embd)?,
18791            cur: e.uninit(n_embd)?,
18792            hq: e.alloc_i8_uninit(n_embd)?,
18793            hd_: e.uninit(n_embd / 32)?,
18794            q0: e.uninit(qmax)?,
18795            k0: e.uninit(kvmax)?,
18796            v0: e.uninit(kvmax)?,
18797            q: e.uninit(qmax)?,
18798            k: e.uninit(kvmax)?,
18799            v: e.uninit(kvmax)?,
18800            attn: e.uninit(qmax)?,
18801            o: e.uninit(n_embd)?,
18802            attn_out: e.uninit(n_embd)?,
18803            zsh: e.uninit(n_embd)?,
18804            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
18805            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
18806            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
18807            zd: e.uninit(n_embd.max(qmax) / 32)?,
18808            gate: e.uninit(ffmax)?,
18809            up: e.uninit(ffmax)?,
18810            act: e.uninit(ffmax)?,
18811            actq: e.alloc_i8_uninit(ffmax)?,
18812            actd: e.uninit(ffmax / 32)?,
18813            f0: e.uninit(n_embd)?,
18814            sn: e.uninit(n_embd)?,
18815            hn: e.uninit(n_embd)?,
18816            logits: e.uninit(n_vocab)?,
18817        })
18818    }
18819
18820    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
18821    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
18822    fn g4_matvec_m1_into(
18823        &self,
18824        e: &Engine,
18825        w: &crate::model::GpuTensor,
18826        aq: &CudaSlice<i8>,
18827        ad: &CudaSlice<f32>,
18828        y: &mut CudaSlice<f32>,
18829    ) -> Result<(), Box<dyn std::error::Error>> {
18830        use crate::model::GpuTensor;
18831        let (bytes, qtype, row_bytes, scale, rp) = match w {
18832            GpuTensor::Quant {
18833                bytes,
18834                qtype,
18835                row_bytes,
18836                scale,
18837                rp,
18838                ..
18839            } => (bytes, *qtype, *row_bytes, *scale, *rp),
18840            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
18841        };
18842        let (mbytes, mrp) = match w {
18843            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
18844            _ => (bytes, rp),
18845        };
18846        e.qmatvec_mmvq_into(
18847            mbytes,
18848            aq,
18849            ad,
18850            1,
18851            w.in_features(),
18852            w.out_features(),
18853            qtype,
18854            row_bytes,
18855            scale,
18856            mrp,
18857            y,
18858        )
18859    }
18860
18861    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
18862    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
18863    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
18864    #[allow(clippy::too_many_arguments)]
18865    pub fn gemma4_decode_step_dc_slotted(
18866        &self,
18867        e: &Engine,
18868        token_d: &CudaSlice<u32>,
18869        pos_d: &mut CudaSlice<i32>,
18870        embd_gpu: &CudaSlice<u8>,
18871        embd_qt: i32,
18872        embd_rb: usize,
18873        cache: &mut Cache,
18874        n_vocab: usize,
18875        cap_bucket_max: Option<(usize, usize)>,
18876        sl: &mut G4DcSlots,
18877        tok_out: &mut CudaSlice<u32>,
18878        ring: Option<(&mut CudaSlice<u32>, usize)>,
18879    ) -> Result<(), Box<dyn std::error::Error>> {
18880        let n_embd = self.cfg.n_embd as usize;
18881        let eps = self.cfg.rms_eps;
18882        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
18883        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
18884        let n_layers = self.layers.len();
18885        let mut has_carry = false;
18886        for il in 0..n_layers {
18887            if !has_carry {
18888                e.rms_norm_q8_1_into(
18889                    &sl.x,
18890                    self.layers[il].attn_norm.float_data(),
18891                    n_embd,
18892                    1,
18893                    eps,
18894                    &mut sl.hq,
18895                    &mut sl.hd_,
18896                )?;
18897            }
18898            has_carry = true;
18899            let layer = &self.layers[il];
18900            let Mixer::Full(fa) = &layer.mixer else {
18901                panic!("gemma4 layer {il} not full-attn")
18902            };
18903            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
18904            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
18905            // the standalone norm only survives on the unfused seam arm.
18906            if !Engine::g4_pnfold_on() {
18907                e.rms_norm(
18908                    &sl.o,
18909                    layer.post_attn_norm.float_data(),
18910                    &mut sl.cur,
18911                    n_embd,
18912                    1,
18913                    eps,
18914                )?;
18915            }
18916            let next_norm = if il + 1 < n_layers {
18917                Some(self.layers[il + 1].attn_norm.float_data())
18918            } else {
18919                None
18920            };
18921            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
18922            std::mem::swap(&mut sl.x, &mut sl.xn);
18923        }
18924        e.rms_norm(
18925            &sl.x,
18926            self.output_norm.float_data(),
18927            &mut sl.hn,
18928            n_embd,
18929            1,
18930            eps,
18931        )?;
18932        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
18933        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
18934        {
18935            let (zq, zd) = (&sl.zq, &sl.zd);
18936            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
18937            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
18938            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
18939        }
18940        self.gemma4_suppress(e, &mut sl.logits, 1)?;
18941        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
18942        if let Some((ring, base)) = ring {
18943            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
18944            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
18945            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
18946            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
18947        }
18948        e.inc_seqlen(pos_d)?;
18949        if cap_bucket_max.is_none() {
18950            cache.pos += 1;
18951        }
18952        Ok(())
18953    }
18954
18955    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
18956    #[allow(clippy::too_many_arguments)]
18957    fn gemma4_decode_attn_dc_slotted(
18958        &self,
18959        e: &Engine,
18960        fa: &crate::hybrid::FullAttnLayer,
18961        il: usize,
18962        pos_d: &CudaSlice<i32>,
18963        cache: &mut Cache,
18964        cap_bucket_max: Option<(usize, usize)>,
18965        sl: &mut G4DcSlots,
18966    ) -> Result<(), Box<dyn std::error::Error>> {
18967        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
18968        let eps = self.cfg.rms_eps;
18969        let aux = self.gemma4_aux.as_ref().unwrap();
18970        let ones = aux.ones(e);
18971        #[cfg(debug_assertions)]
18972        crate::debug_assert_tensor_stream_device(
18973            ones,
18974            &e.stream(),
18975            "gemma4_decode_attn_dc_slotted.ones",
18976        );
18977        {
18978            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
18979            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
18980            if swa {
18981                if !e.matmul_q4_fused3_into(
18982                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
18983                )? {
18984                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
18985                    // (q,k) pair, v through the generic m1 slot matvec — the same two
18986                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
18987                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18988                    {
18989                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
18990                    } else {
18991                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
18992                    }
18993                }
18994            } else {
18995                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18996                    && !e
18997                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
18998                {
18999                    return Err("slotted step: fused2 unavailable".into());
19000                }
19001                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
19002                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
19003            }
19004        }
19005        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
19006        // kernel-for-kernel (graph stream-identity gate).
19007        let ff = if swa {
19008            None
19009        } else {
19010            Some(
19011                aux.rope_freqs(e)
19012                    .expect("gemma4 global rope needs rope_freqs.weight"),
19013            )
19014        };
19015        #[cfg(debug_assertions)]
19016        if let Some(ff) = ff {
19017            crate::debug_assert_tensor_stream_device(
19018                ff,
19019                &e.stream(),
19020                "gemma4_decode_attn_dc_slotted.rope_freqs",
19021            );
19022        }
19023        let kvl = cache.kv[il].as_mut().unwrap();
19024        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19025        if crate::Engine::qkv_append_on() {
19026            // append fold (2026-07-23): mirrors dc_into.
19027            e.rms_norm_qkv_rope_append_dc(
19028                &sl.q0,
19029                &sl.k0,
19030                &sl.v0,
19031                fa.q_norm.float_data(),
19032                fa.k_norm.float_data(),
19033                ones,
19034                &mut sl.q,
19035                &mut sl.k,
19036                &mut sl.v,
19037                hd,
19038                self.gemma4_rope_dims(il),
19039                nh,
19040                nkv,
19041                pos_d,
19042                nh,
19043                nkv,
19044                base,
19045                1.0,
19046                ff,
19047                eps,
19048                &mut kvl.k,
19049                &mut kvl.v,
19050                &kvl.len_d,
19051                kvl.k_tok_bytes,
19052                kvl.v_tok_bytes,
19053                kv_fp8,
19054            )?;
19055        } else {
19056            e.rms_norm_qkv_rope(
19057                &sl.q0,
19058                &sl.k0,
19059                &sl.v0,
19060                fa.q_norm.float_data(),
19061                fa.k_norm.float_data(),
19062                ones,
19063                &mut sl.q,
19064                &mut sl.k,
19065                &mut sl.v,
19066                hd,
19067                self.gemma4_rope_dims(il),
19068                nh,
19069                nkv,
19070                pos_d,
19071                nh,
19072                nkv,
19073                base,
19074                1.0,
19075                ff,
19076                eps,
19077            )?;
19078            e.append_kv_quantized_dc(
19079                &sl.k,
19080                &sl.v,
19081                &mut kvl.k,
19082                &mut kvl.v,
19083                &kvl.len_d,
19084                kvl.kv_dim_k,
19085                kvl.kv_dim_v,
19086                kvl.k_tok_bytes,
19087                kvl.v_tok_bytes,
19088                kv_fp8,
19089            )?;
19090        }
19091        e.inc_seqlen(&mut kvl.len_d)?;
19092        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
19093        let k_view = e.view_u8(&kvl.k, kvl.k.len());
19094        let v_view = e.view_u8(&kvl.v, kvl.v.len());
19095        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19096        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19097        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
19098        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
19099        // the dc_into arm branch-for-branch (stream gate).
19100        let mut fa_q8 = false;
19101        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19102            e.fa_decode_rows(
19103                &sl.q,
19104                &k_view,
19105                &v_view,
19106                &mut sl.attn,
19107                hd,
19108                nh,
19109                nkv,
19110                b_glob - 1,
19111                1,
19112                scale,
19113                kvl.k_tok_bytes,
19114                kvl.v_tok_bytes,
19115                Some((&kvl.len_d, -1)),
19116                false,
19117                false,
19118                Some((&mut sl.zq, &mut sl.zd)),
19119            )?;
19120            fa_q8 = true;
19121        } else if swa && b_swa > win && hd == 256 && rows_on {
19122            e.fa_decode_rows_w(
19123                &sl.q,
19124                &k_view,
19125                &v_view,
19126                &mut sl.attn,
19127                hd,
19128                nh,
19129                nkv,
19130                &kvl.len_d,
19131                -1,
19132                1,
19133                scale,
19134                win,
19135                kvl.k_tok_bytes,
19136                kvl.v_tok_bytes,
19137                Some((&mut sl.zq, &mut sl.zd)),
19138            )?;
19139            fa_q8 = true;
19140        } else {
19141            let b = if swa { b_swa } else { b_glob };
19142            e.fa_decode_dc(
19143                &sl.q,
19144                &k_view,
19145                &v_view,
19146                &mut sl.attn,
19147                hd,
19148                nh,
19149                nkv,
19150                &kvl.len_d,
19151                b,
19152                scale,
19153                kvl.k_tok_bytes,
19154                kvl.v_tok_bytes,
19155                swa && crate::Engine::wkv_on(),
19156            )?;
19157        }
19158        if !fa_q8 {
19159            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
19160            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
19161        }
19162        {
19163            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19164            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19165            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
19166        }
19167        Ok(())
19168    }
19169
19170    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
19171    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
19172    fn gemma4_layer_tail_slotted(
19173        &self,
19174        e: &Engine,
19175        layer: &crate::hybrid::HybridLayer,
19176        next_norm: Option<&CudaSlice<f32>>,
19177        sl: &mut G4DcSlots,
19178    ) -> Result<(), Box<dyn std::error::Error>> {
19179        let n_embd = self.cfg.n_embd as usize;
19180        let eps = self.cfg.rms_eps;
19181        let bits = layer.gemma4.as_ref().unwrap();
19182        let crate::hybrid::Ffn::Dense {
19183            ffn_gate,
19184            ffn_up,
19185            ffn_down,
19186        } = &layer.ffn
19187        else {
19188            return Err("slotted tail: dense ffn only".into());
19189        };
19190        let pnfold = Engine::g4_pnfold_on();
19191        if pnfold {
19192            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
19193            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
19194            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
19195            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
19196            e.rms_pre_add_rms_norm_q8z_into(
19197                or,
19198                layer.post_attn_norm.float_data(),
19199                xr,
19200                bits.ffn_norm.float_data(),
19201                &mut sl.attn_out,
19202                &mut sl.zsh,
19203                n_embd,
19204                1,
19205                eps,
19206                &mut sl.zq,
19207                &mut sl.zd,
19208            )?;
19209        } else {
19210            e.add_rms_norm(
19211                &sl.cur,
19212                &sl.x,
19213                bits.ffn_norm.float_data(),
19214                &mut sl.attn_out,
19215                &mut sl.zsh,
19216                n_embd,
19217                1,
19218                eps,
19219            )?;
19220        }
19221        let n_ff = ffn_gate.out_features();
19222        if !pnfold {
19223            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
19224            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
19225        }
19226        {
19227            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
19228            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
19229            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
19230                && !e.matmul_nvfp4_fused2_into(
19231                    ffn_gate,
19232                    ffn_up,
19233                    zq,
19234                    zd,
19235                    &mut sl.gate,
19236                    &mut sl.up,
19237                )?
19238            {
19239                return Err("slotted tail: ffn fused2 unavailable".into());
19240            }
19241        }
19242        debug_assert!(e.uses_q8_1_fast(ffn_down));
19243        {
19244            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
19245            let upv = e.view(upr, n_ff);
19246            let up_all = upv.slice(0..n_ff);
19247            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
19248            e.gelu_tanh_mul_q8_1_into(
19249                gr,
19250                &up_all,
19251                &mut sl.act,
19252                n_ff,
19253                1,
19254                &mut sl.actq,
19255                &mut sl.actd,
19256            )?;
19257        }
19258        {
19259            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
19260            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
19261            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
19262        }
19263        if pnfold {
19264            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
19265            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
19266            if let Some(w) = next_norm {
19267                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
19268                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
19269                e.rms_pre_add_scale_rms_norm_q8_1_into(
19270                    f0r,
19271                    bits.post_ffw_norm.float_data(),
19272                    aor,
19273                    bits.layer_scale,
19274                    w,
19275                    &mut sl.xn,
19276                    n_embd,
19277                    1,
19278                    eps,
19279                    &mut sl.hq,
19280                    &mut sl.hd_,
19281                )?;
19282                return Ok(());
19283            }
19284        }
19285        e.rms_norm(
19286            &sl.f0,
19287            bits.post_ffw_norm.float_data(),
19288            &mut sl.sn,
19289            n_embd,
19290            1,
19291            eps,
19292        )?;
19293        match next_norm {
19294            Some(w) => {
19295                e.add_scale_rms_norm_q8_1_into(
19296                    &sl.sn,
19297                    &sl.attn_out,
19298                    bits.layer_scale,
19299                    w,
19300                    &mut sl.xn,
19301                    n_embd,
19302                    1,
19303                    eps,
19304                    &mut sl.hq,
19305                    &mut sl.hd_,
19306                )?;
19307            }
19308            None => {
19309                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
19310            }
19311        }
19312        Ok(())
19313    }
19314
19315    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
19316    #[allow(clippy::too_many_arguments)]
19317    fn gemma4_decode_attn_dc(
19318        &self,
19319        e: &Engine,
19320        fa: &crate::hybrid::FullAttnLayer,
19321        il: usize,
19322        hq: &CudaSlice<i8>,
19323        hdq: &CudaSlice<f32>,
19324        pos_d: &CudaSlice<i32>,
19325        cache: &mut Cache,
19326        cap_bucket_max: Option<(usize, usize)>,
19327    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19328        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
19329        let eps = self.cfg.rms_eps;
19330        let aux = self.gemma4_aux.as_ref().unwrap();
19331        let ones = aux.ones(e);
19332        #[cfg(debug_assertions)]
19333        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
19334        let (q0, k0, v0) = if swa {
19335            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
19336                Some(t3) => t3,
19337                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
19338                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19339                    Some((q0, k0)) => {
19340                        let h0 = e.zeros(0)?;
19341                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
19342                        (q0, k0, v0)
19343                    }
19344                    None => {
19345                        let h0 = e.zeros(0)?;
19346                        (
19347                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19348                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19349                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
19350                        )
19351                    }
19352                },
19353            }
19354        } else {
19355            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
19356                Some(p) => p,
19357                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
19358                    Some(p) => p,
19359                    None => {
19360                        let h0 = e.zeros(0)?;
19361                        (
19362                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
19363                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
19364                        )
19365                    }
19366                },
19367            };
19368            let v0 = e.clone_dtod(&k0)?;
19369            (q0, k0, v0)
19370        };
19371        let mut q = e.uninit(nh * hd)?;
19372        let mut k = e.uninit(nkv * hd)?;
19373        let mut v = e.uninit(nkv * hd)?;
19374        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
19375        let ff = if swa {
19376            None
19377        } else {
19378            Some(
19379                aux.rope_freqs(e)
19380                    .expect("gemma4 global rope needs rope_freqs.weight"),
19381            )
19382        };
19383        #[cfg(debug_assertions)]
19384        if let Some(ff) = ff {
19385            crate::debug_assert_tensor_stream_device(
19386                ff,
19387                &e.stream(),
19388                "gemma4_decode_attn_dc.rope_freqs",
19389            );
19390        }
19391        let kvl = cache.kv[il].as_mut().unwrap();
19392        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
19393        if crate::Engine::qkv_append_on() {
19394            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
19395            e.rms_norm_qkv_rope_append_dc(
19396                &q0,
19397                &k0,
19398                &v0,
19399                fa.q_norm.float_data(),
19400                fa.k_norm.float_data(),
19401                ones,
19402                &mut q,
19403                &mut k,
19404                &mut v,
19405                hd,
19406                self.gemma4_rope_dims(il),
19407                nh,
19408                nkv,
19409                pos_d,
19410                nh,
19411                nkv,
19412                base,
19413                1.0,
19414                ff,
19415                eps,
19416                &mut kvl.k,
19417                &mut kvl.v,
19418                &kvl.len_d,
19419                kvl.k_tok_bytes,
19420                kvl.v_tok_bytes,
19421                kv_fp8,
19422            )?;
19423        } else {
19424            e.rms_norm_qkv_rope(
19425                &q0,
19426                &k0,
19427                &v0,
19428                fa.q_norm.float_data(),
19429                fa.k_norm.float_data(),
19430                ones,
19431                &mut q,
19432                &mut k,
19433                &mut v,
19434                hd,
19435                self.gemma4_rope_dims(il),
19436                nh,
19437                nkv,
19438                pos_d,
19439                nh,
19440                nkv,
19441                base,
19442                1.0,
19443                ff,
19444                eps,
19445            )?;
19446            e.append_kv_quantized_dc(
19447                &k,
19448                &v,
19449                &mut kvl.k,
19450                &mut kvl.v,
19451                &kvl.len_d,
19452                kvl.kv_dim_k,
19453                kvl.kv_dim_v,
19454                kvl.k_tok_bytes,
19455                kvl.v_tok_bytes,
19456                kv_fp8,
19457            )?;
19458        }
19459        e.inc_seqlen(&mut kvl.len_d)?;
19460        let mut attn = e.uninit(nh * hd)?;
19461        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
19462        // rides g4_matvec_m1_into instead of matmul's internal quantize.
19463        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19464        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
19465        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
19466        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
19467        // (gemma4_e4b_attn, +0.65% valid window).
19468        match cap_bucket_max {
19469            None => {
19470                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
19471                // decode (SWA layers attend the last `sliding_window` keys); the device
19472                // counters carry only the append slot + the graph seam.
19473                kvl.len += 1;
19474                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19475                if !swa
19476                    && hd == 512
19477                    && kvl.len >= crate::fa512_min_tkv()
19478                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19479                {
19480                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
19481                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
19482                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19483                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19484                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19485                    e.fa_decode_rows(
19486                        &q,
19487                        &kp,
19488                        &vp,
19489                        &mut attn,
19490                        hd,
19491                        nh,
19492                        nkv,
19493                        kvl.len - 1,
19494                        1,
19495                        scale,
19496                        kvl.k_tok_bytes,
19497                        kvl.v_tok_bytes,
19498                        Some((&kvl.len_d, -1)),
19499                        false,
19500                        false,
19501                        Some((&mut aq8, &mut ad8)),
19502                    )?;
19503                    fa_q8 = Some((aq8, ad8));
19504                } else if swa
19505                    && kvl.len > win
19506                    && hd == 256
19507                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
19508                {
19509                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
19510                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
19511                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
19512                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19513                    e.fa_decode_rows_w(
19514                        &q,
19515                        &kp,
19516                        &vp,
19517                        &mut attn,
19518                        hd,
19519                        nh,
19520                        nkv,
19521                        &kvl.len_d,
19522                        -1,
19523                        1,
19524                        scale,
19525                        win,
19526                        kvl.k_tok_bytes,
19527                        kvl.v_tok_bytes,
19528                        Some((&mut aq8, &mut ad8)),
19529                    )?;
19530                    fa_q8 = Some((aq8, ad8));
19531                } else {
19532                    let (off_tok, t_kv) = if swa && kvl.len > win {
19533                        (kvl.len - win, win)
19534                    } else {
19535                        (0, kvl.len)
19536                    };
19537                    let k_view = e.view_u8_range(
19538                        &kvl.k,
19539                        off_tok * kvl.k_tok_bytes,
19540                        (off_tok + t_kv) * kvl.k_tok_bytes,
19541                    );
19542                    let v_view = e.view_u8_range(
19543                        &kvl.v,
19544                        off_tok * kvl.v_tok_bytes,
19545                        (off_tok + t_kv) * kvl.v_tok_bytes,
19546                    );
19547                    e.fa_decode_kvmod(
19548                        &q,
19549                        &k_view,
19550                        &v_view,
19551                        &mut attn,
19552                        hd,
19553                        nh,
19554                        nkv,
19555                        t_kv,
19556                        scale,
19557                        kvl.k_tok_bytes,
19558                        kvl.v_tok_bytes,
19559                        swa && crate::Engine::wkv_on(),
19560                    )?;
19561                }
19562            }
19563            Some((b_swa, b_glob)) => {
19564                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
19565                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
19566                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
19567                // the RUNG max for the rows family (kernels derive per-replay splits from
19568                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
19569                let k_view = e.view_u8(&kvl.k, kvl.k.len());
19570                let v_view = e.view_u8(&kvl.v, kvl.v.len());
19571                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
19572                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19573                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
19574                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19575                    e.fa_decode_rows(
19576                        &q,
19577                        &k_view,
19578                        &v_view,
19579                        &mut attn,
19580                        hd,
19581                        nh,
19582                        nkv,
19583                        b_glob - 1,
19584                        1,
19585                        scale,
19586                        kvl.k_tok_bytes,
19587                        kvl.v_tok_bytes,
19588                        Some((&kvl.len_d, -1)),
19589                        false,
19590                        false,
19591                        Some((&mut aq8, &mut ad8)),
19592                    )?;
19593                    fa_q8 = Some((aq8, ad8));
19594                } else if swa && b_swa > win && hd == 256 && rows_on {
19595                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
19596                    e.fa_decode_rows_w(
19597                        &q,
19598                        &k_view,
19599                        &v_view,
19600                        &mut attn,
19601                        hd,
19602                        nh,
19603                        nkv,
19604                        &kvl.len_d,
19605                        -1,
19606                        1,
19607                        scale,
19608                        win,
19609                        kvl.k_tok_bytes,
19610                        kvl.v_tok_bytes,
19611                        Some((&mut aq8, &mut ad8)),
19612                    )?;
19613                    fa_q8 = Some((aq8, ad8));
19614                } else {
19615                    let b = if swa { b_swa } else { b_glob };
19616                    e.fa_decode_dc(
19617                        &q,
19618                        &k_view,
19619                        &v_view,
19620                        &mut attn,
19621                        hd,
19622                        nh,
19623                        nkv,
19624                        &kvl.len_d,
19625                        b,
19626                        scale,
19627                        kvl.k_tok_bytes,
19628                        kvl.v_tok_bytes,
19629                        swa && crate::Engine::wkv_on(),
19630                    )?;
19631                }
19632            }
19633        }
19634        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
19635        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
19636        if let Some((aq8, ad8)) = fa_q8 {
19637            let mut y = e.uninit(fa.wo.out_features())?;
19638            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
19639            return Ok(y);
19640        }
19641        e.matmul(&fa.wo, &attn, 1)
19642    }
19643
19644    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
19645    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
19646    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
19647    /// views in-graph); caller gates and falls back to the dc-eager loop.
19648    #[allow(clippy::too_many_arguments)]
19649    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
19650    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
19651    pub fn gemma4_generate_graph(
19652        &self,
19653        e: &Engine,
19654        prompt_pos: usize,
19655        first_token: u32,
19656        cache: &mut Cache,
19657        max_new: usize,
19658        eos: &[u32],
19659        mut on_token: impl FnMut(u32) -> bool,
19660    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
19661        if self.is_gemma4_e4b() {
19662            return Err(
19663                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
19664                    .into(),
19665            );
19666        }
19667        use crate::decode::StopReason;
19668        let n_vocab = self.output.out_features();
19669        let n_embd = self.cfg.n_embd as usize;
19670        let embd_gpu = self
19671            .embd_gpu
19672            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
19673        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
19674        for kvl in cache.kv.iter_mut().flatten() {
19675            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
19676        }
19677        let mut token_d = e.stream().clone_htod(&[first_token])?;
19678        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
19679        let g4 = self.cfg.gemma4.as_ref().unwrap();
19680        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
19681        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
19682        let nkv_s = g4
19683            .head_count_kv
19684            .iter()
19685            .zip(g4.swa_pattern.iter())
19686            .find(|p| *p.1)
19687            .map(|p| *p.0 as usize)
19688            .unwrap_or(8);
19689        let nkv_g = g4
19690            .head_count_kv
19691            .iter()
19692            .zip(g4.swa_pattern.iter())
19693            .find(|p| !*p.1)
19694            .map(|p| *p.0 as usize)
19695            .unwrap_or(2);
19696        #[allow(clippy::type_complexity)]
19697        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
19698        let mut graphs: std::collections::HashMap<
19699            ((bool, usize), (bool, usize), bool, bool),
19700            (
19701                cudarc::driver::CudaGraph,
19702                Vec<Box<dyn std::any::Any + Send>>,
19703            ),
19704        > = Default::default();
19705        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
19706        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
19707        let mut slots = self.g4_dc_slots(e)?;
19708        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
19709        // baked at the door entry (the modulo keeps every capture valid indefinitely).
19710        const RING: usize = 64;
19711        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
19712        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
19713        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
19714        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
19715        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
19716        const DRAIN: usize = 1;
19717        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
19718        let ring_base = prompt_pos;
19719        let mut out = Vec::with_capacity(max_new);
19720        let mut reason = StopReason::MaxNew;
19721        let mut next = first_token;
19722        let mut captures = 0usize;
19723        for _ in 0..max_new {
19724            out.push(next);
19725            if eos.contains(&next) {
19726                reason = StopReason::Eos;
19727                break;
19728            }
19729            if !on_token(next) {
19730                reason = StopReason::Callback;
19731                break;
19732            }
19733            let t_kv = cache.pos + 1;
19734            // Bucket key per ARM (graph arc step 3):
19735            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
19736            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
19737            //    the component collapses to a single marker).
19738            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
19739            //    at/above it — the kernel derives splits from len_d per replay, so buckets
19740            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
19741            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
19742            let f512 = crate::fa512_min_tkv();
19743            let key_s = if t_kv > win {
19744                (true, usize::MAX)
19745            } else {
19746                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
19747            };
19748            let (key_g, rung_end) = if t_kv >= f512 {
19749                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
19750                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
19751                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
19752                ((true, end), end)
19753            } else {
19754                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
19755            };
19756            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
19757            if !graphs.contains_key(&key) {
19758                let bucket_max = (t_kv, rung_end);
19759                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
19760                let snap = cache.snapshot(e)?;
19761                let pos_save = e.dtoh_i32_one(&pos_d)?;
19762                let len_save: Vec<Option<i32>> = cache
19763                    .kv
19764                    .iter()
19765                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
19766                    .collect();
19767                let tok_save = e.dtoh_u32_one(&token_d)?;
19768                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
19769                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
19770                // regression class, and this door's measured -8.8%. The keeper pins warmup
19771                // transients so the captured graph holds kernel nodes only.
19772                let graph = {
19773                    let tok_ref = &mut token_d;
19774                    let pos_ref = &mut pos_d;
19775                    let cache_ref = &mut *cache;
19776                    let slots_ref = &mut slots;
19777                    let ring_ref = &mut ring;
19778                    e.capture_graph_retained_flags(
19779                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
19780                        |e| {
19781                        // self-feeding: the argmax writes token_d itself.
19782                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
19783                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
19784                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
19785                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
19786                                                           cache_ref, n_vocab, Some(bucket_max),
19787                                                           sl, tok_ref, Some((rg, ring_base)))
19788                    })?
19789                };
19790                cache.rollback(e, &snap, 0)?;
19791                e.set_i32_one(&mut pos_d, pos_save)?;
19792                for (il, ls) in len_save.iter().enumerate() {
19793                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
19794                        e.set_i32_one(&mut kvl.len_d, *v)?;
19795                    }
19796                }
19797                e.set_u32_one(&mut token_d, tok_save)?;
19798                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1")
19799                    && let Ok(c) = crate::graph_update::node_census(&graph.0)
19800                {
19801                    eprintln!("[graph-census] {c:?}");
19802                }
19803                graphs.insert(key, graph);
19804                captures += 1;
19805            }
19806            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
19807            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
19808            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
19809            // the budget; capture warmups already emitted their tokens through the ring.
19810            let mut chunk = 1usize;
19811            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
19812                .ok()
19813                .and_then(|v| v.parse().ok())
19814                .unwrap_or(DRAIN);
19815            while chunk < drain_cap && out.len() + chunk < max_new {
19816                let t_next = cache.pos + 1 + chunk;
19817                let key_s2 = if t_next > win {
19818                    (true, usize::MAX)
19819                } else {
19820                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
19821                };
19822                let key_g2 = if t_next >= f512 {
19823                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
19824                } else {
19825                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
19826                };
19827                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
19828                    break;
19829                }
19830                chunk += 1;
19831            }
19832            let g = &graphs.get(&key).unwrap().0;
19833            for _ in 0..chunk {
19834                g.launch()?;
19835            }
19836            e.stream().synchronize()?;
19837            let ringh = e.dtoh_u32(&ring)?;
19838            for j in 0..chunk {
19839                let pos_j = cache.pos + j;
19840                let tok_j = ringh[(pos_j - ring_base) % RING];
19841                cache.pos += 0; // advanced below in one shot
19842                if j + 1 == chunk {
19843                    next = tok_j;
19844                } else {
19845                    out.push(tok_j);
19846                    if eos.contains(&tok_j) || !on_token(tok_j) {
19847                        reason = if eos.contains(&tok_j) {
19848                            StopReason::Eos
19849                        } else {
19850                            StopReason::Callback
19851                        };
19852                        // roll device/host state back to the stop point.
19853                        let keep = cache.pos + j + 1;
19854                        e.set_i32_one(&mut pos_d, keep as i32)?;
19855                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
19856                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
19857                            kvl.len = keep;
19858                        }
19859                        cache.pos = keep;
19860                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
19861                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
19862                        }
19863                        return Ok((out, reason));
19864                    }
19865                }
19866            }
19867            cache.pos += chunk;
19868            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
19869                kvl.len += chunk;
19870            }
19871        }
19872        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
19873            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
19874        }
19875        Ok((out, reason))
19876    }
19877
19878    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
19879    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
19880    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
19881    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
19882    /// logits (host) + advances cache.pos by t.
19883    pub(crate) fn gemma4_decode_step_t(
19884        &self,
19885        e: &Engine,
19886        tokens: &[u32],
19887        pos0: usize,
19888        cache: &mut Cache,
19889    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
19890        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
19891    }
19892
19893    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
19894    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
19895    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
19896    pub(crate) fn gemma4_decode_step_t_am(
19897        &self,
19898        e: &Engine,
19899        tokens: &[u32],
19900        pos0: usize,
19901        cache: &mut Cache,
19902    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19903        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
19904        let t = tokens.len();
19905        let n_vocab = self.output.out_features();
19906        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
19907        for i in 0..t {
19908            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
19909        }
19910        Ok((e.dtoh_u32(&toks)?, hn))
19911    }
19912
19913    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
19914    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
19915    pub(crate) fn gemma4_decode_step_t_am_dev(
19916        &self,
19917        e: &Engine,
19918        tok_d: &CudaSlice<u32>,
19919        t: usize,
19920        pos0: usize,
19921        cache: &mut Cache,
19922    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19923        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
19924        let n_vocab = self.output.out_features();
19925        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
19926        for i in 0..t {
19927            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
19928        }
19929        Ok((vam, hn))
19930    }
19931
19932    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
19933    /// llama's h_nextn convention).
19934    pub(crate) fn gemma4_decode_step_t_h(
19935        &self,
19936        e: &Engine,
19937        tokens: &[u32],
19938        pos0: usize,
19939        cache: &mut Cache,
19940    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19941        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
19942        let t = tokens.len();
19943        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
19944        e.softcap(&mut ld, cap, t * self.output.out_features())?;
19945        Ok((e.dtoh(&ld)?, hn))
19946    }
19947
19948    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
19949    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
19950    pub(crate) fn verify_stream_scratch(
19951        &self,
19952        e: &Engine,
19953        cap: usize,
19954    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
19955        Ok(VerifyStreamScratch {
19956            pos_d: e.htod_i32(&vec![0i32; cap])?,
19957            row_ctrs: (0..cap)
19958                .map(|_| e.htod_i32(&[0]))
19959                .collect::<Result<_, _>>()?,
19960        })
19961    }
19962
19963    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
19964    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
19965    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
19966    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
19967    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
19968    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
19969    /// sync, exactly the turnaround the burst exists to remove.
19970    #[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
19971    pub(crate) fn gemma4_verify_t_am_stream(
19972        &self,
19973        e: &Engine,
19974        tok_d: &CudaSlice<u32>,
19975        t: usize,
19976        ctr: &CudaSlice<i32>,
19977        hint: usize,
19978        cache: &mut Cache,
19979        scr: &mut VerifyStreamScratch,
19980    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19981        let n_embd = self.cfg.n_embd as usize;
19982        let eps = self.cfg.rms_eps;
19983        assert!(t <= scr.row_ctrs.len() && t <= 64);
19984        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
19985        for i in 0..t {
19986            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
19987        }
19988        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
19989        let embd_gpu = self
19990            .embd_gpu
19991            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
19992        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
19993        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
19994        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
19995        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
19996        let n_layers = self.layers.len();
19997        for (il, layer) in self.layers.iter().enumerate() {
19998            let (hq, hdq) = match h_carry.take() {
19999                Some(p) => p,
20000                None => {
20001                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20002                }
20003            };
20004            let Mixer::Full(fa) = &layer.mixer else {
20005                panic!("gemma4 layer {il} not full-attn")
20006            };
20007            let o = self
20008                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
20009            let next_norm = if il + 1 < n_layers {
20010                Some(self.layers[il + 1].attn_norm.float_data())
20011            } else {
20012                None
20013            };
20014            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20015            x = xn;
20016            h_carry = hn;
20017            self.dflash_tap(e, cache, il, &x, t)?;
20018        }
20019        let mut hn = e.uninit(t * n_embd)?;
20020        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20021        let ld = e.matmul(&self.output, &hn, t)?;
20022        let n_vocab = self.output.out_features();
20023        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
20024        for i in 0..t {
20025            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
20026        }
20027        Ok((vam, hn))
20028    }
20029
20030    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
20031    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
20032    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
20033    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
20034    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
20035    /// kernel later if it shows in the profile).
20036    pub(crate) fn dflash_tap(
20037        &self,
20038        e: &Engine,
20039        cache: &mut Cache,
20040        il: usize,
20041        x: &CudaSlice<f32>,
20042        t: usize,
20043    ) -> Result<(), Box<dyn std::error::Error>> {
20044        let Some(taps) = cache.dflash_taps.as_mut() else {
20045            return Ok(());
20046        };
20047        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
20048            return Ok(());
20049        };
20050        let h = taps.hidden;
20051        let n_taps = taps.layer_ids.len();
20052        let base = taps.base;
20053        debug_assert!(
20054            base + t <= taps.t,
20055            "tap window {base}+{t} exceeds sink {}",
20056            taps.t
20057        );
20058        let xv = e.view(x, t * h);
20059        for r in 0..t {
20060            let row = xv.slice(r * h..(r + 1) * h);
20061            e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
20062        }
20063        Ok(())
20064    }
20065
20066    fn gemma4_verify_trunk(
20067        &self,
20068        e: &Engine,
20069        tokens: &[u32],
20070        pos0: usize,
20071        cache: &mut Cache,
20072        tok_dev: Option<&CudaSlice<u32>>,
20073    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20074        let n_embd = self.cfg.n_embd as usize;
20075        let eps = self.cfg.rms_eps;
20076        let t = tokens.len();
20077        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
20078        let pos_d = e.htod_i32(&pos)?;
20079        let mut x = match tok_dev {
20080            Some(td) => {
20081                let embd_gpu = self
20082                    .embd_gpu
20083                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
20084                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
20085                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
20086            }
20087            None => e.htod(&self.embd.gather(n_embd, tokens))?,
20088        };
20089        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
20090        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20091        let n_layers = self.layers.len();
20092        for (il, layer) in self.layers.iter().enumerate() {
20093            let (hq, hdq) = match h_carry.take() {
20094                Some(p) => p,
20095                None => {
20096                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
20097                }
20098            };
20099            let Mixer::Full(fa) = &layer.mixer else {
20100                panic!("gemma4 layer {il} not full-attn")
20101            };
20102            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
20103            let next_norm = if il + 1 < n_layers {
20104                Some(self.layers[il + 1].attn_norm.float_data())
20105            } else {
20106                None
20107            };
20108            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
20109            x = xn;
20110            h_carry = hn;
20111            self.dflash_tap(e, cache, il, &x, t)?;
20112        }
20113        let mut hn = e.uninit(t * n_embd)?;
20114        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
20115        let mut ld = e.matmul(&self.output, &hn, t)?;
20116        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
20117        cache.pos += t;
20118        Ok((ld, hn))
20119    }
20120
20121    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
20122    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
20123    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
20124    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
20125    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
20126    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
20127    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
20128    #[allow(clippy::too_many_arguments)]
20129    fn gemma4_verify_attn_stream(
20130        &self,
20131        e: &Engine,
20132        fa: &crate::hybrid::FullAttnLayer,
20133        il: usize,
20134        hq: &CudaSlice<i8>,
20135        hdq: &CudaSlice<f32>,
20136        pos_d: &CudaSlice<i32>,
20137        t: usize,
20138        cache: &mut Cache,
20139        hint: usize,
20140        row_ctrs: &[CudaSlice<i32>],
20141    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20142        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20143        let eps = self.cfg.rms_eps;
20144        let aux = self.gemma4_aux.as_ref().unwrap();
20145        let ones = aux.ones(e);
20146        #[cfg(debug_assertions)]
20147        crate::debug_assert_tensor_stream_device(
20148            ones,
20149            &e.stream(),
20150            "gemma4_verify_attn_stream.ones",
20151        );
20152        let h0 = e.zeros(0)?;
20153        let h = &h0;
20154        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
20155        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
20156        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20157        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20158        let fused_qkv = if f2b {
20159            if swa {
20160                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20161                    .map(|(a, b, c)| (a, b, Some(c)))
20162            } else {
20163                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20164                    .map(|(a, b)| (a, b, None))
20165            }
20166        } else {
20167            None
20168        };
20169        let (q0, k0, v0) = match fused_qkv {
20170            Some((a, b, cv)) => {
20171                let v = match cv {
20172                    Some(c) => c,
20173                    None => e.clone_dtod(&b)?,
20174                };
20175                (a, b, v)
20176            }
20177            None => {
20178                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20179                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20180                let v0 = if swa {
20181                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20182                } else {
20183                    e.clone_dtod(&k0)?
20184                };
20185                (q0, k0, v0)
20186            }
20187        };
20188        let mut q = e.uninit(t * nh * hd)?;
20189        let mut k = e.uninit(t * nkv * hd)?;
20190        let mut v = e.uninit(t * nkv * hd)?;
20191        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
20192        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
20193        let ff = if swa {
20194            None
20195        } else {
20196            Some(
20197                aux.rope_freqs(e)
20198                    .expect("gemma4 global rope needs rope_freqs.weight"),
20199            )
20200        };
20201        #[cfg(debug_assertions)]
20202        if let Some(ff) = ff {
20203            crate::debug_assert_tensor_stream_device(
20204                ff,
20205                &e.stream(),
20206                "gemma4_verify_attn_stream.rope_freqs",
20207            );
20208        }
20209        e.rms_norm_qkv_rope(
20210            &q0,
20211            &k0,
20212            &v0,
20213            fa.q_norm.float_data(),
20214            fa.k_norm.float_data(),
20215            ones,
20216            &mut q,
20217            &mut k,
20218            &mut v,
20219            hd,
20220            self.gemma4_rope_dims(il),
20221            nh * t,
20222            nkv * t,
20223            pos_d,
20224            nh,
20225            nkv,
20226            base,
20227            1.0,
20228            ff,
20229            eps,
20230        )?;
20231        let kvl = cache.kv[il].as_mut().unwrap();
20232        // append at the DEVICE slot; the counter advances by t on-device.
20233        e.append_kv_quantized_rows_dc(
20234            &k,
20235            &v,
20236            &mut kvl.k,
20237            &mut kvl.v,
20238            &kvl.len_d,
20239            t,
20240            kvl.kv_dim_k,
20241            kvl.kv_dim_v,
20242            kvl.k_tok_bytes,
20243            kvl.v_tok_bytes,
20244            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20245        )?;
20246        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
20247        // the sole len writer after this round's attention (base stays = old len, plus = 0).
20248        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20249        let mut attn = e.uninit(t * nh * hd)?;
20250        let k_view = e.view_u8(&kvl.k, kvl.k.len());
20251        let v_view = e.view_u8(&kvl.v, kvl.v.len());
20252        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
20253        // and a stable window regime — the same rung/regime keys as the draft graph).
20254        if swa && hint + 1 >= win {
20255            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
20256            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
20257            e.fa_decode_rows_w(
20258                &q,
20259                &k_view,
20260                &v_view,
20261                &mut attn,
20262                hd,
20263                nh,
20264                nkv,
20265                &kvl.len_d,
20266                0,
20267                t,
20268                scale,
20269                win,
20270                kvl.k_tok_bytes,
20271                kvl.v_tok_bytes,
20272                None,
20273            )?;
20274        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
20275            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
20276            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
20277            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
20278            // Burst entry gates the horizon onto one side of the crossover, so hint decides
20279            // for every row.
20280            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
20281            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
20282            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
20283            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
20284            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
20285            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
20286            // any bucket >= the live length is exact.
20287            let bucket = (hint + t + 2)
20288                .next_power_of_two()
20289                .min(crate::fa512_min_tkv().saturating_sub(1));
20290            let qv = e.view(&q, t * nh * hd);
20291            #[allow(clippy::needless_range_loop)]
20292            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
20293            for i in 0..t {
20294                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
20295                let mut q_one = e.uninit(nh * hd)?;
20296                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20297                let mut a_one = e.uninit(nh * hd)?;
20298                e.fa_decode_dc(
20299                    &q_one,
20300                    &k_view,
20301                    &v_view,
20302                    &mut a_one,
20303                    hd,
20304                    nh,
20305                    nkv,
20306                    &row_ctrs[i],
20307                    bucket,
20308                    scale,
20309                    kvl.k_tok_bytes,
20310                    kvl.v_tok_bytes,
20311                    false,
20312                )?;
20313                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20314            }
20315        } else if hd == 512 {
20316            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
20317            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
20318            e.fa_decode_rows(
20319                &q,
20320                &k_view,
20321                &v_view,
20322                &mut attn,
20323                hd,
20324                nh,
20325                nkv,
20326                hint,
20327                t,
20328                scale,
20329                kvl.k_tok_bytes,
20330                kvl.v_tok_bytes,
20331                Some((&kvl.len_d, 0)),
20332                false,
20333                false,
20334                None,
20335            )?;
20336        } else {
20337            // hd256 under-window: v4 device-len rows twin.
20338            e.fa_decode_rows_dc(
20339                &q,
20340                &k_view,
20341                &v_view,
20342                &mut attn,
20343                hd,
20344                nh,
20345                nkv,
20346                &kvl.len_d,
20347                hint + t,
20348                t,
20349                scale,
20350                kvl.k_tok_bytes,
20351                kvl.v_tok_bytes,
20352                0,
20353                swa && crate::Engine::wkv_on(),
20354            )?;
20355        }
20356        e.matmul(&fa.wo, &attn, t)
20357    }
20358
20359    #[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
20360    fn gemma4_verify_attn(
20361        &self,
20362        e: &Engine,
20363        fa: &crate::hybrid::FullAttnLayer,
20364        il: usize,
20365        hq: &CudaSlice<i8>,
20366        hdq: &CudaSlice<f32>,
20367        pos_d: &CudaSlice<i32>,
20368        t: usize,
20369        cache: &mut Cache,
20370    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20371        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
20372        let eps = self.cfg.rms_eps;
20373        let aux = self.gemma4_aux.as_ref().unwrap();
20374        let ones = aux.ones(e);
20375        #[cfg(debug_assertions)]
20376        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
20377        let n_embd = self.cfg.n_embd as usize;
20378        let _ = n_embd;
20379
20380        let h0 = e.zeros(0)?;
20381        let h = &h0;
20382        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
20383        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
20384        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20385        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
20386        let fused_qkv = if f2b {
20387            if swa {
20388                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
20389                    .map(|(a, b, c)| (a, b, Some(c)))
20390            } else {
20391                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
20392                    .map(|(a, b)| (a, b, None))
20393            }
20394        } else {
20395            None
20396        };
20397        let (q0, k0, v0) = match fused_qkv {
20398            Some((a, b, cv)) => {
20399                let v = match cv {
20400                    Some(c) => c,
20401                    None => e.clone_dtod(&b)?,
20402                };
20403                (a, b, v)
20404            }
20405            None => {
20406                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
20407                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
20408                let v0 = if swa {
20409                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
20410                } else {
20411                    e.clone_dtod(&k0)?
20412                };
20413                (q0, k0, v0)
20414            }
20415        };
20416        let mut q = e.uninit(t * nh * hd)?;
20417        let mut k = e.uninit(t * nkv * hd)?;
20418        let mut v = e.uninit(t * nkv * hd)?;
20419        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
20420        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
20421        let ff = if swa {
20422            None
20423        } else {
20424            Some(
20425                aux.rope_freqs(e)
20426                    .expect("gemma4 global rope needs rope_freqs.weight"),
20427            )
20428        };
20429        #[cfg(debug_assertions)]
20430        if let Some(ff) = ff {
20431            crate::debug_assert_tensor_stream_device(
20432                ff,
20433                &e.stream(),
20434                "gemma4_verify_attn.rope_freqs",
20435            );
20436        }
20437        e.rms_norm_qkv_rope(
20438            &q0,
20439            &k0,
20440            &v0,
20441            fa.q_norm.float_data(),
20442            fa.k_norm.float_data(),
20443            ones,
20444            &mut q,
20445            &mut k,
20446            &mut v,
20447            hd,
20448            self.gemma4_rope_dims(il),
20449            nh * t,
20450            nkv * t,
20451            pos_d,
20452            nh,
20453            nkv,
20454            base,
20455            1.0,
20456            ff,
20457            eps,
20458        )?;
20459        let kvl = cache.kv[il].as_mut().unwrap();
20460        let base_len = kvl.len;
20461        e.append_kv_quantized_rows(
20462            &k,
20463            &v,
20464            &mut kvl.k,
20465            &mut kvl.v,
20466            base_len,
20467            t,
20468            kvl.kv_dim_k,
20469            kvl.kv_dim_v,
20470            kvl.k_tok_bytes,
20471            kvl.v_tok_bytes,
20472            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
20473        )?;
20474        kvl.len += t;
20475        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
20476        let mut attn = e.uninit(t * nh * hd)?;
20477        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
20478        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
20479        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
20480            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
20481            // decode rides the SAME symbol at t=1 (parity law).
20482            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
20483        if rows_ok && (!swa || base_len + t <= win) {
20484            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20485            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20486            if hd == 512 {
20487                // device-len twin: sync the counter to the verify base (async arg-store).
20488                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20489                e.fa_decode_rows(
20490                    &q,
20491                    &k_view,
20492                    &v_view,
20493                    &mut attn,
20494                    hd,
20495                    nh,
20496                    nkv,
20497                    base_len,
20498                    t,
20499                    scale,
20500                    kvl.k_tok_bytes,
20501                    kvl.v_tok_bytes,
20502                    Some((&kvl.len_d, 0)),
20503                    false,
20504                    swa && crate::Engine::wkv_on(),
20505                    None,
20506                )?;
20507            } else {
20508                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
20509                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
20510                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
20511                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20512                e.fa_decode_rows_dc(
20513                    &q,
20514                    &k_view,
20515                    &v_view,
20516                    &mut attn,
20517                    hd,
20518                    nh,
20519                    nkv,
20520                    &kvl.len_d,
20521                    base_len + t,
20522                    t,
20523                    scale,
20524                    kvl.k_tok_bytes,
20525                    kvl.v_tok_bytes,
20526                    0,
20527                    swa && crate::Engine::wkv_on(),
20528                )?;
20529            }
20530            return e.matmul(&fa.wo, &attn, t);
20531        }
20532        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
20533        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
20534        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
20535        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
20536        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
20537        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
20538        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
20539        if hd == 256
20540            && swa
20541            && base_len + 1 >= win
20542            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20543        {
20544            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
20545            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
20546            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
20547            e.fa_decode_rows_w(
20548                &q,
20549                &k_view,
20550                &v_view,
20551                &mut attn,
20552                hd,
20553                nh,
20554                nkv,
20555                &kvl.len_d,
20556                0,
20557                t,
20558                scale,
20559                win,
20560                kvl.k_tok_bytes,
20561                kvl.v_tok_bytes,
20562                None,
20563            )?;
20564            return e.matmul(&fa.wo, &attn, t);
20565        }
20566        for i in 0..t {
20567            let avail = base_len + i + 1;
20568            let (off_tok, t_kv) = if swa && avail > win {
20569                (avail - win, win)
20570            } else {
20571                (0, avail)
20572            };
20573            let k_view = e.view_u8_range(
20574                &kvl.k,
20575                off_tok * kvl.k_tok_bytes,
20576                (off_tok + t_kv) * kvl.k_tok_bytes,
20577            );
20578            let v_view = e.view_u8_range(
20579                &kvl.v,
20580                off_tok * kvl.v_tok_bytes,
20581                (off_tok + t_kv) * kvl.v_tok_bytes,
20582            );
20583            let qi = e.view(&q, t * nh * hd);
20584            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
20585            let mut q_one = e.uninit(nh * hd)?;
20586            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
20587            let mut a_one = e.uninit(nh * hd)?;
20588            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
20589            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
20590            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
20591            if swa
20592                && avail > win
20593                && hd == 256
20594                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20595            {
20596                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20597                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20598                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20599                e.fa_decode_rows_w(
20600                    &q_one,
20601                    &kp,
20602                    &vp,
20603                    &mut a_one,
20604                    hd,
20605                    nh,
20606                    nkv,
20607                    &kvl.len_d,
20608                    0,
20609                    1,
20610                    scale,
20611                    win,
20612                    kvl.k_tok_bytes,
20613                    kvl.v_tok_bytes,
20614                    None,
20615                )?;
20616            } else if !swa
20617                && hd == 512
20618                && avail >= crate::fa512_min_tkv()
20619                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
20620            {
20621                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
20622                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
20623                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
20624                e.fa_decode_rows(
20625                    &q_one,
20626                    &kp,
20627                    &vp,
20628                    &mut a_one,
20629                    hd,
20630                    nh,
20631                    nkv,
20632                    avail - 1,
20633                    1,
20634                    scale,
20635                    kvl.k_tok_bytes,
20636                    kvl.v_tok_bytes,
20637                    Some((&kvl.len_d, 0)),
20638                    false,
20639                    false,
20640                    None,
20641                )?;
20642            } else {
20643                e.fa_decode_kvmod(
20644                    &q_one,
20645                    &k_view,
20646                    &v_view,
20647                    &mut a_one,
20648                    hd,
20649                    nh,
20650                    nkv,
20651                    t_kv,
20652                    scale,
20653                    kvl.k_tok_bytes,
20654                    kvl.v_tok_bytes,
20655                    swa && crate::Engine::wkv_on(),
20656                )?;
20657            }
20658            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
20659        }
20660        e.matmul(&fa.wo, &attn, t)
20661    }
20662
20663    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
20664    /// h_seed = pre-output_norm hidden). Advances cache.pos.
20665    pub(crate) fn gemma4_decode_step_h(
20666        &self,
20667        e: &Engine,
20668        token: u32,
20669        cache: &mut Cache,
20670    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20671        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
20672        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
20673        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
20674        // unsplit rather than guessing a fence.
20675        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
20676            let rt = crate::pp::Pp2Rt::get(e)?;
20677            let _walk = rt.acquire_walk("gemma4_decode_step_h_pp2")?;
20678            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
20679        }
20680        if crate::pp::pp_cuts(self.layers.len()).is_some() {
20681            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
20682        }
20683        let n_embd = self.cfg.n_embd as usize;
20684        let eps = self.cfg.rms_eps;
20685        let pos_d = e.htod_i32(&[cache.pos as i32])?;
20686        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20687        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20688        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
20689        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
20690        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20691        let n_layers = self.layers.len();
20692        for (il, layer) in self.layers.iter().enumerate() {
20693            let (hq, hdq) = match h_carry.take() {
20694                Some(p) => p,
20695                None => {
20696                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
20697                }
20698            };
20699            let Mixer::Full(fa) = &layer.mixer else {
20700                panic!("gemma4 layer {il} not full-attn")
20701            };
20702            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
20703            let next_norm = if il + 1 < n_layers {
20704                Some(self.layers[il + 1].attn_norm.float_data())
20705            } else {
20706                None
20707            };
20708            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20709            x = xn;
20710            h_carry = hn;
20711        }
20712        let mut hn = e.uninit(n_embd)?;
20713        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20714        let h_seed = e.clone_dtod(&x)?;
20715        let mut ld = e.matmul(&self.output, &hn, 1)?;
20716        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20717        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
20718        self.gemma4_suppress(e, &mut ld, 1)?;
20719        let logits = e.dtoh(&ld)?;
20720        cache.pos += 1;
20721        Ok((logits, h_seed))
20722    }
20723
20724    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
20725    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
20726    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
20727    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
20728    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
20729    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
20730    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
20731    fn gemma4_decode_layers(
20732        &self,
20733        e: &Engine,
20734        mut x: CudaSlice<f32>,
20735        lo: usize,
20736        hi: usize,
20737        pos_d: &CudaSlice<i32>,
20738        cache: &mut Cache,
20739    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20740        let n_embd = self.cfg.n_embd as usize;
20741        let eps = self.cfg.rms_eps;
20742        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
20743        for il in lo..hi {
20744            let layer = &self.layers[il];
20745            let (hq, hdq) = match h_carry.take() {
20746                Some(p) => p,
20747                // range head: il == lo — norm against THIS layer's attn_norm.
20748                None => {
20749                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
20750                }
20751            };
20752            let Mixer::Full(fa) = &layer.mixer else {
20753                panic!("gemma4 layer {il} not full-attn")
20754            };
20755            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
20756            let next_norm = if il + 1 < hi {
20757                Some(self.layers[il + 1].attn_norm.float_data())
20758            } else {
20759                None
20760            };
20761            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
20762            x = xn;
20763            h_carry = hn;
20764        }
20765        Ok(x)
20766    }
20767
20768    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
20769    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
20770    /// boundary handoff — same choreography as the generic arm (decode.rs), same
20771    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
20772    /// stage 1 = layers [split, n) + output_norm + softcapped head.
20773    /// Each stage uploads its own copy of the step's position scalar on its own stream.
20774    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
20775    fn gemma4_decode_step_h_pp2(
20776        &self,
20777        e: &Engine,
20778        token: u32,
20779        cache: &mut Cache,
20780        split: usize,
20781    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20782        if crate::pp::pp2_streams_off() {
20783            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
20784        }
20785        let rt = crate::pp::Pp2Rt::get(e)?;
20786        let e0 = rt.engine(0, e);
20787        let e1 = rt.engine(1, e);
20788        let n_embd = self.cfg.n_embd as usize;
20789        let eps = self.cfg.rms_eps;
20790        let pos = cache.pos as i32;
20791
20792        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
20793        let slot = {
20794            let _st0 = rt.enter(0);
20795            let pos_d = e0.htod_i32(&[pos])?;
20796            #[cfg(debug_assertions)]
20797            crate::debug_assert_tensor_stream_device(
20798                &pos_d,
20799                &e0.stream(),
20800                "gemma4_decode_step_h_pp2.stage0.pos_d",
20801            );
20802            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
20803            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20804            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
20805            rt.tx(0, &x, n_embd)?
20806        };
20807
20808        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
20809        let _st1 = rt.enter(1);
20810        let pos_d = e1.htod_i32(&[pos])?;
20811        #[cfg(debug_assertions)]
20812        crate::debug_assert_tensor_stream_device(
20813            &pos_d,
20814            &e1.stream(),
20815            "gemma4_decode_step_h_pp2.stage1.pos_d",
20816        );
20817        let x = rt.rx(0, slot, n_embd)?;
20818        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
20819
20820        let mut hn = e1.uninit(n_embd)?;
20821        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20822        let h_seed = e1.clone_dtod(&x)?;
20823        let mut ld = e1.matmul(&self.output, &hn, 1)?;
20824        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20825        e1.softcap(&mut ld, cap, self.output.out_features())?;
20826        self.gemma4_suppress(e1, &mut ld, 1)?;
20827        let logits = e1.dtoh(&ld)?;
20828        cache.pos += 1;
20829        Ok((logits, h_seed))
20830    }
20831
20832    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
20833    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
20834    fn gemma4_decode_step_h_pp2_samestream(
20835        &self,
20836        e: &Engine,
20837        token: u32,
20838        cache: &mut Cache,
20839        split: usize,
20840    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
20841        let n_embd = self.cfg.n_embd as usize;
20842        let eps = self.cfg.rms_eps;
20843        let pos_d = e.htod_i32(&[cache.pos as i32])?;
20844
20845        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
20846        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
20847        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
20848        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
20849
20850        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
20851        let boundary_tx = e.clone_dtod(&x)?;
20852        let boundary_rx = e.clone_dtod(&boundary_tx)?;
20853
20854        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
20855        let x =
20856            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
20857
20858        let mut hn = e.uninit(n_embd)?;
20859        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
20860        let h_seed = e.clone_dtod(&x)?;
20861        let mut ld = e.matmul(&self.output, &hn, 1)?;
20862        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
20863        e.softcap(&mut ld, cap, self.output.out_features())?;
20864        self.gemma4_suppress(e, &mut ld, 1)?;
20865        let logits = e.dtoh(&ld)?;
20866        cache.pos += 1;
20867        Ok((logits, h_seed))
20868    }
20869}
20870
20871// ============================ step35 (Step-3.7-Flash) ==================================
20872// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
20873// FAMILY and not a few branches inside the generic `full_attn*` chain:
20874//
20875//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
20876//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
20877//      shapes and the FA head counts would be wrong on 33 of 45 layers.
20878//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
20879//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
20880//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
20881//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
20882//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
20883//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
20884//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
20885//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
20886//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
20887//
20888// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
20889impl HybridModel {
20890    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
20891    /// synthesize a drafter or trunk layer from a neighboring class.
20892    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
20893        let geometry = self
20894            .cfg
20895            .layer_geometry(il as u32)
20896            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
20897        debug_assert_eq!(
20898            geometry.attention_gate,
20899            memra_gguf::config::AttentionGateKind::SeparateHead
20900        );
20901        geometry
20902    }
20903
20904    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
20905    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
20906    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
20907    ///
20908    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
20909    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
20910    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
20911    /// `cache`:
20912    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
20913    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
20914    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
20915    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
20916    ///     contract, lane/chunkinv-flip).
20917    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
20918    ///     q/k/v, no cache side effect.
20919    ///
20920    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
20921    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
20922    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
20923    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
20924    /// still contains must be masked per query. memra's window convention
20925    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
20926    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
20927    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
20928    ///
20929    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
20930    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
20931    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
20932    ///
20933    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
20934    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
20935    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
20936    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
20937    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
20938    /// hidden rows, and the generated text — a function of the chunk size:
20939    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
20940    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
20941    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
20942    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
20943    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
20944    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
20945    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
20946    ///   one-token change in a documented machine-config knob changed the answer.
20947    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
20948    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
20949    /// the same rows moves the logits by ~1.8.
20950    ///
20951    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
20952    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
20953    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
20954    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
20955    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
20956    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
20957    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
20958    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
20959    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
20960    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
20961    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
20962    /// those with t_kv <= win = 512.
20963    #[allow(clippy::too_many_arguments)]
20964    fn step35_attn_pre_wo(
20965        &self,
20966        e: &Engine,
20967        fa: &FullAttnLayer,
20968        mut g3: Vec<CudaSlice<f32>>,
20969        hg: Option<&CudaSlice<f32>>,
20970        gt_pre: Option<&CudaSlice<f32>>,
20971        pos_d: &CudaSlice<i32>,
20972        t: usize,
20973        cache: Option<&mut Cache>,
20974        il: usize,
20975        seq_end: usize,
20976    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
20977        let geometry = self.cfg.full_attention_geometry_at(il as u32);
20978        let hd = geometry.head_dim_k as usize;
20979        let nkv = geometry.n_head_kv as usize;
20980        let nh = geometry.n_head as usize;
20981        let rbase = geometry.rope_base;
20982        let scale = geometry.attention_scale();
20983        let swa = geometry.window.is_some();
20984        let eps = self.cfg.rms_eps;
20985        let win = geometry.window.unwrap_or(0) as usize;
20986        let n_rot = geometry.n_rot as usize;
20987
20988        let v = g3.pop().unwrap();
20989        let k0 = g3.pop().unwrap();
20990        let q0 = g3.pop().unwrap();
20991
20992        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
20993        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
20994        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
20995        let mut q = e.uninit(t * nh * hd)?;
20996        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
20997        let mut k = e.uninit(t * nkv * hd)?;
20998        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
20999        let ff = if geometry.rope_factors {
21000            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
21001        } else {
21002            None
21003        };
21004        #[cfg(debug_assertions)]
21005        if let Some(ff) = ff {
21006            crate::debug_assert_tensor_stream_device(
21007                ff,
21008                &e.stream(),
21009                "step35_attn_pre_wo.rope_freqs",
21010            );
21011        }
21012        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
21013
21014        let mut attn = e.uninit(t * nh * hd)?;
21015        match cache {
21016            Some(cache) => {
21017                let base_len = cache.kv[il].as_ref().unwrap().len;
21018                // Read per layer call, never in a measured default.
21019                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
21020                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
21021                let off = if swa {
21022                    let raw = base_len.saturating_sub(win - 1);
21023                    if legacy_tkv || legacy_calllocal {
21024                        raw
21025                    } else {
21026                        raw & !31usize
21027                    }
21028                } else {
21029                    0
21030                };
21031                {
21032                    let kvl = cache.kv[il].as_mut().unwrap();
21033                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
21034                    let write_row = e.prepare_kv_append(kvl, off, t)?;
21035                    e.append_kv_quantized_rows(
21036                        &k,
21037                        &v,
21038                        &mut kvl.k,
21039                        &mut kvl.v,
21040                        write_row,
21041                        t,
21042                        kvl.kv_dim_k,
21043                        kvl.kv_dim_v,
21044                        kvl.k_tok_bytes,
21045                        kvl.v_tok_bytes,
21046                        crate::Engine::kv_fp8_on(),
21047                    )?;
21048                    kvl.len += t;
21049                    let new_len = kvl.len as i32;
21050                    e.set_i32_one(&mut kvl.len_d, new_len)?;
21051                }
21052                let kvl = cache.kv[il].as_ref().unwrap();
21053                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
21054                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
21055                // unaligned view offset here. Both halves are load-bearing for the canaries:
21056                // on the FA default the predicate arms agree bitwise wherever they can differ
21057                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
21058                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
21059                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
21060                // on the current FA path: its tile grid starts at the chunk/call boundary.
21061                // SWA: trim the view to the oldest key any query in this chunk can reach —
21062                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
21063                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
21064                // kernel's online-softmax recurrence groups keys into BK tiles relative to
21065                // the VIEW START — so an unaligned off regroups the same absolute keys into
21066                // different tiles at different chunk sizes = different (m,l) rounding =
21067                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
21068                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
21069                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
21070                // size; the <=31 extra leading keys are older than EVERY query's window
21071                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
21072                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
21073                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
21074                // the floor arm's bits do not move either (gated: G2f, battery 2).
21075                let t_kv = base_len + t - off;
21076                let physical = kvl.physical_rows(off, off + t_kv)?;
21077                let k_view = e.view_u8_range(
21078                    &kvl.k,
21079                    physical.start * kvl.k_tok_bytes,
21080                    physical.end * kvl.k_tok_bytes,
21081                );
21082                let v_view = e.view_u8_range(
21083                    &kvl.v,
21084                    physical.start * kvl.v_tok_bytes,
21085                    physical.end * kvl.v_tok_bytes,
21086                );
21087                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
21088                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
21089                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
21090                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
21091                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
21092                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
21093                // construction, so the invariance assertion MUST break under it (the seam whose
21094                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
21095                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
21096                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
21097                // cached (probes flip it in-process). Never on in a measured default run.
21098                let swa_naive = if legacy_tkv {
21099                    t_kv > win
21100                } else {
21101                    seq_end > win
21102                };
21103                if swa && swa_naive {
21104                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
21105                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
21106                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
21107                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
21108                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
21109                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
21110                    // identically to the unwindowed one modulo the mask, which is the point.
21111                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
21112                    // selected on `seq_end` like every arm here, so the class is uniform for
21113                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
21114                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
21115                    // the f32 floor (the previous numeric config, kept as the A/B seam).
21116                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21117                        e.sdpa_naive_w_quantized_view(
21118                            &q,
21119                            &k_view,
21120                            &v_view,
21121                            &mut attn,
21122                            hd,
21123                            nh,
21124                            nkv,
21125                            t,
21126                            t_kv,
21127                            scale,
21128                            true,
21129                            win,
21130                            kvl.k_tok_bytes,
21131                            kvl.v_tok_bytes,
21132                        )?;
21133                    } else {
21134                        e.fa_prefill_view_ws_w_hd128(
21135                            &q,
21136                            &k_view,
21137                            &v_view,
21138                            &mut attn,
21139                            hd,
21140                            nh,
21141                            nkv,
21142                            t,
21143                            t_kv,
21144                            scale,
21145                            true,
21146                            win,
21147                            kvl.k_tok_bytes,
21148                            kvl.v_tok_bytes,
21149                        )?;
21150                    }
21151                } else if std::env::var("MEMRA_NOFA").is_ok() {
21152                    e.sdpa_naive_quantized_view(
21153                        &q,
21154                        &k_view,
21155                        &v_view,
21156                        &mut attn,
21157                        hd,
21158                        nh,
21159                        nkv,
21160                        t,
21161                        t_kv,
21162                        scale,
21163                        true,
21164                        kvl.k_tok_bytes,
21165                        kvl.v_tok_bytes,
21166                    )?;
21167                } else {
21168                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
21169                    // reach past the window, so the window mask is a no-op under causal and every
21170                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
21171                    // request either way, which is what makes the chunk size arithmetic-free.
21172                    e.fa_prefill_view_ws(
21173                        &q,
21174                        &k_view,
21175                        &v_view,
21176                        &mut attn,
21177                        hd,
21178                        nh,
21179                        nkv,
21180                        t,
21181                        t_kv,
21182                        scale,
21183                        true,
21184                        kvl.k_tok_bytes,
21185                        kvl.v_tok_bytes,
21186                        crate::Engine::kv_fp8_on(),
21187                    )?;
21188                }
21189            }
21190            None => {
21191                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
21192                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
21193                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
21194                // seq_end here too or it re-opens the same door.
21195                debug_assert_eq!(
21196                    seq_end, t,
21197                    "step35 cacheless prefill is monolithic (seq_end == t)"
21198                );
21199                if swa && seq_end > win {
21200                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
21201                } else if std::env::var("MEMRA_NOFA").is_ok() {
21202                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21203                } else {
21204                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
21205                }
21206            }
21207        }
21208
21209        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
21210        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
21211        let gw = fa
21212            .attn_gate
21213            .as_ref()
21214            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
21215        let gt_owned = if gt_pre.is_none() {
21216            Some(e.matmul(
21217                gw,
21218                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
21219                t,
21220            )?)
21221        } else {
21222            None
21223        };
21224        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
21225        let mut ag = e.uninit(t * nh * hd)?;
21226        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
21227        Ok(ag)
21228    }
21229
21230    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
21231    /// `forward_last`, t2probe). Post-`wo`.
21232    pub(crate) fn step35_attn(
21233        &self,
21234        e: &Engine,
21235        fa: &FullAttnLayer,
21236        h: &CudaSlice<f32>,
21237        pos_d: &CudaSlice<i32>,
21238        t: usize,
21239        il: usize,
21240    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21241        let g3 = match self.full_attn_tp_qkv(e, fa, h, t)? {
21242            Some(g3) => g3,
21243            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21244        };
21245        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
21246        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
21247        self.full_attn_o(e, fa, &ag, t)
21248    }
21249
21250    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
21251    /// resident quantized cache, attend through the cache view). Post-`wo`.
21252    ///
21253    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
21254    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
21255    /// own extent.
21256    #[allow(clippy::too_many_arguments)]
21257    pub(crate) fn step35_attn_prime(
21258        &self,
21259        e: &Engine,
21260        fa: &FullAttnLayer,
21261        h: &CudaSlice<f32>,
21262        hx: Option<&CudaSlice<u8>>,
21263        pos_d: &CudaSlice<i32>,
21264        t: usize,
21265        cache: &mut Cache,
21266        il: usize,
21267        seq_end: usize,
21268    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21269        if step_tp_prefill_enabled()? && fa.step_tp_qkv.is_some() {
21270            if hx.is_some() {
21271                return Err(
21272                    "rank-local Step prefill preserves BF16 activations and refuses the q8_1 \
21273                     pre-quantized prime path"
21274                        .into(),
21275                );
21276            }
21277            return self.step35_tp_prefill_attn_resident(e, fa, il, h, pos_d, t, cache, seq_end);
21278        }
21279        let g3 = if fa.step_tp_qkv.is_some() {
21280            if hx.is_some() {
21281                return Err(
21282                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
21283                     pre-quantized prime path"
21284                        .into(),
21285                );
21286            }
21287            self.full_attn_tp_qkv(e, fa, h, t)?
21288                .expect("Step Q/K/V TP disappeared after the presence check")
21289        } else {
21290            match hx {
21291                Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
21292                None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
21293            }
21294        };
21295        let ag =
21296            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
21297        self.full_attn_o(e, fa, &ag, t)
21298    }
21299
21300    fn ensure_step_tp_kv_cache(
21301        &self,
21302        e: &Engine,
21303        fa: &FullAttnLayer,
21304        il: usize,
21305        cache: &mut Cache,
21306    ) -> Result<bool, Box<dyn std::error::Error>> {
21307        let tp = fa
21308            .step_tp_qkv
21309            .as_ref()
21310            .ok_or("Step TP cache hydration lost its resident projections")?;
21311        let geometry = self.cfg.full_attention_geometry_at(il as u32);
21312        let window = geometry.window.map(|window| window as usize);
21313        let ranks = tp.runtime.devices().len();
21314        let head_dim = geometry.head_dim_k as usize;
21315        let kv_heads = geometry.n_head_kv as usize;
21316        let max_ctx = cache.max_ctx;
21317
21318        if cache.tp_kv[il].is_some() {
21319            return Ok(false);
21320        }
21321        let local = cache.kv[il]
21322            .as_ref()
21323            .ok_or_else(|| format!("Step TP layer {il} has no owning-stage KV cache"))?;
21324        if local.kv_dim_k != kv_heads * head_dim || local.kv_dim_v != kv_heads * head_dim {
21325            return Err(format!(
21326                "Step TP layer {il} local KV geometry k={} v={} != {}",
21327                local.kv_dim_k,
21328                local.kv_dim_v,
21329                kv_heads * head_dim
21330            )
21331            .into());
21332        }
21333        let resident_start = window
21334            .map(|window| local.len.saturating_sub(window.saturating_sub(1)) & !31usize)
21335            .unwrap_or(0);
21336        let resident_rows = local.len - resident_start;
21337        let physical = local.physical_rows(resident_start, local.len)?;
21338        let k_rows = if resident_rows == 0 {
21339            Vec::new()
21340        } else {
21341            e.dtoh_u8_view(&e.view_u8_range(
21342                &local.k,
21343                physical.start * local.k_tok_bytes,
21344                physical.end * local.k_tok_bytes,
21345            ))?
21346        };
21347        let v_rows = if resident_rows == 0 {
21348            Vec::new()
21349        } else {
21350            e.dtoh_u8_view(&e.view_u8_range(
21351                &local.v,
21352                physical.start * local.v_tok_bytes,
21353                physical.end * local.v_tok_bytes,
21354            ))?
21355        };
21356        let mut distributed = match window {
21357            Some(window) => tp.runtime.allocate_tp_swa_kv_cache(
21358                kv_heads * head_dim,
21359                kv_heads * head_dim,
21360                max_ctx,
21361                window,
21362            )?,
21363            None => tp.runtime.allocate_tp_kv_cache(
21364                kv_heads * head_dim,
21365                kv_heads * head_dim,
21366                max_ctx,
21367            )?,
21368        };
21369        if distributed.k_tok_bytes() * ranks != local.k_tok_bytes
21370            || distributed.v_tok_bytes() * ranks != local.v_tok_bytes
21371        {
21372            return Err(format!(
21373                "Step TP layer {il} distributed/local KV token bytes disagree: \
21374                 k={}x{ranks}/{} v={}x{ranks}/{}",
21375                distributed.k_tok_bytes(),
21376                local.k_tok_bytes,
21377                distributed.v_tok_bytes(),
21378                local.v_tok_bytes,
21379            )
21380            .into());
21381        }
21382        tp.runtime.hydrate_tp_kv_cache_from(
21383            &mut distributed,
21384            local.len,
21385            resident_start,
21386            &k_rows,
21387            &v_rows,
21388        )?;
21389        cache.tp_kv[il] = Some(distributed);
21390        Ok(true)
21391    }
21392
21393    #[allow(clippy::too_many_arguments)]
21394    #[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
21395    fn step35_tp_prefill_attn_resident(
21396        &self,
21397        e: &Engine,
21398        fa: &FullAttnLayer,
21399        il: usize,
21400        h: &CudaSlice<f32>,
21401        pos_d: &CudaSlice<i32>,
21402        tokens: usize,
21403        cache: &mut Cache,
21404        seq_end: usize,
21405    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21406        let tp = fa
21407            .step_tp_qkv
21408            .as_ref()
21409            .ok_or("Step TP prefill lost its resident projections")?;
21410        let attention = tp
21411            .attention
21412            .as_ref()
21413            .ok_or("Step TP prefill lost its resident attention auxiliaries")?;
21414        let ranks = tp.runtime.devices().len();
21415        if !step_tp_prefill_shape(
21416            true,
21417            tokens,
21418            ranks,
21419            tp.runtime.native_p2p(),
21420            true,
21421            crate::Engine::kv_fp8_on(),
21422        ) {
21423            return Err(format!(
21424                "rank-local Step prefill requires tokens>={PRIME_MIN_T}, TP2/TP4 native P2P, \
21425                 rank-local attention, and q8_0/q5_1 KV; got tokens={tokens} ranks={ranks} \
21426                 native_p2p={} fp8_kv={}",
21427                tp.runtime.native_p2p(),
21428                crate::Engine::kv_fp8_on(),
21429            )
21430            .into());
21431        }
21432        for seam in [
21433            "MEMRA_STEP35_SWA_TKV",
21434            "MEMRA_PRIME_CALLLOCAL",
21435            "MEMRA_PRIME_F32CHUNK0",
21436        ] {
21437            if std::env::var(seam).as_deref() == Ok("1") {
21438                return Err(format!(
21439                    "rank-local Step prefill has not qualified the legacy seam {seam}=1"
21440                )
21441                .into());
21442            }
21443        }
21444
21445        let geometry = self.cfg.full_attention_geometry_at(il as u32);
21446        let window = geometry.window.map(|window| window as usize);
21447        let head_dim = geometry.head_dim_k as usize;
21448        let heads = geometry.n_head as usize;
21449        let kv_heads = geometry.n_head_kv as usize;
21450        if heads % ranks != 0 || kv_heads % ranks != 0 {
21451            return Err(format!(
21452                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
21453            )
21454            .into());
21455        }
21456        let local_heads = heads / ranks;
21457        let local_kv_heads = kv_heads / ranks;
21458        let local_kv_dim = local_kv_heads * head_dim;
21459        let hidden = self.cfg.n_embd as usize;
21460        let expected_input = tokens
21461            .checked_mul(hidden)
21462            .ok_or("Step TP prefill input size overflow")?;
21463        if h.len() < expected_input {
21464            return Err(format!(
21465                "Step TP prefill input {} is shorter than {tokens}x{hidden}",
21466                h.len()
21467            )
21468            .into());
21469        }
21470        let positions = e.dtoh_i32(pos_d)?;
21471        if positions.len() != tokens {
21472            return Err(format!(
21473                "rank-local Step prefill positions {} != tokens {tokens}",
21474                positions.len()
21475            )
21476            .into());
21477        }
21478
21479        let mut active_input = e.uninit(expected_input)?;
21480        e.copy_view_into(
21481            &mut active_input,
21482            0,
21483            &h.slice(0..expected_input),
21484            expected_input,
21485        )?;
21486        let mut input = tp.runtime.allocate_replicated_device_rows(tokens, hidden)?;
21487        // PRODUCER FENCE (2026-08-20 flake fix): active_input was written on the MODEL engine's
21488        // stream; the refresh below reads it from the runtime root engine's stream (same device,
21489        // different stream). Unfenced, the peer read can overtake the in-flight copy — the
21490        // layer-count-amplified arm of the boot flake.
21491        e.stream().synchronize()?;
21492        tp.runtime
21493            .refresh_replicated_device_rows_from_root(&mut input, &active_input)?;
21494        let q_raw = tp
21495            .runtime
21496            .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &input)?;
21497        let k_raw = tp
21498            .runtime
21499            .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &input)?;
21500        let v_raw = tp
21501            .runtime
21502            .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &input)?;
21503        let mut q = Vec::with_capacity(ranks);
21504        let mut k = Vec::with_capacity(ranks);
21505        for rank in 0..ranks {
21506            let engine = tp
21507                .runtime
21508                .rank_engine(rank)
21509                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21510            let _main = engine.gpu.enter_main()?;
21511            let mut q_rank = engine.uninit(tokens * local_heads * head_dim)?;
21512            engine.rms_norm(
21513                &q_raw[rank],
21514                &attention.q_norm[rank],
21515                &mut q_rank,
21516                head_dim,
21517                tokens * local_heads,
21518                self.cfg.rms_eps,
21519            )?;
21520            let mut k_rank = engine.uninit(tokens * local_kv_dim)?;
21521            engine.rms_norm(
21522                &k_raw[rank],
21523                &attention.k_norm[rank],
21524                &mut k_rank,
21525                head_dim,
21526                tokens * local_kv_heads,
21527                self.cfg.rms_eps,
21528            )?;
21529            let position = engine.htod_i32(&positions)?;
21530            let rope_freqs = if geometry.rope_factors {
21531                self.step35_aux
21532                    .as_ref()
21533                    .and_then(|aux| aux.rope_freqs(engine))
21534            } else {
21535                None
21536            };
21537            engine.rope_neox2(
21538                &mut q_rank,
21539                &mut k_rank,
21540                &position,
21541                head_dim,
21542                geometry.n_rot as usize,
21543                local_heads,
21544                local_kv_heads,
21545                tokens,
21546                geometry.rope_base,
21547                1.0,
21548                rope_freqs,
21549            )?;
21550            q.push(q_rank);
21551            k.push(k_rank);
21552        }
21553
21554        let gate_weight = fa
21555            .attn_gate
21556            .as_ref()
21557            .ok_or("step35 layer is missing attn_gate.weight")?;
21558        let gate = e.dtoh(&e.matmul(gate_weight, h, tokens)?)?;
21559        if gate.len() != tokens * heads {
21560            return Err(format!(
21561                "Step TP layer {il} gate output {} != {tokens}x{heads}",
21562                gate.len()
21563            )
21564            .into());
21565        }
21566
21567        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
21568        let base_len = cache.kv[il]
21569            .as_ref()
21570            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
21571            .len;
21572        let distributed = cache.tp_kv[il]
21573            .as_ref()
21574            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
21575        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
21576            return Err(format!(
21577                "Step TP layer {il} cache lengths diverged before prefill: \
21578                 local={base_len} distributed={}/{}",
21579                distributed.committed_len(),
21580                distributed.staged_len()
21581            )
21582            .into());
21583        }
21584        let target_len = base_len
21585            .checked_add(tokens)
21586            .ok_or("Step TP prefill cache length overflow")?;
21587        if target_len > cache.max_ctx {
21588            return Err(format!(
21589                "Step TP layer {il} prefill exceeds cache: {base_len}+{tokens}>{}",
21590                cache.max_ctx
21591            )
21592            .into());
21593        }
21594        if seq_end < target_len {
21595            return Err(format!(
21596                "Step TP layer {il} request end {seq_end} precedes chunk end {target_len}"
21597            )
21598            .into());
21599        }
21600
21601        let transaction = cache.tp_kv[il]
21602            .as_mut()
21603            .expect("distributed cache checked above")
21604            .begin_transaction()?;
21605        if let Err(error) = tp.runtime.append_tp_kv_transaction(
21606            cache.tp_kv[il]
21607                .as_mut()
21608                .expect("distributed cache checked above"),
21609            transaction,
21610            &k,
21611            &v_raw,
21612            tokens,
21613        ) {
21614            let _ = tp.runtime.rollback_tp_kv_transaction(
21615                cache.tp_kv[il]
21616                    .as_mut()
21617                    .expect("distributed cache checked above"),
21618                transaction,
21619            );
21620            return Err(error);
21621        }
21622
21623        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21624            let distributed = cache.tp_kv[il]
21625                .as_ref()
21626                .expect("distributed cache checked above");
21627            let staged_len = distributed.staged_len();
21628            let view_start = window
21629                .map(|window| base_len.saturating_sub(window.saturating_sub(1)) & !31usize)
21630                .unwrap_or(0);
21631            let physical = distributed.physical_range(view_start, staged_len)?;
21632            let t_kv = staged_len - view_start;
21633            let swa_naive = window.is_some_and(|window| seq_end > window);
21634            let mut gated = Vec::with_capacity(ranks);
21635            #[allow(clippy::needless_range_loop)]
21636            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
21637            for rank in 0..ranks {
21638                let engine = tp
21639                    .runtime
21640                    .rank_engine(rank)
21641                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21642                let _main = engine.gpu.enter_main()?;
21643                let rank_cache = distributed
21644                    .rank(rank)
21645                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
21646                let k_view = engine.view_u8_range(
21647                    rank_cache.k(),
21648                    physical.start * distributed.k_tok_bytes(),
21649                    physical.end * distributed.k_tok_bytes(),
21650                );
21651                let v_view = engine.view_u8_range(
21652                    rank_cache.v(),
21653                    physical.start * distributed.v_tok_bytes(),
21654                    physical.end * distributed.v_tok_bytes(),
21655                );
21656                let mut attention_out = engine.uninit(tokens * local_heads * head_dim)?;
21657                if swa_naive {
21658                    let window = window.expect("SWA predicate requires a window");
21659                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
21660                        engine.sdpa_naive_w_quantized_view(
21661                            &q[rank],
21662                            &k_view,
21663                            &v_view,
21664                            &mut attention_out,
21665                            head_dim,
21666                            local_heads,
21667                            local_kv_heads,
21668                            tokens,
21669                            t_kv,
21670                            geometry.attention_scale(),
21671                            true,
21672                            window,
21673                            distributed.k_tok_bytes(),
21674                            distributed.v_tok_bytes(),
21675                        )?;
21676                    } else {
21677                        engine.fa_prefill_view_ws_w_hd128(
21678                            &q[rank],
21679                            &k_view,
21680                            &v_view,
21681                            &mut attention_out,
21682                            head_dim,
21683                            local_heads,
21684                            local_kv_heads,
21685                            tokens,
21686                            t_kv,
21687                            geometry.attention_scale(),
21688                            true,
21689                            window,
21690                            distributed.k_tok_bytes(),
21691                            distributed.v_tok_bytes(),
21692                        )?;
21693                    }
21694                } else if std::env::var("MEMRA_NOFA").is_ok() {
21695                    engine.sdpa_naive_quantized_view(
21696                        &q[rank],
21697                        &k_view,
21698                        &v_view,
21699                        &mut attention_out,
21700                        head_dim,
21701                        local_heads,
21702                        local_kv_heads,
21703                        tokens,
21704                        t_kv,
21705                        geometry.attention_scale(),
21706                        true,
21707                        distributed.k_tok_bytes(),
21708                        distributed.v_tok_bytes(),
21709                    )?;
21710                } else {
21711                    engine.fa_prefill_view_ws(
21712                        &q[rank],
21713                        &k_view,
21714                        &v_view,
21715                        &mut attention_out,
21716                        head_dim,
21717                        local_heads,
21718                        local_kv_heads,
21719                        tokens,
21720                        t_kv,
21721                        geometry.attention_scale(),
21722                        true,
21723                        distributed.k_tok_bytes(),
21724                        distributed.v_tok_bytes(),
21725                        false,
21726                    )?;
21727                }
21728
21729                let gate_start = rank * local_heads;
21730                let mut gate_rank = Vec::with_capacity(tokens * local_heads);
21731                for token in 0..tokens {
21732                    let start = token * heads + gate_start;
21733                    gate_rank.extend_from_slice(&gate[start..start + local_heads]);
21734                }
21735                let gate_rank = engine.htod(&gate_rank)?;
21736                let mut gated_rank = engine.uninit(tokens * local_heads * head_dim)?;
21737                engine.attn_head_gate(
21738                    &attention_out,
21739                    &gate_rank,
21740                    &mut gated_rank,
21741                    None,
21742                    head_dim,
21743                    local_heads,
21744                    tokens,
21745                )?;
21746                gated.push(gated_rank);
21747            }
21748            for rank in 1..ranks {
21749                let engine = tp
21750                    .runtime
21751                    .rank_engine(rank)
21752                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
21753                let _main = engine.gpu.enter_main()?;
21754                engine.stream().synchronize()?;
21755            }
21756
21757            let (output, k_shadow, v_shadow) = if tp.runtime.bulk_p2p() {
21758                let output = tp
21759                    .runtime
21760                    .step_bf16_row_parallel_resident_root_device(&tp.o, &gated, tokens)?;
21761                let k_shadow =
21762                    tp.runtime
21763                        .gather_native_column_shards_device(&k, tokens, local_kv_dim)?;
21764                let v_shadow =
21765                    tp.runtime
21766                        .gather_native_column_shards_device(&v_raw, tokens, local_kv_dim)?;
21767                let root = tp
21768                    .runtime
21769                    .rank_engine(0)
21770                    .ok_or("Step TP prefill lost its root engine")?;
21771                let _main = root.gpu.enter_main()?;
21772                root.stream().synchronize()?;
21773                (output, k_shadow, v_shadow)
21774            } else {
21775                let attention = tp.runtime.gather_native_column_shards(
21776                    &gated,
21777                    tokens,
21778                    local_heads * head_dim,
21779                )?;
21780                let output = tp
21781                    .runtime
21782                    .step_bf16_row_parallel_resident_native(&tp.o, &attention, tokens)?;
21783                let k_shadow = tp
21784                    .runtime
21785                    .gather_native_column_shards(&k, tokens, local_kv_dim)?;
21786                let v_shadow =
21787                    tp.runtime
21788                        .gather_native_column_shards(&v_raw, tokens, local_kv_dim)?;
21789                (e.htod(&output)?, e.htod(&k_shadow)?, e.htod(&v_shadow)?)
21790            };
21791            let local = cache.kv[il]
21792                .as_mut()
21793                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
21794            if local.len != base_len {
21795                return Err(format!(
21796                    "Step TP layer {il} local cache changed during prefill: \
21797                     len={} base={base_len}",
21798                    local.len
21799                )
21800                .into());
21801            }
21802            let retain_from = window
21803                .map(|window| {
21804                    let staged_retain = staged_len.saturating_sub(window) & !31usize;
21805                    let rollback_retain =
21806                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
21807                    staged_retain.min(rollback_retain)
21808                })
21809                .unwrap_or(0);
21810            let write_row = e.prepare_kv_append(local, retain_from, tokens)?;
21811            e.append_kv_quantized_rows(
21812                &k_shadow,
21813                &v_shadow,
21814                &mut local.k,
21815                &mut local.v,
21816                write_row,
21817                tokens,
21818                local.kv_dim_k,
21819                local.kv_dim_v,
21820                local.k_tok_bytes,
21821                local.v_tok_bytes,
21822                false,
21823            )?;
21824            local.len = staged_len;
21825            e.set_i32_one(&mut local.len_d, staged_len as i32)?;
21826            Ok(output)
21827        })();
21828
21829        let output = match staged {
21830            Ok(output) => output,
21831            Err(error) => {
21832                let _ = tp.runtime.rollback_tp_kv_transaction(
21833                    cache.tp_kv[il]
21834                        .as_mut()
21835                        .expect("distributed cache checked above"),
21836                    transaction,
21837                );
21838                if let Some(local) = cache.kv[il].as_mut() {
21839                    local.len = base_len;
21840                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
21841                }
21842                return Err(error);
21843            }
21844        };
21845        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
21846            cache.tp_kv[il]
21847                .as_mut()
21848                .expect("distributed cache checked above"),
21849            transaction,
21850            tokens,
21851        ) {
21852            let _ = tp.runtime.rollback_tp_kv_transaction(
21853                cache.tp_kv[il]
21854                    .as_mut()
21855                    .expect("distributed cache checked above"),
21856                transaction,
21857            );
21858            let local = cache.kv[il].as_mut().expect("local cache checked above");
21859            local.len = base_len;
21860            e.set_i32_one(&mut local.len_d, base_len as i32)?;
21861            return Err(error);
21862        }
21863
21864        let committed = cache.tp_kv[il]
21865            .as_ref()
21866            .expect("distributed cache checked above")
21867            .committed_len();
21868        let local_len = cache.kv[il]
21869            .as_ref()
21870            .expect("local cache checked above")
21871            .len;
21872        if committed != local_len {
21873            return Err(format!(
21874                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
21875            )
21876            .into());
21877        }
21878        eprintln!(
21879            "[step-tp-prefill-attn] execute layer={} devices={:?} tokens={tokens} \
21880             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
21881             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
21882             attention_scope={} input_path=root-device-replicated gate_tensor_parallel=false \
21883             gate_shards=host-canonical o_tensor_parallel=true local_cache_shadow=true \
21884             cache_commit=chunk transport={} native_p2p=true bulk_p2p={} \
21885             output={} performance_claim=false",
21886            tp.layer,
21887            tp.devices,
21888            hydrated,
21889            if window.is_some() {
21890                "rank-local-swa-ring"
21891            } else {
21892                "rank-local-global"
21893            },
21894            tp.runtime.transport_label(),
21895            tp.runtime.bulk_p2p(),
21896            if tp.runtime.bulk_p2p() {
21897                "root-device"
21898            } else {
21899                "root-readback"
21900            },
21901        );
21902        Ok(output)
21903    }
21904
21905    fn step35_tp_decode_attn_resident(
21906        &self,
21907        e: &Engine,
21908        fa: &FullAttnLayer,
21909        il: usize,
21910        h: &CudaSlice<f32>,
21911        pos_d: &CudaSlice<i32>,
21912        cache: &mut Cache,
21913    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21914        // MEMRA_STEP_TP_TIMING=1: cumulative wall-clock of the rank-local TP attention decode,
21915        // printed every ~10 decode steps' worth of layers — the wall-decomposition twin of the
21916        // nvfp4-dev-routes counter.
21917        static ATTN_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21918        static ATTN_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21919        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
21920        let started = timing.then(std::time::Instant::now);
21921        let result = if crate::tp::step_tp_decode_v2_enabled()? {
21922            self.step35_tp_decode_attn_resident_v2(e, fa, il, h, pos_d, cache)
21923        } else {
21924            self.step35_tp_decode_attn_resident_inner(e, fa, il, h, pos_d, cache)
21925        };
21926        if let Some(started) = started {
21927            use std::sync::atomic::Ordering;
21928            let ns = ATTN_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
21929                + started.elapsed().as_nanos() as u64;
21930            let calls = ATTN_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
21931            if calls.is_multiple_of(430) {
21932                eprintln!(
21933                    "[step-tp-attn-timing] calls={calls} total_ms={:.1} avg_us={:.1}",
21934                    ns as f64 / 1.0e6,
21935                    ns as f64 / calls as f64 / 1.0e3,
21936                );
21937            }
21938        }
21939        result
21940    }
21941
21942    #[allow(clippy::too_many_arguments)]
21943    fn step35_tp_decode_attn_resident_inner(
21944        &self,
21945        e: &Engine,
21946        fa: &FullAttnLayer,
21947        il: usize,
21948        h: &CudaSlice<f32>,
21949        pos_d: &CudaSlice<i32>,
21950        cache: &mut Cache,
21951    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
21952        // MEMRA_STEP_TP_TIMING=1 phase decomposition of the 1550us/layer decode wall. Each lap
21953        // drains every stream so queued async work is billed to the phase that queued it — the
21954        // drains perturb absolute wall, but v1 already ends most phases on a host sync, so the
21955        // relative split is honest. Timing OFF is the measured configuration: zero extra syncs.
21956        static T_POS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21957        static T_QKV: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21958        static T_NORMROPE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21959        static T_GATE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21960        static T_APPEND: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21961        static T_ATTN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21962        static T_OPROJ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21963        static T_SHADOW: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21964        static T_PHASE_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
21965        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
21966        #[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
21967        fn lap(
21968            runtime: &crate::tp::TpE4m3HostBounce,
21969            e: &Engine,
21970            timer: &std::sync::atomic::AtomicU64,
21971            started: &mut Option<std::time::Instant>,
21972        ) -> Result<(), Box<dyn std::error::Error>> {
21973            let Some(start) = started.as_mut() else {
21974                return Ok(());
21975            };
21976            for rank in 0..runtime.devices().len() {
21977                if let Some(engine) = runtime.rank_engine(rank) {
21978                    let _main = engine.gpu.enter_main()?;
21979                    engine.stream().synchronize()?;
21980                }
21981            }
21982            e.stream().synchronize()?;
21983            timer.fetch_add(
21984                start.elapsed().as_nanos() as u64,
21985                std::sync::atomic::Ordering::Relaxed,
21986            );
21987            *start = std::time::Instant::now();
21988            Ok(())
21989        }
21990        let tp = fa
21991            .step_tp_qkv
21992            .as_ref()
21993            .ok_or("Step TP decode lost its resident projections")?;
21994        let attention = tp
21995            .attention
21996            .as_ref()
21997            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
21998        if !tp.runtime.native_p2p() {
21999            return Err("rank-local Step attention requires native P2P".into());
22000        }
22001        if crate::Engine::kv_fp8_on() {
22002            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
22003        }
22004
22005        let geometry = self.step35_geom(il);
22006        let window = geometry.window.map(|window| window as usize);
22007        let ranks = tp.runtime.devices().len();
22008        let head_dim = geometry.head_dim_k as usize;
22009        let heads = geometry.n_head as usize;
22010        let kv_heads = geometry.n_head_kv as usize;
22011        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
22012            return Err(format!(
22013                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
22014            )
22015            .into());
22016        }
22017        let local_heads = heads / ranks;
22018        let local_kv_heads = kv_heads / ranks;
22019        let local_kv_dim = local_kv_heads * head_dim;
22020        let max_ctx = cache.max_ctx;
22021
22022        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
22023
22024        let base_len = cache.kv[il]
22025            .as_ref()
22026            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
22027            .len;
22028        let distributed = cache.tp_kv[il]
22029            .as_ref()
22030            .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
22031        if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
22032            return Err(format!(
22033                "Step TP layer {il} cache lengths diverged before decode: \
22034                 local={base_len} distributed={}/{}",
22035                distributed.committed_len(),
22036                distributed.staged_len()
22037            )
22038            .into());
22039        }
22040
22041        let mut lap_start = timing.then(std::time::Instant::now);
22042        let positions = e.dtoh_i32(pos_d)?;
22043        if positions.len() != 1 {
22044            return Err(format!(
22045                "rank-local Step decode requires one position, got {}",
22046                positions.len()
22047            )
22048            .into());
22049        }
22050        lap(&tp.runtime, e, &T_POS, &mut lap_start)?;
22051        let (q_raw, k_raw, v_raw, input_path) = if let Some(decode_input) =
22052            attention.decode_input.as_ref()
22053        {
22054            let mut decode_input = decode_input
22055                .lock()
22056                .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
22057            // PRODUCER FENCE (2026-08-20 flake fix): h is the layer hidden written on the MODEL
22058            // engine's stream; the refresh reads it from the runtime root engine's stream. This
22059            // per-layer-per-token seam is the layer-count-amplified arm of the boot flake.
22060            e.stream().synchronize()?;
22061            tp.runtime
22062                .refresh_replicated_device_rows_from_root(&mut decode_input, h)?;
22063            let q_raw = tp
22064                .runtime
22065                .bf16_column_parallel_resident_replicated_device_shards(&tp.q, &decode_input)?;
22066            let k_raw = tp
22067                .runtime
22068                .bf16_column_parallel_resident_replicated_device_shards(&tp.k, &decode_input)?;
22069            let v_raw = tp
22070                .runtime
22071                .bf16_column_parallel_resident_replicated_device_shards(&tp.v, &decode_input)?;
22072            (q_raw, k_raw, v_raw, "root-device-replicated")
22073        } else {
22074            let activation = e.dtoh(h)?;
22075            let q_raw =
22076                tp.runtime
22077                    .bf16_column_parallel_resident_device_shards(&tp.q, &activation, 1)?;
22078            let k_raw =
22079                tp.runtime
22080                    .bf16_column_parallel_resident_device_shards(&tp.k, &activation, 1)?;
22081            let v_raw =
22082                tp.runtime
22083                    .bf16_column_parallel_resident_device_shards(&tp.v, &activation, 1)?;
22084            (q_raw, k_raw, v_raw, "host-replicated")
22085        };
22086        lap(&tp.runtime, e, &T_QKV, &mut lap_start)?;
22087        let mut q = Vec::with_capacity(ranks);
22088        let mut k = Vec::with_capacity(ranks);
22089        for rank in 0..ranks {
22090            let engine = tp
22091                .runtime
22092                .rank_engine(rank)
22093                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22094            let _main = engine.gpu.enter_main()?;
22095            let mut q_rank = engine.uninit(local_heads * head_dim)?;
22096            engine.rms_norm(
22097                &q_raw[rank],
22098                &attention.q_norm[rank],
22099                &mut q_rank,
22100                head_dim,
22101                local_heads,
22102                self.cfg.rms_eps,
22103            )?;
22104            let mut k_rank = engine.uninit(local_kv_dim)?;
22105            engine.rms_norm(
22106                &k_raw[rank],
22107                &attention.k_norm[rank],
22108                &mut k_rank,
22109                head_dim,
22110                local_kv_heads,
22111                self.cfg.rms_eps,
22112            )?;
22113            let position = engine.htod_i32(&positions)?;
22114            let rope_freqs = if geometry.rope_factors {
22115                self.step35_aux
22116                    .as_ref()
22117                    .and_then(|aux| aux.rope_freqs(engine))
22118            } else {
22119                None
22120            };
22121            engine.rope_neox2(
22122                &mut q_rank,
22123                &mut k_rank,
22124                &position,
22125                head_dim,
22126                geometry.n_rot as usize,
22127                local_heads,
22128                local_kv_heads,
22129                1,
22130                geometry.rope_base,
22131                1.0,
22132                rope_freqs,
22133            )?;
22134            q.push(q_rank);
22135            k.push(k_rank);
22136        }
22137        lap(&tp.runtime, e, &T_NORMROPE, &mut lap_start)?;
22138
22139        let gate_weight = fa
22140            .attn_gate
22141            .as_ref()
22142            .ok_or("step35 layer is missing attn_gate.weight")?;
22143        let gate = e.matmul(gate_weight, h, 1)?;
22144        let gate = e.dtoh(&gate)?;
22145        if gate.len() != heads {
22146            return Err(format!("Step TP layer {il} gate output {} != {heads}", gate.len()).into());
22147        }
22148        lap(&tp.runtime, e, &T_GATE, &mut lap_start)?;
22149
22150        let transaction = cache.tp_kv[il]
22151            .as_mut()
22152            .expect("distributed cache checked above")
22153            .begin_transaction()?;
22154        if let Err(error) = tp.runtime.append_tp_kv_transaction(
22155            cache.tp_kv[il]
22156                .as_mut()
22157                .expect("distributed cache checked above"),
22158            transaction,
22159            &k,
22160            &v_raw,
22161            1,
22162        ) {
22163            let _ = tp.runtime.rollback_tp_kv_transaction(
22164                cache.tp_kv[il]
22165                    .as_mut()
22166                    .expect("distributed cache checked above"),
22167                transaction,
22168            );
22169            return Err(error);
22170        }
22171        lap(&tp.runtime, e, &T_APPEND, &mut lap_start)?;
22172
22173        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22174            let distributed = cache.tp_kv[il]
22175                .as_ref()
22176                .expect("distributed cache checked above");
22177            let staged_len = distributed.staged_len();
22178            let view_start = window
22179                .map(|window| staged_len.saturating_sub(window))
22180                .unwrap_or(0);
22181            let physical = distributed.physical_range(view_start, staged_len)?;
22182            let t_kv = staged_len - view_start;
22183            let mut gated = Vec::with_capacity(ranks);
22184            #[allow(clippy::needless_range_loop)]
22185            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
22186            for rank in 0..ranks {
22187                let engine = tp
22188                    .runtime
22189                    .rank_engine(rank)
22190                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
22191                let _main = engine.gpu.enter_main()?;
22192                let rank_cache = distributed
22193                    .rank(rank)
22194                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
22195                let k_view = engine.view_u8_range(
22196                    rank_cache.k(),
22197                    physical.start * distributed.k_tok_bytes(),
22198                    physical.end * distributed.k_tok_bytes(),
22199                );
22200                let v_view = engine.view_u8_range(
22201                    rank_cache.v(),
22202                    physical.start * distributed.v_tok_bytes(),
22203                    physical.end * distributed.v_tok_bytes(),
22204                );
22205                let mut attention_out = engine.uninit(local_heads * head_dim)?;
22206                engine.fa_decode_kvmod(
22207                    &q[rank],
22208                    &k_view,
22209                    &v_view,
22210                    &mut attention_out,
22211                    head_dim,
22212                    local_heads,
22213                    local_kv_heads,
22214                    t_kv,
22215                    geometry.attention_scale(),
22216                    distributed.k_tok_bytes(),
22217                    distributed.v_tok_bytes(),
22218                    false,
22219                )?;
22220                let gate_start = rank * local_heads;
22221                let gate_rank = engine.htod(&gate[gate_start..gate_start + local_heads])?;
22222                let mut gated_rank = engine.uninit(local_heads * head_dim)?;
22223                engine.attn_head_gate(
22224                    &attention_out,
22225                    &gate_rank,
22226                    &mut gated_rank,
22227                    None,
22228                    head_dim,
22229                    local_heads,
22230                    1,
22231                )?;
22232                gated.push(gated_rank);
22233            }
22234            lap(&tp.runtime, e, &T_ATTN, &mut lap_start)?;
22235
22236            let gathered =
22237                tp.runtime
22238                    .gather_native_column_shards(&gated, 1, local_heads * head_dim)?;
22239            let output = tp
22240                .runtime
22241                .step_bf16_row_parallel_resident_native(&tp.o, &gathered, 1)?;
22242            let output = e.htod(&output)?;
22243            lap(&tp.runtime, e, &T_OPROJ, &mut lap_start)?;
22244
22245            let k_shadow = tp
22246                .runtime
22247                .gather_native_column_shards(&k, 1, local_kv_dim)?;
22248            let v_shadow = tp
22249                .runtime
22250                .gather_native_column_shards(&v_raw, 1, local_kv_dim)?;
22251            let k_shadow = e.htod(&k_shadow)?;
22252            let v_shadow = e.htod(&v_shadow)?;
22253            let local = cache.kv[il]
22254                .as_mut()
22255                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
22256            if local.len != base_len || base_len + 1 > max_ctx {
22257                return Err(format!(
22258                    "Step TP layer {il} local cache changed during decode: \
22259                     len={} base={base_len} max={max_ctx}",
22260                    local.len
22261                )
22262                .into());
22263            }
22264            let retain_from = window
22265                .map(|window| {
22266                    let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
22267                    let rollback_retain =
22268                        base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
22269                    staged_retain.min(rollback_retain)
22270                })
22271                .unwrap_or(0);
22272            let write_row = e.prepare_kv_append(local, retain_from, 1)?;
22273            e.append_kv_quantized(
22274                &k_shadow,
22275                &v_shadow,
22276                &mut local.k,
22277                &mut local.v,
22278                write_row,
22279                local.kv_dim_k,
22280                local.kv_dim_v,
22281                local.k_tok_bytes,
22282                local.v_tok_bytes,
22283                false,
22284            )?;
22285            local.len = base_len + 1;
22286            e.set_i32_one(&mut local.len_d, local.len as i32)?;
22287            Ok(output)
22288        })();
22289
22290        let output = match staged {
22291            Ok(output) => output,
22292            Err(error) => {
22293                let _ = tp.runtime.rollback_tp_kv_transaction(
22294                    cache.tp_kv[il]
22295                        .as_mut()
22296                        .expect("distributed cache checked above"),
22297                    transaction,
22298                );
22299                if let Some(local) = cache.kv[il].as_mut() {
22300                    local.len = base_len;
22301                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
22302                }
22303                return Err(error);
22304            }
22305        };
22306        if let Err(error) = tp.runtime.commit_tp_kv_transaction(
22307            cache.tp_kv[il]
22308                .as_mut()
22309                .expect("distributed cache checked above"),
22310            transaction,
22311            1,
22312        ) {
22313            let _ = tp.runtime.rollback_tp_kv_transaction(
22314                cache.tp_kv[il]
22315                    .as_mut()
22316                    .expect("distributed cache checked above"),
22317                transaction,
22318            );
22319            let local = cache.kv[il].as_mut().expect("local cache checked above");
22320            local.len = base_len;
22321            e.set_i32_one(&mut local.len_d, base_len as i32)?;
22322            return Err(error);
22323        }
22324
22325        let committed = cache.tp_kv[il]
22326            .as_ref()
22327            .expect("distributed cache checked above")
22328            .committed_len();
22329        let local_len = cache.kv[il]
22330            .as_ref()
22331            .expect("local cache checked above")
22332            .len;
22333        if committed != local_len {
22334            return Err(format!(
22335                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
22336            )
22337            .into());
22338        }
22339        lap(&tp.runtime, e, &T_SHADOW, &mut lap_start)?;
22340        if timing {
22341            use std::sync::atomic::Ordering;
22342            let calls = T_PHASE_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
22343            if calls.is_multiple_of(430) {
22344                let avg = |t: &std::sync::atomic::AtomicU64| {
22345                    t.load(Ordering::Relaxed) as f64 / calls as f64 / 1.0e3
22346                };
22347                eprintln!(
22348                    "[step-tp-attn-phase] calls={calls} avg_us pos={:.1} qkv={:.1} \
22349                     normrope={:.1} gate={:.1} append={:.1} attn={:.1} oproj={:.1} shadow={:.1}",
22350                    avg(&T_POS),
22351                    avg(&T_QKV),
22352                    avg(&T_NORMROPE),
22353                    avg(&T_GATE),
22354                    avg(&T_APPEND),
22355                    avg(&T_ATTN),
22356                    avg(&T_OPROJ),
22357                    avg(&T_SHADOW),
22358                );
22359            }
22360        }
22361        eprintln!(
22362            "[step-tp-attn] execute layer={} devices={:?} tokens=1 \
22363             qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
22364             kv_cache_distributed=true kv_cache_hydrated={} attention_tensor_parallel=true \
22365             attention_scope={} input_path={} kv_physical_rows={} \
22366             gate_tensor_parallel=false gate_shards=host-canonical o_tensor_parallel=true \
22367             local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
22368             bulk_p2p={} output=root-readback performance_claim=false",
22369            tp.layer,
22370            tp.devices,
22371            hydrated,
22372            if window.is_some() {
22373                "rank-local-swa-ring"
22374            } else {
22375                "rank-local-global"
22376            },
22377            input_path,
22378            cache.tp_kv[il]
22379                .as_ref()
22380                .expect("distributed cache checked above")
22381                .physical_capacity(),
22382            tp.runtime.transport_label(),
22383            tp.runtime.bulk_p2p(),
22384        );
22385        Ok(output)
22386    }
22387
22388    /// v2 rank-local decode attention (MEMRA_STEP_TP_DECODE_V2): the same kernels, operand
22389    /// values, and canonical reduction order as `step35_tp_decode_attn_resident_inner`,
22390    /// restructured onto a persistent per-runtime workspace with evented cross-stream ordering.
22391    /// Per layer per token this path performs exactly one cuMemAlloc (the returned e-context
22392    /// output row), no host round-trip, and no host stream synchronize — the phase timers
22393    /// measured v1 spending 81% of its 1550us/layer wall on those three classes.
22394    #[allow(clippy::too_many_arguments)]
22395    /// T-COLUMN verify precompute for layer `il`: weight-amortized QKV(+gate) over the T
22396    /// verify columns into the ws slabs (per-column rope/append/fa run later through the
22397    /// unmodified t=1 program via the col-select door). Ok(false) when the layer is not on
22398    /// the resident fused TP2 class (caller falls back to the per-row walk).
22399    pub(crate) fn step35_verify_qkv_precompute(
22400        &self,
22401        e: &Engine,
22402        il: usize,
22403        h_t: &CudaSlice<f32>,
22404        t: usize,
22405    ) -> Result<bool, Box<dyn std::error::Error>> {
22406        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22407            return Ok(false);
22408        };
22409        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22410            return Ok(false);
22411        };
22412        let Some(attention) = tp.attention.as_ref() else {
22413            return Ok(false);
22414        };
22415        if !tp.runtime.native_p2p() || !crate::tp::step_tp_qkv_fused_enabled()? {
22416            return Ok(false);
22417        }
22418        let geometry = self.step35_geom(il);
22419        let heads = geometry.n_head as usize;
22420        let ws_index = tp
22421            .runtime
22422            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22423        let gate_shards = attention
22424            .gate_shards_bf16
22425            .as_deref()
22426            .map(crate::tp::StepTpGateShards::Bf16);
22427        tp.runtime.decode_v2_input_qkv_tcol(
22428            ws_index,
22429            e,
22430            h_t,
22431            t,
22432            &tp.q,
22433            &tp.k,
22434            &tp.v,
22435            gate_shards,
22436        )?;
22437        Ok(true)
22438    }
22439
22440    /// MEMRA_TCOL_OPROJ join for the verify walk: after every column of layer `il`
22441    /// stashed its `gated` rows, produce the [t, o_out] `mixed` slab on `e` via the
22442    /// weight-amortized b4_tcol + slab join. Callers only reach this after the stash
22443    /// flag confirmed the defer engaged for every column.
22444    pub(crate) fn step35_verify_oproj_tcol(
22445        &self,
22446        e: &Engine,
22447        il: usize,
22448        t: usize,
22449    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22450        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22451            return Err("tcol o_proj join expects full attention".into());
22452        };
22453        let tp = fa
22454            .step_tp_qkv
22455            .as_ref()
22456            .ok_or("tcol o_proj join lost its resident projections")?;
22457        let heads = self.step35_geom(il).n_head as usize;
22458        let ws_index = tp
22459            .runtime
22460            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22461        tp.runtime.decode_v2_oproj_tcol(ws_index, e, &tp.o, t)
22462    }
22463
22464    /// MEMRA_SPEC_FA2 precheck: decide BEFORE arming the defer whether both verify
22465    /// columns of layer `il` will take the dcw arm AND the T=2 launch is bit-safe —
22466    /// stashing is unrecoverable (no per-column output exists), so every dynamic input
22467    /// to the engine-side dcw decision is evaluated here, plus the equal-partition
22468    /// guard fa_decode_dcw2's contract requires. Boundary rounds return false and the
22469    /// walk runs the ordinary per-column program.
22470    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
22471    pub(crate) fn step35_spec_fa2_precheck(
22472        &self,
22473        cache: &Cache,
22474        il: usize,
22475        pos0: usize,
22476    ) -> Result<bool, Box<dyn std::error::Error>> {
22477        // MEMRA_SPEC_FA2_DEBUG=1: print the first failing clause once per clause id —
22478        // a silently-vacuous door is indistinguishable from a slow one without this.
22479        fn nope(clause: &str, il: usize, pos0: usize) -> bool {
22480            static DBG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22481            static SEEN: std::sync::Mutex<Vec<&'static str>> = std::sync::Mutex::new(Vec::new());
22482            if *DBG.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1")) {
22483                let mut seen = SEEN.lock().unwrap();
22484                if !seen.contains(&clause) {
22485                    // leak: bounded by the clause-id set
22486                    seen.push(Box::leak(clause.to_string().into_boxed_str()));
22487                    eprintln!("[spec-fa2] precheck FAIL clause={clause} il={il} pos0={pos0}");
22488                }
22489            }
22490            false
22491        }
22492        // MEMRA_SPEC_FA2_LAYER=<il>: engage on ONE layer only (divergence bisection).
22493        static ONLY: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
22494        if let Some(only) =
22495            ONLY.get_or_init(|| std::env::var("MEMRA_SPEC_FA2_LAYER").ok()?.parse().ok())
22496            && *only != il
22497        {
22498            return Ok(false);
22499        }
22500        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22501            return Ok(nope("mixer", il, pos0));
22502        };
22503        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22504            return Ok(nope("step_tp", il, pos0));
22505        };
22506        let Some(attention) = tp.attention.as_ref() else {
22507            return Ok(nope("attention", il, pos0));
22508        };
22509        if !tp.runtime.native_p2p()
22510            || crate::Engine::kv_fp8_on()
22511            || !crate::tp::step_tp_dcw_enabled()?
22512            || !crate::tp::step_tp_qkv_fused_enabled()?
22513            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22514        {
22515            return Ok(nope("runtime-doors", il, pos0));
22516        }
22517        let geometry = self.step35_geom(il);
22518        let head_dim = geometry.head_dim_k as usize;
22519        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22520            return Ok(nope("fa-class", il, pos0));
22521        }
22522        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22523            return Ok(nope("tp-kv", il, pos0));
22524        };
22525        if distributed.staged_len() != pos0 {
22526            return Ok(nope("staged-len", il, pos0));
22527        }
22528        // Both appends must land without a ring rebase (rebase columns take the
22529        // host-row path, which cannot stash).
22530        let (_, would_rebase) = distributed.peek_append_ring(2)?;
22531        if would_rebase {
22532            return Ok(nope("rebase", il, pos0));
22533        }
22534        let window = geometry.window.map(|w| w as usize);
22535        // Capped SWA is REDUCTION-CLASS in the joined kernel (the two rows' windows
22536        // shift by one key, so one shared tile grid cannot reproduce both rows'
22537        // per-column FP grouping) — and drifted verify logits change accept decisions,
22538        // breaking the spec==target contract. Engage only when BOTH rows' views start
22539        // at 0 (global, or SWA still inside its window): bitwise per row under the
22540        // partition guard below. At agentic ctx this keeps the global layers — ~3/4 of
22541        // the per-key fa work — and leaves capped-SWA layers on the per-column program.
22542        if let Some(w) = window
22543            && pos0 + 2 > w
22544        {
22545            return Ok(nope("swa-capped", il, pos0));
22546        }
22547        // Row r's own per-column launch sees the POST-append view: T_r = pos0 + 1 + r
22548        // (kernel T_kv = len_dev - lstart; the host bucket matches it — the one-partition
22549        // law). Both dcw eligibility (t_kv_eff >= 96) and the vec floor key off T0.
22550        let (t0, t1) = (pos0 + 1, pos0 + 2);
22551        if t0 < 96 {
22552            return Ok(nope("dcw-floor", il, pos0));
22553        }
22554        if std::env::var("MEMRA_NO_FA_VEC").is_ok() || t0 < crate::fa_vec_min_tkv() {
22555            return Ok(nope("vec-floor", il, pos0));
22556        }
22557        // Equal-partition guard, on the KERNEL's derivation: split width (sp), effective
22558        // count (ns = ceil(T/sp)) and stride (per = ceil(T/ns)) must all match between
22559        // the two rows' own launches — the joined kernel derives one grid from T1 and
22560        // row0 inherits it, so any difference shifts row0's split boundaries and changes
22561        // the combine's merge rounding. Boundary rounds fall back per column.
22562        let ranks = tp.runtime.devices().len();
22563        let local_kv_heads = (geometry.n_head_kv as usize / ranks).max(1);
22564        let sp0 = crate::fa_split_keys_pub(t0, local_kv_heads);
22565        let sp1 = crate::fa_split_keys_pub(t1, local_kv_heads);
22566        if sp0 != sp1 {
22567            return Ok(nope("partition-sp", il, pos0));
22568        }
22569        let (ns0, ns1) = (t0.div_ceil(sp0), t1.div_ceil(sp1));
22570        if ns0 != ns1 {
22571            return Ok(nope("partition-ns", il, pos0));
22572        }
22573        if t0.div_ceil(ns0) != t1.div_ceil(ns1) {
22574            return Ok(nope("partition-per", il, pos0));
22575        }
22576        Ok(true)
22577    }
22578
22579    /// T-ROW fa precheck (the rows kernel supersedes the dcw2 pair-join): every dynamic
22580    /// input of the engine-side dcw decision must hold for EVERY row — stashing is
22581    /// unrecoverable — plus the rows-launcher guards (big-rig ladder, no env split
22582    /// overrides). No partition or capped-SWA clauses: each row derives its OWN geometry.
22583    pub(crate) fn step35_fa_rows_precheck(
22584        &self,
22585        cache: &Cache,
22586        il: usize,
22587        pos0: usize,
22588        t: usize,
22589    ) -> Result<bool, Box<dyn std::error::Error>> {
22590        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22591            return Ok(false);
22592        };
22593        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22594            return Ok(false);
22595        };
22596        let Some(attention) = tp.attention.as_ref() else {
22597            return Ok(false);
22598        };
22599        if !tp.runtime.native_p2p()
22600            || crate::Engine::kv_fp8_on()
22601            || !crate::tp::step_tp_dcw_enabled()?
22602            || !crate::tp::step_tp_qkv_fused_enabled()?
22603            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22604        {
22605            return Ok(false);
22606        }
22607        let geometry = self.step35_geom(il);
22608        let head_dim = geometry.head_dim_k as usize;
22609        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22610            return Ok(false);
22611        }
22612        if crate::fa_sm_count() < 128
22613            || std::env::var("MEMRA_FA_SPLIT").is_ok()
22614            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22615            || std::env::var("MEMRA_FA_SP16").is_ok()
22616            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22617        {
22618            return Ok(false);
22619        }
22620        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22621            return Ok(false);
22622        };
22623        if distributed.staged_len() != pos0 {
22624            return Ok(false);
22625        }
22626        let (_, would_rebase) = distributed.peek_append_ring(t)?;
22627        if would_rebase {
22628            return Ok(false);
22629        }
22630        // Row 0 sees the smallest view: its post-append effective t_kv must clear both
22631        // the dcw floor and the vec-class floor (later rows only grow).
22632        let window = geometry.window.map(|w| w as usize);
22633        let t0 = window.map(|w| (pos0 + 1).min(w)).unwrap_or(pos0 + 1);
22634        if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22635            return Ok(false);
22636        }
22637        Ok(true)
22638    }
22639
22640    /// T-ROW fa join for the verify walk (same-session rows: shared ring/len with
22641    /// len_back = t-1-r). Tables stage once per (layer, rank, ring, t) and live on the
22642    /// owning rank.
22643    pub(crate) fn step35_verify_fa_rows_join(
22644        &self,
22645        e: &Engine,
22646        il: usize,
22647        cache: &Cache,
22648        pos0: usize,
22649        t: usize,
22650    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
22651        use cudarc::driver::DevicePtr;
22652        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22653            return Err("fa rows join expects full attention".into());
22654        };
22655        let tp = fa
22656            .step_tp_qkv
22657            .as_ref()
22658            .ok_or("fa rows join lost its resident projections")?;
22659        let geometry = self.step35_geom(il);
22660        let heads = geometry.n_head as usize;
22661        let head_dim = geometry.head_dim_k as usize;
22662        let window = geometry.window.map(|w| w as usize);
22663        let distributed = cache.tp_kv[il]
22664            .as_ref()
22665            .ok_or("fa rows join lost its distributed KV cache")?;
22666        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
22667        // Host mirror of the kernel's big-rig ladder (launcher-guarded identical).
22668        let ladder = |t_kv: usize| -> usize {
22669            if t_kv <= 2048 {
22670                16
22671            } else if t_kv <= 16384 {
22672                64
22673            } else {
22674                128
22675            }
22676        };
22677        let mut max_ns = 1usize;
22678        for r in 0..t {
22679            let t_kv = window
22680                .map(|w| (pos0 + r + 1).min(w))
22681                .unwrap_or(pos0 + r + 1);
22682            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
22683        }
22684        // Rebuild the tiny raw-pointer table from the live distributed cache immediately
22685        // before launch. A process-lifetime map cannot prove allocation generation: CUDA may
22686        // recycle len/base independently of the large K/V rings, making a pointer-key cache
22687        // hit refer to another session (Hermes `11339f5cd3c132a3`).
22688        let ranks = tp.runtime.devices().len();
22689        let mut tables = Vec::with_capacity(ranks);
22690        for rank in 0..ranks {
22691            let engine = tp
22692                .runtime
22693                .rank_engine(rank)
22694                .ok_or("fa rows join lost a rank engine")?;
22695            let rank_cache = distributed
22696                .rank(rank)
22697                .ok_or("fa rows join lost a KV cache rank")?;
22698            let _main = engine.gpu.enter_main()?;
22699            let s = engine.stream();
22700            let (kp, _g0) = rank_cache.k().device_ptr(&s);
22701            let (vp, _g1) = rank_cache.v().device_ptr(&s);
22702            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
22703            let bp = match rank_cache.base_d() {
22704                Some(b) => {
22705                    let (p, _g) = b.device_ptr(&s);
22706                    p
22707                }
22708                None => 0u64,
22709            };
22710            let mut host = Vec::with_capacity(t * 6);
22711            for r in 0..t {
22712                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, (t - 1 - r) as u64]);
22713            }
22714            tables.push(engine.stream().clone_htod(&host)?);
22715        }
22716        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
22717        let ws_index = tp
22718            .runtime
22719            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22720        tp.runtime.decode_v2_fa_rows_join(
22721            ws_index,
22722            e,
22723            &tp.o,
22724            &tabs,
22725            t,
22726            head_dim,
22727            window.unwrap_or(0),
22728            max_ns,
22729            geometry.attention_scale(),
22730            k_tok_bytes,
22731            v_tok_bytes,
22732        )
22733    }
22734
22735    /// Multi-session t-row fa precheck (the batched serving walk): the static doors of
22736    /// the rows kernel plus per-SESSION dynamic checks — every row's own cache must be
22737    /// hydrated, in sync, rebase-free and above both floors.
22738    pub(crate) fn step35_batch_fa_rows_precheck(
22739        &self,
22740        caches: &[&mut Cache],
22741        row_to_cache: impl Fn(usize) -> usize,
22742        positions: &[i32],
22743        il: usize,
22744    ) -> Result<bool, Box<dyn std::error::Error>> {
22745        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22746            return Ok(false);
22747        };
22748        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22749            return Ok(false);
22750        };
22751        let Some(attention) = tp.attention.as_ref() else {
22752            return Ok(false);
22753        };
22754        if !tp.runtime.native_p2p()
22755            || crate::Engine::kv_fp8_on()
22756            || !crate::tp::step_tp_dcw_enabled()?
22757            || !crate::tp::step_tp_qkv_fused_enabled()?
22758            || (attention.gate_shards.is_none() && attention.gate_shards_bf16.is_none())
22759        {
22760            return Ok(false);
22761        }
22762        let geometry = self.step35_geom(il);
22763        let head_dim = geometry.head_dim_k as usize;
22764        if head_dim > 256 || !head_dim.is_multiple_of(32) || !crate::fa_v3_on() {
22765            return Ok(false);
22766        }
22767        if crate::fa_sm_count() < 128
22768            || std::env::var("MEMRA_FA_SPLIT").is_ok()
22769            || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
22770            || std::env::var("MEMRA_FA_SP16").is_ok()
22771            || std::env::var("MEMRA_NO_FA_VEC").is_ok()
22772        {
22773            return Ok(false);
22774        }
22775        let window = geometry.window.map(|w| w as usize);
22776        for (r, &pos) in positions.iter().enumerate() {
22777            let cache = &caches[row_to_cache(r)];
22778            let Some(distributed) = cache.tp_kv[il].as_ref() else {
22779                return Ok(false);
22780            };
22781            if distributed.staged_len() != pos as usize {
22782                return Ok(false);
22783            }
22784            if distributed.peek_append_ring(1)?.1 {
22785                return Ok(false);
22786            }
22787            let t0 = window
22788                .map(|w| (pos as usize + 1).min(w))
22789                .unwrap_or(pos as usize + 1);
22790            if t0 < 96 || t0 < crate::fa_vec_min_tkv() {
22791                return Ok(false);
22792            }
22793        }
22794        Ok(true)
22795    }
22796
22797    /// FULL t-row attention pass for the VERIFY walk (same-session rows): rope/append +
22798    /// fa + combine + o_proj join in 3 launches/rank/layer. Row r appends at slot
22799    /// len-base+r and one last block advances len by t; the fa rows read len_back =
22800    /// t-1-r. Returns None when the fused-rope class does not hold (the walk keeps the
22801    /// per-column stash flow). Caller has passed `step35_fa_rows_precheck`.
22802    pub(crate) fn step35_verify_rope_fa_pass(
22803        &self,
22804        e: &Engine,
22805        il: usize,
22806        cache: &Cache,
22807        pos0: usize,
22808        t: usize,
22809        stage_pos: bool,
22810    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22811        use cudarc::driver::DevicePtr;
22812        if !crate::tp::fuse_rope_append_on() {
22813            return Ok(None);
22814        }
22815        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22816            return Ok(None);
22817        };
22818        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22819            return Ok(None);
22820        };
22821        let Some(attention) = tp.attention.as_ref() else {
22822            return Ok(None);
22823        };
22824        let geometry = self.step35_geom(il);
22825        let head_dim = geometry.head_dim_k as usize;
22826        if head_dim != 128 {
22827            return Ok(None);
22828        }
22829        let heads = geometry.n_head as usize;
22830        let window = geometry.window.map(|w| w as usize);
22831        let ranks = tp.runtime.devices().len();
22832        let Some(distributed) = cache.tp_kv[il].as_ref() else {
22833            return Ok(None);
22834        };
22835        if distributed.kv_dim_k() != distributed.kv_dim_v() {
22836            return Ok(None);
22837        }
22838        {
22839            let rank0 = distributed.rank(0).ok_or("verify rope pass lost rank 0")?;
22840            if rank0.base_d().is_none()
22841                && distributed.staged_len() + t > distributed.physical_capacity()
22842            {
22843                return Ok(None);
22844            }
22845        }
22846        let mut rope_freqs = Vec::with_capacity(ranks);
22847        for rank in 0..ranks {
22848            let engine = tp
22849                .runtime
22850                .rank_engine(rank)
22851                .ok_or("verify rope pass lost a rank engine")?;
22852            rope_freqs.push(if geometry.rope_factors {
22853                match self
22854                    .step35_aux
22855                    .as_ref()
22856                    .and_then(|aux| aux.rope_freqs(engine))
22857                {
22858                    Some(f) => Some(f),
22859                    None => return Ok(None),
22860                }
22861            } else {
22862                None
22863            });
22864        }
22865        let ladder = |t_kv: usize| -> usize {
22866            if t_kv <= 2048 {
22867                16
22868            } else if t_kv <= 16384 {
22869                64
22870            } else {
22871                128
22872            }
22873        };
22874        let (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
22875        let mut max_ns = 1usize;
22876        let mut positions = Vec::with_capacity(t);
22877        for r in 0..t {
22878            positions.push((pos0 + r) as i32);
22879            let t_kv = window
22880                .map(|w| (pos0 + r + 1).min(w))
22881                .unwrap_or(pos0 + r + 1);
22882            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
22883        }
22884        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
22885        let mut tab_keys = vec![0u64; ranks];
22886        for rank in 0..ranks {
22887            let engine = tp
22888                .runtime
22889                .rank_engine(rank)
22890                .ok_or("verify rope pass lost a rank engine")?;
22891            let rank_cache = distributed
22892                .rank(rank)
22893                .ok_or("verify rope pass lost a KV cache rank")?;
22894            let _main = engine.gpu.enter_main()?;
22895            let s = engine.stream();
22896            let (kp, _g0) = rank_cache.k().device_ptr(&s);
22897            let (vp, _g1) = rank_cache.v().device_ptr(&s);
22898            let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
22899            let bp = match rank_cache.base_d() {
22900                Some(b) => {
22901                    let (p, _g) = b.device_ptr(&s);
22902                    p
22903                }
22904                None => 0u64,
22905            };
22906            tab_keys[rank] = kp
22907                .rotate_left(17)
22908                .wrapping_add(bp)
22909                .wrapping_add((il as u64) << 32)
22910                .wrapping_add(t as u64)
22911                .wrapping_add(1 << 63);
22912            for _r in 0..t {
22913                session_parts[rank].push([kp, vp, lp, bp]);
22914            }
22915        }
22916        let ws_index = tp
22917            .runtime
22918            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
22919        tp.runtime
22920            .decode_v2_rope_fa_rows(
22921                ws_index,
22922                e,
22923                &tp.o,
22924                &session_parts,
22925                &tab_keys,
22926                &positions,
22927                stage_pos,
22928                true,
22929                &attention.q_norm,
22930                &attention.k_norm,
22931                &rope_freqs,
22932                t,
22933                head_dim,
22934                geometry.n_rot as usize,
22935                window.unwrap_or(0),
22936                max_ns,
22937                geometry.attention_scale(),
22938                k_tok_bytes,
22939                v_tok_bytes,
22940                self.cfg.rms_eps,
22941                geometry.rope_base,
22942            )
22943            .map(Some)
22944    }
22945
22946    /// FULL t-row attention pass for the batched walk (rope/append + fa + combine +
22947    /// o_proj join, 3 launches/rank/layer): returns None when the fused-rope class does
22948    /// not hold — the caller falls back to the per-row stash flow. The caller has
22949    /// already passed `step35_batch_fa_rows_precheck`.
22950    #[allow(clippy::too_many_arguments)]
22951    pub(crate) fn step35_batch_rope_fa_pass(
22952        &self,
22953        e: &Engine,
22954        il: usize,
22955        caches: &[&mut Cache],
22956        row_to_cache: impl Fn(usize) -> usize,
22957        positions: &[i32],
22958        t: usize,
22959        stage_pos: bool,
22960    ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
22961        use cudarc::driver::DevicePtr;
22962        if !crate::tp::fuse_rope_append_on() {
22963            return Ok(None);
22964        }
22965        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
22966            return Ok(None);
22967        };
22968        let Some(tp) = fa.step_tp_qkv.as_ref() else {
22969            return Ok(None);
22970        };
22971        let Some(attention) = tp.attention.as_ref() else {
22972            return Ok(None);
22973        };
22974        let geometry = self.step35_geom(il);
22975        let head_dim = geometry.head_dim_k as usize;
22976        if head_dim != 128 {
22977            return Ok(None);
22978        }
22979        let heads = geometry.n_head as usize;
22980        let window = geometry.window.map(|w| w as usize);
22981        let ranks = tp.runtime.devices().len();
22982        // The rows kernels never arm base_d; refuse once a ring could have rebased
22983        // without an armed base (the table would read base=0 after a real rebase).
22984        for r in 0..t {
22985            let cache = &caches[row_to_cache(r)];
22986            let Some(distributed) = cache.tp_kv[il].as_ref() else {
22987                return Ok(None);
22988            };
22989            if distributed.kv_dim_k() != distributed.kv_dim_v() {
22990                return Ok(None);
22991            }
22992            let rank0 = distributed.rank(0).ok_or("rope fa pass lost rank 0")?;
22993            if rank0.base_d().is_none()
22994                && distributed.staged_len() + t > distributed.physical_capacity()
22995            {
22996                return Ok(None);
22997            }
22998        }
22999        let mut rope_freqs = Vec::with_capacity(ranks);
23000        for rank in 0..ranks {
23001            let engine = tp
23002                .runtime
23003                .rank_engine(rank)
23004                .ok_or("rope fa pass lost a rank engine")?;
23005            rope_freqs.push(if geometry.rope_factors {
23006                match self
23007                    .step35_aux
23008                    .as_ref()
23009                    .and_then(|aux| aux.rope_freqs(engine))
23010                {
23011                    Some(f) => Some(f),
23012                    None => return Ok(None),
23013                }
23014            } else {
23015                None
23016            });
23017        }
23018        let ladder = |t_kv: usize| -> usize {
23019            if t_kv <= 2048 {
23020                16
23021            } else if t_kv <= 16384 {
23022                64
23023            } else {
23024                128
23025            }
23026        };
23027        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23028        let mut session_parts: Vec<Vec<[u64; 4]>> = vec![Vec::with_capacity(t); ranks];
23029        let mut tab_keys = vec![0u64; ranks];
23030        for (r, &pos) in positions.iter().enumerate().take(t) {
23031            let cache = &caches[row_to_cache(r)];
23032            let distributed = cache.tp_kv[il]
23033                .as_ref()
23034                .ok_or("rope fa pass lost a distributed KV cache")?;
23035            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23036            let t_kv = window
23037                .map(|w| (pos as usize + 1).min(w))
23038                .unwrap_or(pos as usize + 1);
23039            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23040            for rank in 0..ranks {
23041                let engine = tp
23042                    .runtime
23043                    .rank_engine(rank)
23044                    .ok_or("rope fa pass lost a rank engine")?;
23045                let rank_cache = distributed
23046                    .rank(rank)
23047                    .ok_or("rope fa pass lost a KV cache rank")?;
23048                let _main = engine.gpu.enter_main()?;
23049                let s = engine.stream();
23050                let (kp, _g0) = rank_cache.k().device_ptr(&s);
23051                let (vp, _g1) = rank_cache.v().device_ptr(&s);
23052                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23053                let bp = match rank_cache.base_d() {
23054                    Some(b) => {
23055                        let (p, _g) = b.device_ptr(&s);
23056                        p
23057                    }
23058                    None => 0u64,
23059                };
23060                tab_keys[rank] = tab_keys[rank]
23061                    .rotate_left(9)
23062                    .wrapping_add(kp)
23063                    .wrapping_add(bp)
23064                    .wrapping_add(il as u64);
23065                session_parts[rank].push([kp, vp, lp, bp]);
23066            }
23067        }
23068        let ws_index = tp
23069            .runtime
23070            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23071        tp.runtime
23072            .decode_v2_rope_fa_rows(
23073                ws_index,
23074                e,
23075                &tp.o,
23076                &session_parts,
23077                &tab_keys,
23078                positions,
23079                stage_pos,
23080                false,
23081                &attention.q_norm,
23082                &attention.k_norm,
23083                &rope_freqs,
23084                t,
23085                head_dim,
23086                geometry.n_rot as usize,
23087                window.unwrap_or(0),
23088                max_ns,
23089                geometry.attention_scale(),
23090                k_tok_bytes,
23091                v_tok_bytes,
23092                self.cfg.rms_eps,
23093                geometry.rope_base,
23094            )
23095            .map(Some)
23096    }
23097
23098    /// Multi-session t-row fa join (batched serving): per-row table entries point at
23099    /// each row's OWN session rings/counters (len_back = 0 — every session appended
23100    /// exactly its one row). Tables stage once per (layer, rank, session-set, t).
23101    #[allow(clippy::too_many_arguments)]
23102    pub(crate) fn step35_batch_fa_rows_join(
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    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23111        use cudarc::driver::DevicePtr;
23112        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23113            return Err("batch fa rows join expects full attention".into());
23114        };
23115        let tp = fa
23116            .step_tp_qkv
23117            .as_ref()
23118            .ok_or("batch fa rows join lost its resident projections")?;
23119        let geometry = self.step35_geom(il);
23120        let heads = geometry.n_head as usize;
23121        let head_dim = geometry.head_dim_k as usize;
23122        let window = geometry.window.map(|w| w as usize);
23123        let ladder = |t_kv: usize| -> usize {
23124            if t_kv <= 2048 {
23125                16
23126            } else if t_kv <= 16384 {
23127                64
23128            } else {
23129                128
23130            }
23131        };
23132        let (mut max_ns, mut k_tok_bytes, mut v_tok_bytes) = (1usize, 0usize, 0usize);
23133        for (r, &pos) in positions.iter().enumerate() {
23134            let cache = &caches[row_to_cache(r)];
23135            let distributed = cache.tp_kv[il]
23136                .as_ref()
23137                .ok_or("batch fa rows join lost a distributed KV cache")?;
23138            (k_tok_bytes, v_tok_bytes) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
23139            let t_kv = window
23140                .map(|w| (pos as usize + 1).min(w))
23141                .unwrap_or(pos as usize + 1);
23142            max_ns = max_ns.max(t_kv.div_ceil(ladder(t_kv)));
23143        }
23144        // Multi-session tables also rebuild from every live K/V/len/base tuple. Keeping a
23145        // process-lifetime raw-pointer cache here omitted V and len identity and had no
23146        // allocation generation, so allocator reuse could bind one request to another.
23147        let ranks = tp.runtime.devices().len();
23148        let mut tables = Vec::with_capacity(ranks);
23149        for rank in 0..ranks {
23150            let engine = tp
23151                .runtime
23152                .rank_engine(rank)
23153                .ok_or("batch fa rows join lost a rank engine")?;
23154            let _main = engine.gpu.enter_main()?;
23155            let s = engine.stream();
23156            let mut host = Vec::with_capacity(t * 6);
23157            for r in 0..t {
23158                let cache = &caches[row_to_cache(r)];
23159                let distributed = cache.tp_kv[il]
23160                    .as_ref()
23161                    .ok_or("batch fa rows join lost a distributed KV cache")?;
23162                let rank_cache = distributed
23163                    .rank(rank)
23164                    .ok_or("batch fa rows join lost a KV cache rank")?;
23165                let (kp, _g0) = rank_cache.k().device_ptr(&s);
23166                let (vp, _g1) = rank_cache.v().device_ptr(&s);
23167                let (lp, _g2) = rank_cache.len_d().device_ptr(&s);
23168                let bp = match rank_cache.base_d() {
23169                    Some(b) => {
23170                        let (p, _g) = b.device_ptr(&s);
23171                        p
23172                    }
23173                    None => 0u64,
23174                };
23175                host.extend_from_slice(&[kp, vp, lp, bp, 0u64, 0u64]);
23176            }
23177            tables.push(engine.stream().clone_htod(&host)?);
23178        }
23179        let tabs: Vec<&CudaSlice<u64>> = tables.iter().collect();
23180        let ws_index = tp
23181            .runtime
23182            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23183        tp.runtime.decode_v2_fa_rows_join(
23184            ws_index,
23185            e,
23186            &tp.o,
23187            &tabs,
23188            t,
23189            head_dim,
23190            window.unwrap_or(0),
23191            max_ns,
23192            geometry.attention_scale(),
23193            k_tok_bytes,
23194            v_tok_bytes,
23195        )
23196    }
23197
23198    /// MEMRA_SPEC_FA2 join for the verify walk: both columns stashed; one shared-KV T=2
23199    /// fa per rank + the weight-amortized o_proj join produce the [2, o_out] `mixed`
23200    /// slab on `e`.
23201    #[allow(dead_code)] // allow: banked MEMRA_SPEC_FA2 arm; kept as the named seam for the spec-FA2 join program
23202    pub(crate) fn step35_verify_spec_fa2_join(
23203        &self,
23204        e: &Engine,
23205        il: usize,
23206        cache: &Cache,
23207        pos0: usize,
23208    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23209        let crate::hybrid::Mixer::Full(fa) = &self.layers[il].mixer else {
23210            return Err("spec fa2 join expects full attention".into());
23211        };
23212        let tp = fa
23213            .step_tp_qkv
23214            .as_ref()
23215            .ok_or("spec fa2 join lost its resident projections")?;
23216        let geometry = self.step35_geom(il);
23217        let heads = geometry.n_head as usize;
23218        let head_dim = geometry.head_dim_k as usize;
23219        let window = geometry.window.map(|w| w as usize);
23220        // POST-append view of the second row (kernel T1 = len - lstart with len =
23221        // pos0 + 2): sp/ns derive from it, and the precheck proved row0 shares them.
23222        let bucket = window.map(|w| (pos0 + 2).min(w)).unwrap_or(pos0 + 2);
23223        let distributed = cache.tp_kv[il]
23224            .as_ref()
23225            .ok_or("spec fa2 join lost its distributed KV cache")?;
23226        let ws_index = tp
23227            .runtime
23228            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23229        tp.runtime.decode_v2_spec_fa2_join(
23230            ws_index,
23231            e,
23232            &tp.o,
23233            distributed,
23234            head_dim,
23235            window.unwrap_or(0),
23236            bucket,
23237            geometry.attention_scale(),
23238        )
23239    }
23240
23241    pub(crate) fn step35_tp_decode_attn_resident_v2(
23242        &self,
23243        e: &Engine,
23244        fa: &FullAttnLayer,
23245        il: usize,
23246        h: &CudaSlice<f32>,
23247        pos_d: &CudaSlice<i32>,
23248        cache: &mut Cache,
23249    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23250        let tp = fa
23251            .step_tp_qkv
23252            .as_ref()
23253            .ok_or("Step TP decode lost its resident projections")?;
23254        let attention = tp
23255            .attention
23256            .as_ref()
23257            .ok_or("Step TP decode lost its resident attention auxiliaries")?;
23258        if !tp.runtime.native_p2p() {
23259            return Err("rank-local Step attention requires native P2P".into());
23260        }
23261        if crate::Engine::kv_fp8_on() {
23262            return Err("rank-local Step attention has not qualified the FP8 KV cache".into());
23263        }
23264
23265        let geometry = self.cfg.full_attention_geometry_at(il as u32);
23266        let window = geometry.window.map(|window| window as usize);
23267        let ranks = tp.runtime.devices().len();
23268        let head_dim = geometry.head_dim_k as usize;
23269        let heads = geometry.n_head as usize;
23270        let kv_heads = geometry.n_head_kv as usize;
23271        if !heads.is_multiple_of(ranks) || !kv_heads.is_multiple_of(ranks) {
23272            return Err(format!(
23273                "Step attention heads q={heads} kv={kv_heads} are not divisible by TP={ranks}"
23274            )
23275            .into());
23276        }
23277        let local_heads = heads / ranks;
23278        let local_kv_heads = kv_heads / ranks;
23279        let max_ctx = cache.max_ctx;
23280
23281        let hydrated = self.ensure_step_tp_kv_cache(e, fa, il, cache)?;
23282
23283        let base_len = cache.kv[il]
23284            .as_ref()
23285            .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?
23286            .len;
23287        {
23288            let distributed = cache.tp_kv[il]
23289                .as_ref()
23290                .ok_or_else(|| format!("Step TP layer {il} lost its distributed KV cache"))?;
23291            if distributed.committed_len() != base_len || distributed.staged_len() != base_len {
23292                return Err(format!(
23293                    "Step TP layer {il} cache lengths diverged before decode: \
23294                     local={base_len} distributed={}/{}",
23295                    distributed.committed_len(),
23296                    distributed.staged_len()
23297                )
23298                .into());
23299            }
23300        }
23301        if pos_d.len() != 1 {
23302            return Err(format!(
23303                "rank-local Step decode requires one position, got {}",
23304                pos_d.len()
23305            )
23306            .into());
23307        }
23308
23309        let decode_input = attention
23310            .decode_input
23311            .as_ref()
23312            .ok_or("Step TP decode v2 requires the replicated decode input")?;
23313        let mut decode_input = decode_input
23314            .lock()
23315            .map_err(|_| "Step TP replicated decode input lock is poisoned")?;
23316
23317        let has_gate = fa.attn_gate.is_some();
23318        // Gate: with per-rank shards loaded (fused door), the fused QKV+gate kernel computes
23319        // it rank-locally and the model-engine matmul (and its staging copies) disappears.
23320        // Otherwise it queues on e's stream BEFORE decode_v2_input_qkv records the entry
23321        // event, so the rank-stream reads of the staged gate are ordered without a host sync.
23322        let use_gate_shards = has_gate
23323            && (attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some())
23324            && crate::tp::step_tp_qkv_fused_enabled()?;
23325        let gate_raw = if !has_gate || use_gate_shards {
23326            None
23327        } else {
23328            let gate_weight = fa
23329                .attn_gate
23330                .as_ref()
23331                .ok_or("step35 layer is missing attn_gate.weight")?;
23332            let gate_raw = e.matmul(gate_weight, h, 1)?;
23333            if gate_raw.len() != heads {
23334                return Err(format!(
23335                    "Step TP layer {il} gate output {} != {heads}",
23336                    gate_raw.len()
23337                )
23338                .into());
23339            }
23340            Some(gate_raw)
23341        };
23342
23343        let ws_index = tp
23344            .runtime
23345            .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
23346        let mut ws_guard = tp
23347            .runtime
23348            .decode_v2_workspace()
23349            .lock()
23350            .map_err(|_| "Step TP decode v2 workspace lock is poisoned")?;
23351        let ws = ws_guard
23352            .get_mut(ws_index)
23353            .ok_or("Step TP decode v2 workspace missing after ensure")?;
23354
23355        let mut rope_freqs = Vec::with_capacity(ranks);
23356        for rank in 0..ranks {
23357            let engine = tp
23358                .runtime
23359                .rank_engine(rank)
23360                .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23361            rope_freqs.push(if geometry.rope_factors {
23362                self.step35_aux
23363                    .as_ref()
23364                    .and_then(|aux| aux.rope_freqs(engine))
23365            } else {
23366                None
23367            });
23368        }
23369        // DCW arm (MEMRA_STEP_TP_DCW=1): device-counter append + counter-derived fa — the
23370        // exact captured-child content, run eagerly. bucket = effective t_kv keeps the fa
23371        // split geometry identical to the kvmod path (one-partition law) -> bit-identical.
23372        // Rebase tokens and sub-vec-floor contexts take the host-row path below.
23373        // (Eligibility computed BEFORE input_qkv so FUSION #1 can defer the norm+rope into
23374        // the fused rope+append+inc launch on dcw tokens.)
23375        let staged_next = base_len + 1;
23376        let t_kv_eff = window
23377            .map(|window| staged_next.min(window))
23378            .unwrap_or(staged_next);
23379        let dcw = crate::tp::step_tp_dcw_enabled()?
23380            && (use_gate_shards || (!has_gate && crate::tp::step_tp_qkv_fused_enabled()?))
23381            && t_kv_eff >= 96
23382            && {
23383                let (write_row, would_rebase) = cache.tp_kv[il]
23384                    .as_ref()
23385                    .expect("distributed cache checked above")
23386                    .peek_append_ring(1)?;
23387                if !would_rebase {
23388                    // Arm the base mirrors on first use: base = logical staged - physical row.
23389                    let base = (base_len - write_row) as i32;
23390                    let distributed = cache.tp_kv[il]
23391                        .as_mut()
23392                        .expect("distributed cache checked above");
23393                    for rank in 0..ranks {
23394                        let engine = tp.runtime.rank_engine(rank).ok_or_else(|| {
23395                            format!("Step TP layer {il} has no engine for rank {rank}")
23396                        })?;
23397                        let _main = engine.gpu.enter_main()?;
23398                        let rank_cache = distributed.rank_mut(rank).ok_or_else(|| {
23399                            format!("Step TP layer {il} has no KV cache rank {rank}")
23400                        })?;
23401                        if rank_cache.base_d().is_none() {
23402                            rank_cache.arm_base_d(engine.htod_i32(&[base])?);
23403                        }
23404                    }
23405                }
23406                !would_rebase
23407            };
23408        let fuse_rope = dcw
23409            && crate::tp::fuse_rope_append_on()
23410            && head_dim == 128
23411            && cache.tp_kv[il]
23412                .as_ref()
23413                .map(|d| d.kv_dim_k() == d.kv_dim_v() && d.kv_dim_k() == local_kv_heads * head_dim)
23414                .unwrap_or(false);
23415
23416        let tcol_col = crate::tp::take_verify_tcol();
23417        // MEMRA_SPEC_FA2 defer: the verify walk armed this column for the shared-KV T=2
23418        // attention. On dcw tokens the per-rank pass still norms/ropes/APPENDS (cache
23419        // state must advance per column) but skips the fa+gate launch; post-rope q and
23420        // gate rows are stashed instead, and ONE fa_decode_dcw2 per rank joins both
23421        // columns after the second append. Non-dcw tokens ignore the defer (the fa runs
23422        // normally and the walk consumes the real output — stash flag stays unset).
23423        let fa2_col = crate::tp::take_spec_fa2_defer();
23424        tp.runtime.decode_v2_input_qkv(
23425            ws,
23426            e,
23427            h,
23428            pos_d,
23429            gate_raw.as_ref(),
23430            if !use_gate_shards {
23431                None
23432            } else if let Some(shards) = attention.gate_shards.as_deref() {
23433                Some(crate::tp::StepTpGateShards::F32(shards))
23434            } else {
23435                attention
23436                    .gate_shards_bf16
23437                    .as_deref()
23438                    .map(crate::tp::StepTpGateShards::Bf16)
23439            },
23440            &mut decode_input,
23441            &tp.q,
23442            &tp.k,
23443            &tp.v,
23444            &attention.q_norm,
23445            &attention.k_norm,
23446            head_dim,
23447            geometry.n_rot as usize,
23448            geometry.rope_base,
23449            &rope_freqs,
23450            self.cfg.rms_eps,
23451            has_gate,
23452            fuse_rope,
23453            tcol_col,
23454        )?;
23455
23456        let transaction = cache.tp_kv[il]
23457            .as_mut()
23458            .expect("distributed cache checked above")
23459            .begin_transaction()?;
23460        let append_result = tp.runtime.append_tp_kv_transaction_inner(
23461            cache.tp_kv[il]
23462                .as_mut()
23463                .expect("distributed cache checked above"),
23464            transaction,
23465            &ws.k,
23466            &ws.v_raw,
23467            1,
23468            dcw,
23469        );
23470        if let Err(error) = append_result {
23471            let _ = tp.runtime.rollback_tp_kv_transaction(
23472                cache.tp_kv[il]
23473                    .as_mut()
23474                    .expect("distributed cache checked above"),
23475                transaction,
23476            );
23477            return Err(error);
23478        }
23479
23480        let staged = (|| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23481            // Copy the view scalars out and DROP the shared borrow — the merged dcw arm
23482            // reborrows the cache mutably per rank.
23483            let (staged_len, physical, k_tok_bytes_c, v_tok_bytes_c, capacity) = {
23484                let distributed = cache.tp_kv[il]
23485                    .as_ref()
23486                    .expect("distributed cache checked above");
23487                let staged_len = distributed.staged_len();
23488                let view_start = window
23489                    .map(|window| staged_len.saturating_sub(window))
23490                    .unwrap_or(0);
23491                (
23492                    staged_len,
23493                    distributed.physical_range(view_start, staged_len)?,
23494                    distributed.k_tok_bytes(),
23495                    distributed.v_tok_bytes(),
23496                    distributed.physical_capacity(),
23497                )
23498            };
23499            let view_start = window
23500                .map(|window| staged_len.saturating_sub(window))
23501                .unwrap_or(0);
23502            let t_kv = staged_len - view_start;
23503            for rank in 0..ranks {
23504                let engine = tp
23505                    .runtime
23506                    .rank_engine(rank)
23507                    .ok_or_else(|| format!("Step TP layer {il} has no engine for rank {rank}"))?;
23508                let _main = engine.gpu.enter_main()?;
23509                if dcw {
23510                    // MERGED per-rank pass (the capture unit): append + inc + fa + gate on ONE
23511                    // stream visit. distributed is borrowed shared here; the planes need mut —
23512                    // reborrow through the cache Option (the closure holds cache mutably).
23513                    {
23514                        let distributed_mut = cache.tp_kv[il]
23515                            .as_mut()
23516                            .expect("distributed cache checked above");
23517                        let (kv_dim_k, kv_dim_v) =
23518                            (distributed_mut.kv_dim_k(), distributed_mut.kv_dim_v());
23519                        let (k_tok_bytes, v_tok_bytes) =
23520                            (distributed_mut.k_tok_bytes(), distributed_mut.v_tok_bytes());
23521                        let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23522                            format!("Step TP layer {il} has no KV cache rank {rank}")
23523                        })?;
23524                        let (k_plane, v_plane, len_d, base_d) =
23525                            rank_cache.planes_and_counters_mut();
23526                        if fuse_rope {
23527                            // FUSION #1: norms + rope (deferred out of input_qkv) + append
23528                            // + last-block len inc in ONE launch. Bit-identical bodies.
23529                            let same_dev = engine.ctx().ordinal() == e.ctx().ordinal();
23530                            let crate::tp::StepTpDecodeV2Ws {
23531                                q_raw,
23532                                k_raw,
23533                                v_raw,
23534                                q,
23535                                k,
23536                                pos,
23537                                pos_stage,
23538                                fuse_ctr,
23539                                ..
23540                            } = &mut *ws;
23541                            // Same-device rank: the staged-copy elision leaves pos[rank]
23542                            // stale — read the e-context pos stage directly (mirrors the
23543                            // rope elision in input_qkv_rank).
23544                            let pos_ref: &CudaSlice<i32> = if same_dev {
23545                                pos_stage
23546                                    .as_ref()
23547                                    .ok_or("step TP decode v2 pos stage not armed")?
23548                            } else {
23549                                &pos[rank]
23550                            };
23551                            engine.qk_norm_rope_append_inc_dcw(
23552                                &q_raw[rank],
23553                                &k_raw[rank],
23554                                &v_raw[rank],
23555                                &attention.q_norm[rank],
23556                                &attention.k_norm[rank],
23557                                &mut q[rank],
23558                                &mut k[rank],
23559                                pos_ref,
23560                                k_plane,
23561                                v_plane,
23562                                len_d,
23563                                base_d,
23564                                &mut fuse_ctr[rank],
23565                                kv_dim_k,
23566                                kv_dim_v,
23567                                k_tok_bytes,
23568                                v_tok_bytes,
23569                                head_dim,
23570                                geometry.n_rot as usize,
23571                                local_heads,
23572                                local_kv_heads,
23573                                self.cfg.rms_eps,
23574                                geometry.rope_base,
23575                                1.0,
23576                                rope_freqs[rank],
23577                            )?;
23578                        } else {
23579                            engine.append_kv_quantized_dcw(
23580                                &ws.k[rank],
23581                                &ws.v_raw[rank],
23582                                k_plane,
23583                                v_plane,
23584                                len_d,
23585                                base_d,
23586                                kv_dim_k,
23587                                kv_dim_v,
23588                                k_tok_bytes,
23589                                v_tok_bytes,
23590                            )?;
23591                        }
23592                        if !fuse_rope {
23593                            let rank_cache = distributed_mut.rank_mut(rank).ok_or_else(|| {
23594                                format!("Step TP layer {il} has no KV cache rank {rank}")
23595                            })?;
23596                            engine.inc_i32(rank_cache.len_d_mut())?;
23597                        }
23598                    }
23599                    let distributed = cache.tp_kv[il]
23600                        .as_ref()
23601                        .expect("distributed cache checked above");
23602                    let rank_cache = distributed
23603                        .rank(rank)
23604                        .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23605                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * k_tok_bytes_c);
23606                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * v_tok_bytes_c);
23607                    if fa2_col.is_some() {
23608                        // SPEC_FA2 defer: append landed above; the fa for this column
23609                        // runs in the T=2 joined launch after the pair's second append.
23610                        continue;
23611                    }
23612                    {
23613                        // FUSION #2d: combine + head gate in the dcw tail — `gated` receives
23614                        // the gated output directly (bit-identical; one launch saved).
23615                        let crate::tp::StepTpDecodeV2Ws { q, gate, gated, .. } = &mut *ws;
23616                        engine.fa_decode_dcw(
23617                            &q[rank],
23618                            &k_ring,
23619                            &v_ring,
23620                            &mut gated[rank],
23621                            head_dim,
23622                            local_heads,
23623                            local_kv_heads,
23624                            rank_cache.len_d(),
23625                            rank_cache.base_d(),
23626                            window.unwrap_or(0),
23627                            t_kv,
23628                            geometry.attention_scale(),
23629                            k_tok_bytes_c,
23630                            v_tok_bytes_c,
23631                            has_gate.then_some(&gate[rank]),
23632                        )?;
23633                    }
23634                    continue;
23635                }
23636                let distributed = cache.tp_kv[il]
23637                    .as_ref()
23638                    .expect("distributed cache checked above");
23639                let rank_cache = distributed
23640                    .rank(rank)
23641                    .ok_or_else(|| format!("Step TP layer {il} has no KV cache rank {rank}"))?;
23642                let k_view = engine.view_u8_range(
23643                    rank_cache.k(),
23644                    physical.start * k_tok_bytes_c,
23645                    physical.end * k_tok_bytes_c,
23646                );
23647                let v_view = engine.view_u8_range(
23648                    rank_cache.v(),
23649                    physical.start * v_tok_bytes_c,
23650                    physical.end * v_tok_bytes_c,
23651                );
23652                if has_gate {
23653                    engine.fa_decode_kvmod(
23654                        &ws.q[rank],
23655                        &k_view,
23656                        &v_view,
23657                        &mut ws.attn_out[rank],
23658                        head_dim,
23659                        local_heads,
23660                        local_kv_heads,
23661                        t_kv,
23662                        geometry.attention_scale(),
23663                        k_tok_bytes_c,
23664                        v_tok_bytes_c,
23665                        false,
23666                    )?;
23667                    engine.attn_head_gate(
23668                        &ws.attn_out[rank],
23669                        &ws.gate[rank],
23670                        &mut ws.gated[rank],
23671                        None,
23672                        head_dim,
23673                        local_heads,
23674                        1,
23675                    )?;
23676                } else {
23677                    engine.fa_decode_kvmod(
23678                        &ws.q[rank],
23679                        &k_view,
23680                        &v_view,
23681                        &mut ws.gated[rank],
23682                        head_dim,
23683                        local_heads,
23684                        local_kv_heads,
23685                        t_kv,
23686                        geometry.attention_scale(),
23687                        k_tok_bytes_c,
23688                        v_tok_bytes_c,
23689                        false,
23690                    )?;
23691                }
23692            }
23693
23694            // MEMRA_TCOL_OPROJ defer: the verify driver armed a column — stash this
23695            // column's `gated` rows and skip the per-column finish choreography entirely
23696            // (the batched b4_tcol + join runs after every column). The returned buffer
23697            // is UNWRITTEN in that mode (oproj-tail precedent); the driver reads the
23698            // stashed flag, never this buffer. Ineligible configs fall back to the
23699            // normal finish and the driver consumes the real `mixed` per column.
23700            let output = if let Some(col) = fa2_col.filter(|_| dcw) {
23701                // SPEC_FA2 stash: q + gate rows to the fa2 slabs; fa, o_proj and the
23702                // finish all run in the joined pass. Returned buffer is UNWRITTEN
23703                // (oproj-defer precedent — the walk reads the stash flag, never this).
23704                tp.runtime.decode_v2_stash_fa2(ws, e, col)?;
23705                crate::tp::set_spec_fa2_stashed();
23706                e.uninit(ws.o_out)?
23707            } else if let Some(col) = crate::tp::take_tcol_oproj_defer() {
23708                if tp.runtime.decode_v2_oproj_tcol_eligible(ws, &tp.o) {
23709                    tp.runtime.decode_v2_stash_gated(ws, e, col)?;
23710                    crate::tp::set_tcol_oproj_stashed();
23711                    e.uninit(ws.o_out)?
23712                } else {
23713                    tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23714                }
23715            } else {
23716                tp.runtime.decode_v2_finish(ws, e, &tp.o)?
23717            };
23718
23719            // Local shadow append: reads ws.k_shadow/ws.v_shadow on e's stream, which
23720            // decode_v2_finish ordered behind the root event. Same math and cache state
23721            // transitions as v1.
23722            let local = cache.kv[il]
23723                .as_mut()
23724                .ok_or_else(|| format!("Step TP layer {il} lost its local KV cache"))?;
23725            if local.len != base_len || base_len + 1 > max_ctx {
23726                return Err(format!(
23727                    "Step TP layer {il} local cache changed during decode: \
23728                     len={} base={base_len} max={max_ctx}",
23729                    local.len
23730                )
23731                .into());
23732            }
23733            if crate::tp::no_local_shadow_on() {
23734                // Lengths advance, contents stay stale (graph-door precedent: decode reads
23735                // only the distributed TP caches; local contents feed spec/MTP scratch).
23736                local.len = base_len + 1;
23737                // MEMRA_LEN_MIRROR_LAZY=1: skip the 4B in-stream htod — nothing reads the
23738                // LOCAL device mirror in TP decode (the dcw fa reads the RANK counters),
23739                // and each tiny copy costs a compute->copy engine turnaround mid-layer.
23740                if !crate::tp::len_mirror_lazy_on() {
23741                    e.set_i32_one(&mut local.len_d, local.len as i32)?;
23742                }
23743            } else {
23744                let retain_from = window
23745                    .map(|window| {
23746                        let staged_retain = (base_len + 1).saturating_sub(window) & !31usize;
23747                        let rollback_retain =
23748                            base_len.saturating_sub(window.saturating_sub(1)) & !31usize;
23749                        staged_retain.min(rollback_retain)
23750                    })
23751                    .unwrap_or(0);
23752                let write_row = e.prepare_kv_append(local, retain_from, 1)?;
23753                e.append_kv_quantized(
23754                    &ws.k_shadow,
23755                    &ws.v_shadow,
23756                    &mut local.k,
23757                    &mut local.v,
23758                    write_row,
23759                    local.kv_dim_k,
23760                    local.kv_dim_v,
23761                    local.k_tok_bytes,
23762                    local.v_tok_bytes,
23763                    false,
23764                )?;
23765                local.len = base_len + 1;
23766                e.set_i32_one(&mut local.len_d, local.len as i32)?;
23767            }
23768            Ok(output)
23769        })();
23770
23771        let output = match staged {
23772            Ok(output) => output,
23773            Err(error) => {
23774                let _ = tp.runtime.rollback_tp_kv_transaction(
23775                    cache.tp_kv[il]
23776                        .as_mut()
23777                        .expect("distributed cache checked above"),
23778                    transaction,
23779                );
23780                if let Some(local) = cache.kv[il].as_mut() {
23781                    local.len = base_len;
23782                    let _ = e.set_i32_one(&mut local.len_d, base_len as i32);
23783                }
23784                return Err(error);
23785            }
23786        };
23787        // MEMRA_LEN_MIRROR_LAZY under FUSE_ROPE_APPEND: the fused append atomicInc owns
23788        // the rank counters (same value as the absolute re-set on full accept), so commit
23789        // host bookkeeping only — kills two 4B in-stream htods per layer. Non-fused dcw
23790        // keeps the absolute set (its appends do NOT inc).
23791        let lazy_commit = fuse_rope && crate::tp::len_mirror_lazy_on();
23792        if lazy_commit {
23793            if let Err(error) = tp.runtime.commit_tp_kv_transaction_external(
23794                cache.tp_kv[il]
23795                    .as_mut()
23796                    .expect("distributed cache checked above"),
23797                transaction,
23798                1,
23799            ) {
23800                let _ = tp.runtime.rollback_tp_kv_transaction(
23801                    cache.tp_kv[il]
23802                        .as_mut()
23803                        .expect("distributed cache checked above"),
23804                    transaction,
23805                );
23806                let local = cache.kv[il].as_mut().expect("local cache checked above");
23807                local.len = base_len;
23808                e.set_i32_one(&mut local.len_d, base_len as i32)?;
23809                return Err(error);
23810            }
23811        } else if let Err(error) = tp.runtime.commit_tp_kv_transaction(
23812            cache.tp_kv[il]
23813                .as_mut()
23814                .expect("distributed cache checked above"),
23815            transaction,
23816            1,
23817        ) {
23818            let _ = tp.runtime.rollback_tp_kv_transaction(
23819                cache.tp_kv[il]
23820                    .as_mut()
23821                    .expect("distributed cache checked above"),
23822                transaction,
23823            );
23824            let local = cache.kv[il].as_mut().expect("local cache checked above");
23825            local.len = base_len;
23826            e.set_i32_one(&mut local.len_d, base_len as i32)?;
23827            return Err(error);
23828        }
23829
23830        let committed = cache.tp_kv[il]
23831            .as_ref()
23832            .expect("distributed cache checked above")
23833            .committed_len();
23834        let local_len = cache.kv[il]
23835            .as_ref()
23836            .expect("local cache checked above")
23837            .len;
23838        if committed != local_len {
23839            return Err(format!(
23840                "Step TP layer {il} committed cache length {committed} != local shadow {local_len}"
23841            )
23842            .into());
23843        }
23844        static V2_LOGGED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
23845        if !V2_LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) {
23846            eprintln!(
23847                "[step-tp-attn-v2] execute layer={} devices={:?} tokens=1 driver=v2 \
23848                 qkv_tensor_parallel=true qk_norm_rank_local=true rope_rank_local=true \
23849                 kv_cache_distributed=true kv_cache_hydrated={hydrated} \
23850                 attention_tensor_parallel=true attention_scope={} \
23851                 input_path=root-device-replicated gate={} gate_tensor_parallel={} \
23852                 gate_shards={} o_tensor_parallel=true o_reduce=root-device \
23853                 local_cache_shadow=true cache_commit=immediate transport={} native_p2p=true \
23854                 bulk_p2p={} workspace=persistent ordering=evented output=e-device \
23855                 performance_claim=false (logged once; every decode layer runs this driver)",
23856                tp.layer,
23857                tp.devices,
23858                if window.is_some() {
23859                    "rank-local-swa-ring"
23860                } else {
23861                    "rank-local-global"
23862                },
23863                has_gate,
23864                use_gate_shards,
23865                if use_gate_shards {
23866                    "device-staged"
23867                } else if has_gate {
23868                    "root-staged"
23869                } else {
23870                    "none"
23871                },
23872                tp.runtime.transport_label(),
23873                tp.runtime.bulk_p2p(),
23874            );
23875        }
23876        Ok(output)
23877    }
23878
23879    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
23880    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
23881    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
23882    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
23883    /// requiring `attn_gate`).
23884    ///
23885    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
23886    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
23887    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
23888    #[allow(clippy::too_many_arguments)]
23889    pub(crate) fn step35_decode_attn(
23890        &self,
23891        e: &Engine,
23892        fa: &FullAttnLayer,
23893        il: usize,
23894        h: &CudaSlice<f32>,
23895        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
23896        pos_d: &CudaSlice<i32>,
23897        cache: &mut Cache,
23898    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
23899        if fa
23900            .step_tp_qkv
23901            .as_ref()
23902            .is_some_and(|tp| tp.attention.is_some())
23903        {
23904            if pre_q.is_some() {
23905                return Err(
23906                    "rank-local Step attention preserves BF16 activations and refuses the q8_1 \
23907                     pre-quantized decode path"
23908                        .into(),
23909                );
23910            }
23911            return self.step35_tp_decode_attn_resident(e, fa, il, h, pos_d, cache);
23912        }
23913
23914        let geometry = self.step35_geom(il);
23915        let hd = geometry.head_dim_k as usize;
23916        let nkv = geometry.n_head_kv as usize;
23917        let nh = geometry.n_head as usize;
23918        let rbase = geometry.rope_base;
23919        let scale = geometry.attention_scale();
23920        let swa = geometry.window.is_some();
23921        let eps = self.cfg.rms_eps;
23922        let win = geometry.window.unwrap_or(0) as usize;
23923        let n_rot = geometry.n_rot as usize;
23924        let n_embd = self.cfg.n_embd as usize;
23925        let gw = fa
23926            .attn_gate
23927            .as_ref()
23928            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
23929
23930        let tp_qkv = if fa.step_tp_qkv.is_some() {
23931            if pre_q.is_some() {
23932                return Err(
23933                    "Step Q/K/V TP preserves BF16 activations and refuses the q8_1 \
23934                     pre-quantized decode path"
23935                        .into(),
23936                );
23937            }
23938            self.full_attn_tp_qkv(e, fa, h, 1)?
23939        } else {
23940            None
23941        };
23942
23943        let (q0, k0, v0, gt) = match tp_qkv {
23944            Some(mut g3) => {
23945                let v = g3.pop().unwrap();
23946                let k = g3.pop().unwrap();
23947                let q = g3.pop().unwrap();
23948                let gt = e.matmul(gw, h, 1)?;
23949                (q, k, v, gt)
23950            }
23951            None => match pre_q {
23952                Some((hq, hdq)) => {
23953                    debug_assert!(
23954                        e.uses_q8_1_fast(gw),
23955                        "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
23956                         (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
23957                    );
23958                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
23959                        Some(t3) => t3,
23960                        None => (
23961                            e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
23962                            e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
23963                            e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
23964                        ),
23965                    };
23966                    let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
23967                    (a, b, c, gt)
23968                }
23969                None => {
23970                    if e.uses_q8_1_fast(&fa.wq)
23971                        && e.uses_q8_1_fast(&fa.wk)
23972                        && e.uses_q8_1_fast(&fa.wv)
23973                        && e.uses_q8_1_fast(gw)
23974                    {
23975                        let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
23976                        let (a, b, c) =
23977                            match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
23978                                Some(t3) => t3,
23979                                None => (
23980                                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
23981                                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
23982                                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
23983                                ),
23984                            };
23985                        let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
23986                        (a, b, c, gt)
23987                    } else {
23988                        (
23989                            e.matmul(&fa.wq, h, 1)?,
23990                            e.matmul(&fa.wk, h, 1)?,
23991                            e.matmul(&fa.wv, h, 1)?,
23992                            e.matmul(gw, h, 1)?,
23993                        )
23994                    }
23995                }
23996            },
23997        };
23998
23999        let mut q = e.uninit(nh * hd)?;
24000        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
24001        let mut k = e.uninit(nkv * hd)?;
24002        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
24003        let ff = if swa {
24004            None
24005        } else {
24006            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
24007        };
24008        #[cfg(debug_assertions)]
24009        if let Some(ff) = ff {
24010            crate::debug_assert_tensor_stream_device(
24011                ff,
24012                &e.stream(),
24013                "step35_decode_attn.rope_freqs",
24014            );
24015        }
24016        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
24017
24018        if std::env::var("MEMRA_NOFA").is_ok() {
24019            return Err(
24020                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
24021                        cache; unset MEMRA_NOFA to use fa_decode"
24022                    .into(),
24023            );
24024        }
24025        let kvl = cache.kv[il].as_mut().unwrap();
24026        let next_len = kvl.len + 1;
24027        let (off, t_kv) = if swa && next_len > win {
24028            (next_len - win, win)
24029        } else {
24030            (0, next_len)
24031        };
24032        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
24033        e.append_kv_quantized(
24034            &k,
24035            &v0,
24036            &mut kvl.k,
24037            &mut kvl.v,
24038            write_row,
24039            kvl.kv_dim_k,
24040            kvl.kv_dim_v,
24041            kvl.k_tok_bytes,
24042            kvl.v_tok_bytes,
24043            crate::Engine::kv_fp8_on(),
24044        )?;
24045        kvl.len = next_len;
24046        let physical = kvl.physical_rows(off, off + t_kv)?;
24047        let k_view = e.view_u8_range(
24048            &kvl.k,
24049            physical.start * kvl.k_tok_bytes,
24050            physical.end * kvl.k_tok_bytes,
24051        );
24052        let v_view = e.view_u8_range(
24053            &kvl.v,
24054            physical.start * kvl.v_tok_bytes,
24055            physical.end * kvl.v_tok_bytes,
24056        );
24057        let mut attn = e.uninit(nh * hd)?;
24058        e.fa_decode_kvmod(
24059            &q,
24060            &k_view,
24061            &v_view,
24062            &mut attn,
24063            hd,
24064            nh,
24065            nkv,
24066            t_kv,
24067            scale,
24068            kvl.k_tok_bytes,
24069            kvl.v_tok_bytes,
24070            crate::Engine::kv_fp8_on(),
24071        )?;
24072
24073        let mut ag = e.uninit(nh * hd)?;
24074        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
24075        self.full_attn_o(e, fa, &ag, 1)
24076    }
24077}
24078
24079// ===================================================================================== //
24080//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
24081//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
24082//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
24083//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
24084//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
24085//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
24086// ===================================================================================== //
24087impl HybridModel {
24088    pub fn is_gemma4_e4b(&self) -> bool {
24089        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
24090    }
24091
24092    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
24093    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
24094    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
24095    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
24096        let g = self.cfg.gemma4.as_ref().unwrap();
24097        let swa = g.swa_pattern[il];
24098        let hd = if swa {
24099            g.key_length_swa
24100        } else {
24101            g.key_length_global
24102        } as usize;
24103        let Mixer::Full(fa) = &self.layers[il].mixer else {
24104            panic!("e4b layer {il} not full-attn")
24105        };
24106        let nh = fa.wq.out_features() / hd;
24107        let nkv = fa.wk.out_features() / hd;
24108        (
24109            hd,
24110            nkv,
24111            nh,
24112            if swa {
24113                g.rope_base_swa
24114            } else {
24115                g.rope_base_global
24116            },
24117            1.0,
24118            swa,
24119        )
24120    }
24121
24122    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
24123    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
24124        self.layers[il]
24125            .gemma4
24126            .as_ref()
24127            .and_then(|b| b.e4b.as_ref())
24128            .and_then(|e4| e4.kv_share.map(|t| t as usize))
24129    }
24130
24131    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
24132    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
24133    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
24134    ///     (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
24135    fn gemma4_e4b_inp_pl(
24136        &self,
24137        e: &Engine,
24138        tokens: &[u32],
24139        x_scaled: &CudaSlice<f32>,
24140        t: usize,
24141    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24142        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
24143        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
24144    }
24145
24146    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
24147    fn gemma4_e4b_inp_pl_dev(
24148        &self,
24149        e: &Engine,
24150        tok_d: &CudaSlice<u32>,
24151        x_scaled: &CudaSlice<f32>,
24152        t: usize,
24153    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24154        let aux = self.gemma4_aux.as_ref().unwrap();
24155        let m = aux.e4b.as_ref().unwrap();
24156        let n_embd = self.cfg.n_embd as usize;
24157        let n_layer = self.layers.len();
24158        let width = m.n_epl * n_layer;
24159        let tbl = m.tok_tbl_gpu.get_or_init(|| {
24160            e.upload_u8(&m.tok_embd_bytes)
24161                .expect("e4b per-layer token table upload")
24162        });
24163        let mut a =
24164            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
24165        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
24166        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
24167        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
24168        let mut pn = e.uninit(t * width)?;
24169        e.rms_norm(
24170            &p,
24171            m.proj_norm.float_data(),
24172            &mut pn,
24173            m.n_epl,
24174            t * n_layer,
24175            self.cfg.rms_eps,
24176        )?;
24177        let mut out = e.uninit(t * width)?;
24178        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
24179        Ok(out)
24180    }
24181
24182    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
24183    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
24184    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
24185    /// already holds this forward's rows — the target runs earlier in the stack).
24186    #[allow(clippy::too_many_arguments)]
24187    fn gemma4_e4b_attn(
24188        &self,
24189        e: &Engine,
24190        il: usize,
24191        hq: &CudaSlice<i8>,
24192        hdq: &CudaSlice<f32>,
24193        pos_d: &CudaSlice<i32>,
24194        t: usize,
24195        cache: &mut Cache,
24196        dc_bucket: Option<usize>,
24197    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
24198        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
24199        let eps = self.cfg.rms_eps;
24200        let aux = self.gemma4_aux.as_ref().unwrap();
24201        let ones = aux.ones(e);
24202        #[cfg(debug_assertions)]
24203        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
24204        let Mixer::Full(fa) = &self.layers[il].mixer else {
24205            unreachable!()
24206        };
24207        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
24208        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
24209        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
24210        let h0 = e.zeros(0)?;
24211        let h = &h0;
24212
24213        let ff = if swa {
24214            None
24215        } else {
24216            Some(
24217                aux.rope_freqs(e)
24218                    .expect("e4b global rope needs rope_freqs.weight"),
24219            )
24220        };
24221        #[cfg(debug_assertions)]
24222        if let Some(ff) = ff {
24223            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
24224        }
24225        let share = self.gemma4_e4b_kv_target(il);
24226        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
24227        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
24228        let mut q;
24229        if let Some(_tgt) = share {
24230            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
24231            q = e.uninit(t * nh * hd)?;
24232            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
24233            // empty; q0 stands in for the unused k/v pointers).
24234            let mut kdummy = e.uninit(1)?;
24235            let mut vdummy = e.uninit(1)?;
24236            e.rms_norm_qkv_rope(
24237                &q0,
24238                &q0,
24239                &q0,
24240                fa.q_norm.float_data(),
24241                fa.q_norm.float_data(),
24242                ones,
24243                &mut q,
24244                &mut kdummy,
24245                &mut vdummy,
24246                hd,
24247                self.gemma4_rope_dims(il),
24248                nh * t,
24249                0,
24250                pos_d,
24251                nh,
24252                1,
24253                base,
24254                1.0,
24255                ff,
24256                eps,
24257            )?;
24258        } else {
24259            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
24260            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
24261            // q|k|v rows — the cat norm+rope twin consumes it directly.
24262            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
24263            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
24264            q = e.uninit(t * nh * hd)?;
24265            let mut k = e.uninit(t * nkv * hd)?;
24266            let mut v = e.uninit(t * nkv * hd)?;
24267            if t == 1 && cat.is_some() {
24268                #[allow(clippy::unnecessary_unwrap)]
24269                // allow: the Some-guard sits in a multi-clause regime gate; if-let would reshape the arm structure
24270                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
24271                e.rms_norm_qkv_rope_cat(
24272                    &qkv0,
24273                    fa.q_norm.float_data(),
24274                    fa.k_norm.float_data(),
24275                    ones,
24276                    &mut q,
24277                    &mut k,
24278                    &mut v,
24279                    hd,
24280                    self.gemma4_rope_dims(il),
24281                    nh,
24282                    nkv,
24283                    pos_d,
24284                    nh,
24285                    nkv,
24286                    base,
24287                    1.0,
24288                    ff,
24289                    eps,
24290                )?;
24291            } else {
24292                let (q0, k0, v0) = match if t == 1 {
24293                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
24294                } else {
24295                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
24296                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
24297                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24298                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
24299                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
24300                    } else {
24301                        None
24302                    }
24303                } {
24304                    Some(triple) => triple,
24305                    None => (
24306                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
24307                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
24308                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
24309                    ), // E4B: real v (K != V)
24310                };
24311                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
24312                // the normed rows; V ones-rms, never roped).
24313                e.rms_norm_qkv_rope(
24314                    &q0,
24315                    &k0,
24316                    &v0,
24317                    fa.q_norm.float_data(),
24318                    fa.k_norm.float_data(),
24319                    ones,
24320                    &mut q,
24321                    &mut k,
24322                    &mut v,
24323                    hd,
24324                    self.gemma4_rope_dims(il),
24325                    nh * t,
24326                    nkv * t,
24327                    pos_d,
24328                    nh,
24329                    nkv,
24330                    base,
24331                    1.0,
24332                    ff,
24333                    eps,
24334                )?;
24335            }
24336            let kvl = cache.kv[il].as_mut().unwrap();
24337            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
24338            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
24339            // degenerate tok-0 stream, 2026-07-12).
24340            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24341            if dc_bucket.is_some() {
24342                // DC arm (graph serving): append at the len_d slot, advance the counter
24343                // in-stream — replay-correct, no host len in the launch args. Host mirrors
24344                // are NOT touched here (the replay loop owns them; a bump at capture-record
24345                // time would double-count the capture iteration).
24346                debug_assert!(t == 1);
24347                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
24348                e.append_kv_quantized_row_dc_inc(
24349                    &k,
24350                    &v,
24351                    &mut kvl.k,
24352                    &mut kvl.v,
24353                    &mut kvl.len_d,
24354                    kvl.kv_dim_k,
24355                    kvl.kv_dim_v,
24356                    kvl.k_tok_bytes,
24357                    kvl.v_tok_bytes,
24358                    cls,
24359                )?;
24360            } else {
24361                e.append_kv_quantized_rows(
24362                    &k,
24363                    &v,
24364                    &mut kvl.k,
24365                    &mut kvl.v,
24366                    kvl.len,
24367                    t,
24368                    kvl.kv_dim_k,
24369                    kvl.kv_dim_v,
24370                    kvl.k_tok_bytes,
24371                    kvl.v_tok_bytes,
24372                    cls,
24373                )?;
24374                kvl.len += t;
24375            }
24376            kv_f32 = Some((k, v));
24377        }
24378        // attention: per-row causal fa over the (own or target) quantized cache. The cache
24379        // already contains this forward's rows in both arms; row i attends [.., base+i].
24380        let kvl_idx = share.unwrap_or(il);
24381        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
24382        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
24383        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
24384        let mut attn = e.uninit(t * nh * hd)?;
24385        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
24386        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
24387        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
24388        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
24389        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
24390        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
24391        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
24392        //     rows (the T=K verify kernel; the target appended this forward's rows already).
24393        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
24394        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
24395        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
24396        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
24397            if let Some((kf, vf)) = &kv_f32 {
24398                if hd == 256 && t <= win {
24399                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24400                    return e.matmul(&fa.wo, &attn, t);
24401                }
24402                if hd == 256 && swa && t > win {
24403                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24404                    return e.matmul(&fa.wo, &attn, t);
24405                }
24406                if hd == 512 && !swa {
24407                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24408                    return e.matmul(&fa.wo, &attn, t);
24409                }
24410            } else if share.is_some() {
24411                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24412                let k_view = e.view_u8(&kvl.k, kvl.k.len());
24413                let v_view = e.view_u8(&kvl.v, kvl.v.len());
24414                if hd == 256 && (!swa || t <= win) {
24415                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
24416                    e.fa_prefill_view(
24417                        &q,
24418                        &k_view,
24419                        &v_view,
24420                        &mut attn,
24421                        hd,
24422                        nh,
24423                        nkv,
24424                        t,
24425                        t,
24426                        scale,
24427                        true,
24428                        kvl.k_tok_bytes,
24429                        kvl.v_tok_bytes,
24430                        g,
24431                    )?;
24432                    return e.matmul(&fa.wo, &attn, t);
24433                }
24434                // remaining shared classes (swa above the window; hd512 globals): dequant
24435                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
24436                let kv_dim = nkv * hd;
24437                let mut kf = e.uninit(t * kv_dim)?;
24438                let mut vf = e.uninit(t * kv_dim)?;
24439                e.fa_dequant_kv_view_f32(
24440                    &k_view,
24441                    &v_view,
24442                    &mut kf,
24443                    &mut vf,
24444                    kv_dim,
24445                    kv_dim,
24446                    t,
24447                    kvl.k_tok_bytes,
24448                    kvl.v_tok_bytes,
24449                    g,
24450                )?;
24451                if hd == 512 {
24452                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
24453                } else {
24454                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
24455                }
24456                return e.matmul(&fa.wo, &attn, t);
24457            }
24458        }
24459        if let Some(bucket) = dc_bucket {
24460            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
24461            // fa_decode_dc over the live counter. len_d already advanced past this token
24462            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
24463            // counter (advanced when the target ran earlier in the stack).
24464            assert!(t == 1);
24465            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
24466            // and under the window every live t_kv sits below it — cap the capture bucket
24467            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
24468            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
24469            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
24470            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
24471                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
24472            } else {
24473                bucket
24474            };
24475            let k_view = e.view_u8(&kvl.k, kvl.k.len());
24476            let v_view = e.view_u8(&kvl.v, kvl.v.len());
24477            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
24478            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
24479            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
24480            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
24481            // captured into the dc graph like any other launch. Extending the cascade to
24482            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
24483            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
24484            // MEMRA_WPF=0 rollback seam.
24485            if crate::Engine::wpf_level() >= 1 {
24486                e.prefetch_weight_l2(&fa.wo)?;
24487            }
24488            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
24489            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
24490            if e.uses_q8_1_fast(&fa.wo) {
24491                let mut oq = e.alloc_i8_uninit(nh * hd)?;
24492                let mut od = e.zeros(nh * hd / 32)?;
24493                e.fa_decode_dc_q8(
24494                    &q,
24495                    &k_view,
24496                    &v_view,
24497                    &mut attn,
24498                    hd,
24499                    nh,
24500                    nkv,
24501                    &kvl.len_d,
24502                    bucket,
24503                    scale,
24504                    kvl.k_tok_bytes,
24505                    kvl.v_tok_bytes,
24506                    g,
24507                    Some((&mut oq, &mut od)),
24508                )?;
24509                return e.matmul_pre(&fa.wo, &oq, &od, &attn, t);
24510            }
24511            e.fa_decode_dc(
24512                &q,
24513                &k_view,
24514                &v_view,
24515                &mut attn,
24516                hd,
24517                nh,
24518                nkv,
24519                &kvl.len_d,
24520                bucket,
24521                scale,
24522                kvl.k_tok_bytes,
24523                kvl.v_tok_bytes,
24524                g,
24525            )?;
24526            return e.matmul(&fa.wo, &attn, t);
24527        }
24528        for i in 0..t {
24529            let avail = base_len + i + 1;
24530            let (off_tok, t_kv) = if swa && avail > win {
24531                (avail - win, win)
24532            } else {
24533                (0, avail)
24534            };
24535            let k_view = e.view_u8_range(
24536                &kvl.k,
24537                off_tok * kvl.k_tok_bytes,
24538                (off_tok + t_kv) * kvl.k_tok_bytes,
24539            );
24540            let v_view = e.view_u8_range(
24541                &kvl.v,
24542                off_tok * kvl.v_tok_bytes,
24543                (off_tok + t_kv) * kvl.v_tok_bytes,
24544            );
24545            let qv = e.view(&q, t * nh * hd);
24546            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
24547            let mut q_one = e.uninit(nh * hd)?;
24548            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
24549            let mut a_one = e.uninit(nh * hd)?;
24550            // read class MUST match the append class (globals are e4m3 under gkv): the
24551            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
24552            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
24553            e.fa_decode_kvmod(
24554                &q_one,
24555                &k_view,
24556                &v_view,
24557                &mut a_one,
24558                hd,
24559                nh,
24560                nkv,
24561                t_kv,
24562                scale,
24563                kvl.k_tok_bytes,
24564                kvl.v_tok_bytes,
24565                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
24566            )?;
24567            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
24568        }
24569        e.matmul(&fa.wo, &attn, t)
24570    }
24571
24572    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
24573    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
24574    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
24575    /// layer; does NOT advance cache.pos (caller owns pos).
24576    fn gemma4_e4b_trunk(
24577        &self,
24578        e: &Engine,
24579        tokens: &[u32],
24580        pos0: usize,
24581        cache: &mut Cache,
24582        head_last: bool,
24583    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24584        let n_embd = self.cfg.n_embd as usize;
24585        let t = tokens.len();
24586        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24587        let pos_d = e.htod_i32(&pos)?;
24588        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
24589        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24590        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
24591        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
24592    }
24593
24594    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
24595    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
24596    /// eager chain by construction: SAME functions, not twins).
24597    #[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
24598    fn gemma4_e4b_trunk_core(
24599        &self,
24600        e: &Engine,
24601        x_in: CudaSlice<f32>,
24602        inp_pl: CudaSlice<f32>,
24603        pos_d: &CudaSlice<i32>,
24604        t: usize,
24605        cache: &mut Cache,
24606        dc_bucket: Option<usize>,
24607        cap_logits: bool,
24608        head_last: bool,
24609    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24610        let n_embd = self.cfg.n_embd as usize;
24611        let eps = self.cfg.rms_eps;
24612        let n_layer = self.layers.len();
24613        let mut x = x_in;
24614        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
24615        let n_epl = aux_e4b.n_epl;
24616
24617        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
24618        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
24619        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
24620        // head rides matmul_pre too. First layer's pair comes from a standalone fused
24621        // norm+quant.
24622        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
24623        for il in 0..n_layer {
24624            let layer = &self.layers[il];
24625            let (hq, hdq) = match h_carry.take() {
24626                Some(p) => p,
24627                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
24628            };
24629            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
24630            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
24631            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
24632            let bits = layer.gemma4.as_ref().unwrap();
24633            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
24634            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
24635            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
24636            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
24637            // the fused single-phase reduction is NOT FP-order-identical to the unfused
24638            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
24639            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
24640            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
24641            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
24642            // gate dropped, decode AND verify ride the same fused chain — parity by
24643            // construction, VERIFY-GATE 0.000e0.
24644            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
24645            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
24646                e,
24647                layer,
24648                &o,
24649                &x,
24650                t,
24651                Some(layer.post_attn_norm.float_data()),
24652                fuse_exit,
24653            )?;
24654            let mut resid = e.uninit(t * n_embd)?;
24655            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
24656            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
24657            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
24658            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
24659            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
24660            let g = if fuse_exit {
24661                // sn here = RAW f0 (post_ffw deferred).
24662                let (rq, rd) = e.rms_pre_add_q8_1(
24663                    &sn,
24664                    bits.post_ffw_norm.float_data(),
24665                    &attn_out,
24666                    &mut resid,
24667                    n_embd,
24668                    t,
24669                    self.cfg.rms_eps,
24670                )?;
24671                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
24672            } else {
24673                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
24674                e.matmul(&e4b.inp_gate, &resid, t)?
24675            };
24676            let mut act = e.uninit(t * n_epl)?;
24677            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
24678                let ipv = e.view(&inp_pl, n_epl * n_layer);
24679                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
24680                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
24681                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
24682            } else {
24683                let mut inp_this = e.uninit(t * n_epl)?;
24684                e.copy_rows_strided(
24685                    &inp_pl,
24686                    &mut inp_this,
24687                    n_epl,
24688                    t,
24689                    n_epl * n_layer,
24690                    il * n_epl,
24691                )?;
24692                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
24693                e.matmul(&e4b.proj, &act, t)?
24694            };
24695            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
24696            // ONE launch (glue-fusion lane; last layer emits through output_norm).
24697            let next_norm = if il + 1 < n_layer {
24698                self.layers[il + 1].attn_norm.float_data()
24699            } else {
24700                self.output_norm.float_data()
24701            };
24702            let mut xn = e.uninit(t * n_embd)?;
24703            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
24704                &y,
24705                e4b.post_norm.float_data(),
24706                &resid,
24707                bits.layer_scale,
24708                next_norm,
24709                &mut xn,
24710                n_embd,
24711                t,
24712                eps,
24713            )?;
24714            h_carry = Some(pair);
24715            x = xn;
24716        }
24717        // the head consumes the last layer's fused (output_norm) emit. head_last callers
24718        // (prime, last_only forward) need only the final row's logits — the all-T head is
24719        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
24720        let (oq, odq) = h_carry.take().unwrap();
24721        let h0 = e.zeros(0)?;
24722        let hm = if head_last { 1 } else { t };
24723        let (hq, hd) = if head_last && t > 1 {
24724            let mut q1 = e.uninit_i8(n_embd)?;
24725            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
24726            let nb = n_embd / 32;
24727            let mut d1 = e.uninit(nb)?;
24728            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
24729            (q1, d1)
24730        } else {
24731            (oq, odq)
24732        };
24733        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
24734        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
24735        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
24736        // Logit-returning callers (host logits / spec prime) keep the capped emit.
24737        if cap_logits {
24738            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
24739            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
24740        }
24741        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
24742        Ok((ld, x))
24743    }
24744
24745    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
24746    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
24747    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
24748    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
24749    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
24750    /// covers exactly the layers that appended).
24751    pub fn gemma4_e4b_decode_step_t_am_dev(
24752        &self,
24753        e: &Engine,
24754        tok_d: &CudaSlice<u32>,
24755        t: usize,
24756        pos0: usize,
24757        cache: &mut Cache,
24758    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24759        let n_embd = self.cfg.n_embd as usize;
24760        let eps = self.cfg.rms_eps;
24761        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
24762        let pos_d = e.htod_i32(&pos)?;
24763        let embd_gpu = self
24764            .embd_gpu
24765            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
24766        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
24767        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
24768        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
24769        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
24770        let (ld, xp) =
24771            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
24772        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
24773        // emit is already capped, matching the eager chain bit-for-bit).
24774        let n_vocab = self.output.out_features();
24775        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
24776        for i in 0..t {
24777            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
24778        }
24779        let mut hn = e.uninit(t * n_embd)?;
24780        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24781        cache.pos += t;
24782        Ok((vam, hn))
24783    }
24784
24785    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
24786    /// prime path — mirror of `gemma4_decode_step_t_h`).
24787    pub(crate) fn gemma4_e4b_decode_step_t_h(
24788        &self,
24789        e: &Engine,
24790        tokens: &[u32],
24791        pos0: usize,
24792        cache: &mut Cache,
24793    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24794        let n_embd = self.cfg.n_embd as usize;
24795        let eps = self.cfg.rms_eps;
24796        let t = tokens.len();
24797        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
24798        let mut hn = e.uninit(t * n_embd)?;
24799        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
24800        cache.pos += t;
24801        Ok((e.dtoh(&ld)?, hn))
24802    }
24803
24804    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
24805    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
24806    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
24807    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
24808    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
24809    #[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
24810    pub fn gemma4_e4b_decode_step_dcg(
24811        &self,
24812        e: &Engine,
24813        token_d: &mut CudaSlice<u32>,
24814        pos_d: &mut CudaSlice<i32>,
24815        embd_gpu: &CudaSlice<u8>,
24816        embd_qt: i32,
24817        embd_rb: usize,
24818        cache: &mut Cache,
24819        n_vocab: usize,
24820        bucket: usize,
24821    ) -> Result<(), Box<dyn std::error::Error>> {
24822        let n_embd = self.cfg.n_embd as usize;
24823        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
24824        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
24825        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
24826        let (ld, _x) =
24827            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
24828        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
24829        e.inc_seqlen(pos_d)?;
24830        Ok(())
24831    }
24832
24833    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
24834    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
24835    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
24836    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
24837    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
24838    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
24839    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
24840    #[allow(clippy::too_many_arguments)]
24841    pub fn gemma4_e4b_decode_step_dc(
24842        &self,
24843        e: &Engine,
24844        token_d: &CudaSlice<u32>,
24845        pos_d: &mut CudaSlice<i32>,
24846        embd_gpu: &CudaSlice<u8>,
24847        embd_qt: i32,
24848        embd_rb: usize,
24849        cache: &mut Cache,
24850        n_vocab: usize,
24851    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
24852        let n_embd = self.cfg.n_embd as usize;
24853        let eps = self.cfg.rms_eps;
24854        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
24855        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
24856        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
24857        let (ld, _x) =
24858            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
24859        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
24860        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
24861        e.inc_seqlen(pos_d)?;
24862        cache.pos += 1;
24863        let _ = eps;
24864        Ok(tok_out)
24865    }
24866
24867    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
24868    /// pre-output_norm hidden). Advances cache.pos.
24869    pub(crate) fn gemma4_e4b_decode_step_h(
24870        &self,
24871        e: &Engine,
24872        token: u32,
24873        cache: &mut Cache,
24874    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24875        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
24876        let logits = e.dtoh(&ld)?;
24877        cache.pos += 1;
24878        Ok((logits, x))
24879    }
24880
24881    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
24882    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
24883    /// fast; the prefill fa arms come later.
24884    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
24885    pub(crate) fn gemma4_e4b_prime(
24886        &self,
24887        e: &Engine,
24888        tokens: &[u32],
24889        cache: &mut Cache,
24890    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24891        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
24892        // process-kill as gemma4_prime — refuse per-request.
24893        if cache.pos != 0 {
24894            return Err(
24895                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
24896                        call or decode tokenwise"
24897                    .into(),
24898            );
24899        }
24900        let n_embd = self.cfg.n_embd as usize;
24901        let t = tokens.len();
24902        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
24903        cache.pos += t;
24904        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
24905        let xv = e.view(&x, t * n_embd);
24906        let row = xv.slice((t - 1) * n_embd..t * n_embd);
24907        let mut h_seed = e.uninit(n_embd)?;
24908        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
24909        Ok((last, h_seed, x))
24910    }
24911
24912    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
24913    pub(crate) fn gemma4_e4b_forward(
24914        &self,
24915        e: &Engine,
24916        tokens: &[u32],
24917        last_only: bool,
24918    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
24919        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
24920        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
24921        e.dtoh(&ld) // head_last already reduced to the final row when last_only
24922    }
24923}
24924
24925#[cfg(test)]
24926mod prime_chunk_schedule_tests {
24927    use super::{
24928        PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, PrimePpSignal, PrimePpStageChannels, PrimePpWaveCredits,
24929        PrimePpWaveSlot, active_matrix_values, align_prime_ranges_to_gdn,
24930        dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring,
24931        move_prime_cache_layers, parse_step_ep_grouped_prefill, parse_step_tp_prefill,
24932        prime_cache_stage_for_layer, recv_prime_pp_signal, restore_prime_cache_layers,
24933        step_grouped_decode_shape, step_grouped_prefill_shape, step_tp_prefill_shape,
24934        validate_step_prime_batch_modes,
24935    };
24936
24937    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
24938        ranges.iter().map(|(start, end)| end - start).collect()
24939    }
24940
24941    #[allow(clippy::manual_clamp)] // allow: the min/max chain mirrors the reference arithmetic order in pinned sizing/quant math
24942    fn auto_chunk(t: usize) -> usize {
24943        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
24944    }
24945
24946    #[test]
24947    fn ppn_prime_cache_partition_moves_and_restores_every_layer() {
24948        let round_trip = |fence: &[usize], layers: usize| {
24949            let original: Vec<Option<usize>> = (0..layers).map(Some).collect();
24950            let mut parent = original.clone();
24951            let mut stages: Vec<Vec<Option<usize>>> =
24952                (0..fence.len() - 1).map(|_| vec![None; layers]).collect();
24953
24954            move_prime_cache_layers(&mut parent, &mut stages, fence);
24955            assert!(parent.iter().all(Option::is_none));
24956            for layer in 0..layers {
24957                let owner = prime_cache_stage_for_layer(fence, layer);
24958                for (stage, values) in stages.iter().enumerate() {
24959                    assert_eq!(values[layer], (stage == owner).then_some(layer));
24960                }
24961            }
24962
24963            restore_prime_cache_layers(&mut parent, &mut stages, fence);
24964            assert_eq!(parent, original);
24965            assert!(stages.iter().flatten().all(Option::is_none));
24966        };
24967
24968        // Layers beyond the trunk fence end model MTP/tail state and remain last-stage owned.
24969        round_trip(&[0, 5, 8], 10);
24970        round_trip(&[0, 2, 5, 8], 10);
24971        round_trip(&[0, 1, 3, 6, 8], 10);
24972    }
24973
24974    #[test]
24975    fn ppn_prime_wave_credit_requires_the_exact_oldest_wave_and_slot() {
24976        let mut credits = PrimePpWaveCredits::default();
24977        let wave0 = PrimePpWaveSlot { wave: 0, slot: 1 };
24978        let wave1 = PrimePpWaveSlot { wave: 1, slot: 0 };
24979        credits.record_send(wave0).unwrap();
24980        assert_eq!(credits.release_required(), None);
24981        credits.record_send(wave1).unwrap();
24982        assert_eq!(credits.release_required(), Some(wave0));
24983
24984        assert!(
24985            credits
24986                .record_release(PrimePpWaveSlot { wave: 0, slot: 0 })
24987                .unwrap_err()
24988                .contains("does not match oldest pending")
24989        );
24990        assert_eq!(credits.release_required(), Some(wave0));
24991        credits.record_release(wave0).unwrap();
24992        credits
24993            .record_send(PrimePpWaveSlot { wave: 2, slot: 1 })
24994            .unwrap();
24995        assert!(
24996            credits
24997                .record_send(PrimePpWaveSlot { wave: 4, slot: 0 })
24998                .unwrap_err()
24999                .contains("while wave 3 was next")
25000        );
25001        assert!(
25002            credits
25003                .record_send(PrimePpWaveSlot { wave: 3, slot: 1 })
25004                .unwrap_err()
25005                .contains("reused slot 1")
25006        );
25007    }
25008
25009    #[test]
25010    fn ppn_prime_wave_signal_reports_order_error_injected_error_and_closure() {
25011        let expected = PrimePpWaveSlot { wave: 2, slot: 1 };
25012
25013        let (sender, receiver) = std::sync::mpsc::channel();
25014        sender.send(PrimePpSignal::Slot(expected)).unwrap();
25015        assert_eq!(
25016            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap(),
25017            expected
25018        );
25019
25020        let (sender, receiver) = std::sync::mpsc::channel();
25021        sender
25022            .send(PrimePpSignal::Slot(PrimePpWaveSlot { wave: 3, slot: 1 }))
25023            .unwrap();
25024        assert!(
25025            recv_prime_pp_signal(&receiver, expected, true, "test")
25026                .unwrap_err()
25027                .contains("expected wave/slot")
25028        );
25029
25030        let (sender, receiver) = std::sync::mpsc::channel();
25031        sender
25032            .send(PrimePpSignal::Error("injected stage failure".into()))
25033            .unwrap();
25034        assert_eq!(
25035            recv_prime_pp_signal(&receiver, expected, true, "test").unwrap_err(),
25036            "injected stage failure"
25037        );
25038
25039        let (upstream_sender, upstream_receiver) = std::sync::mpsc::channel();
25040        let (outgoing_sender, outgoing_receiver) = std::sync::mpsc::channel();
25041        let (_release_sender, released_downstream) = std::sync::mpsc::channel();
25042        PrimePpStageChannels {
25043            incoming: None,
25044            release_upstream: Some(upstream_sender),
25045            outgoing: outgoing_sender,
25046            released_downstream,
25047        }
25048        .notify_failure("injected worker error");
25049        assert_eq!(
25050            recv_prime_pp_signal(&upstream_receiver, expected, false, "test").unwrap_err(),
25051            "injected worker error"
25052        );
25053        assert_eq!(
25054            recv_prime_pp_signal(&outgoing_receiver, expected, false, "test").unwrap_err(),
25055            "injected worker error"
25056        );
25057
25058        let (sender, receiver) = std::sync::mpsc::channel::<PrimePpSignal>();
25059        drop(sender);
25060        assert!(
25061            recv_prime_pp_signal(&receiver, expected, true, "test")
25062                .unwrap_err()
25063                .contains("channel closed while waiting for wave 2")
25064        );
25065    }
25066
25067    /// TOOTH for the PP-auto-ranges GDN grid law (lane/hermes-perf-fixes, 2026-08-23;
25068    /// primegrid pattern): the AUTO schedules put internal prime-call boundaries OFF the
25069    /// WY-chunk grid — the broken arm must be demonstrably off-grid, and the aligned twin
25070    /// must land every boundary on it without changing coverage.
25071    #[test]
25072    fn auto_prime_ranges_align_to_the_gdn_grid() {
25073        let c = 32usize; // shipped MEMRA_GDN_CHUNK default/clamp floor
25074        let assert_covers = |ranges: &[(usize, usize)], t: usize| {
25075            assert_eq!(ranges.first().map(|&(s, _)| s), Some(0));
25076            assert_eq!(ranges.last().map(|&(_, e)| e), Some(t));
25077            for w in ranges.windows(2) {
25078                assert_eq!(w[0].1, w[1].0, "ranges must stay contiguous");
25079            }
25080            assert!(ranges.iter().all(|&(s, e)| e > s), "no empty range");
25081        };
25082
25083        // The PP-2 auto geometry at a real agentic length: t=9510 -> fill = 1189 (div_ceil
25084        // by 8), every internal boundary off the 32 grid — the falsified-identity arm.
25085        let t = 9510usize;
25086        let fill = auto_chunk(t);
25087        let fixed = fixed_prime_chunk_ranges(t, fill);
25088        assert!(
25089            fixed[..fixed.len() - 1].iter().any(|&(_, e)| e % c != 0),
25090            "broken arm vanished: fixed auto boundaries all landed on-grid"
25091        );
25092        let dynamic = dynamic_prime_chunk_ranges(t, fill, &fixed);
25093        assert!(
25094            dynamic[..dynamic.len() - 1]
25095                .iter()
25096                .any(|&(_, e)| e % c != 0),
25097            "broken arm vanished: dynamic auto boundaries all landed on-grid"
25098        );
25099
25100        for ranges in [&fixed, &dynamic] {
25101            let aligned = align_prime_ranges_to_gdn(ranges, t, c);
25102            assert_covers(&aligned, t);
25103            for &(_, e) in &aligned[..aligned.len() - 1] {
25104                assert_eq!(e % c, 0, "internal boundary {e} off the {c}-grid");
25105            }
25106            // boundaries only move DOWN, at most c-1 tokens.
25107            for (&(_, a), &(_, b)) in aligned.iter().zip(ranges.iter()) {
25108                assert!(a <= b && b - a < c);
25109            }
25110        }
25111
25112        // Collapse/merge: boundaries inside one grid cell fuse instead of emitting an
25113        // empty range; the schedule survives degenerate short fills.
25114        let tight = vec![(0usize, 33usize), (33, 40), (40, 200)];
25115        let aligned = align_prime_ranges_to_gdn(&tight, 200, c);
25116        assert_covers(&aligned, 200);
25117        assert_eq!(aligned, vec![(0, 32), (32, 200)]);
25118
25119        // No-ops: single range, c=0 (grid off), already-aligned schedules.
25120        assert_eq!(align_prime_ranges_to_gdn(&[(0, 200)], 200, c), [(0, 200)]);
25121        assert_eq!(align_prime_ranges_to_gdn(&tight, 200, 0), tight.as_slice());
25122        let on_grid = vec![(0usize, 128usize), (128, 256), (256, 300)];
25123        assert_eq!(
25124            align_prime_ranges_to_gdn(&on_grid, 300, c),
25125            on_grid.as_slice()
25126        );
25127    }
25128
25129    #[test]
25130    fn active_matrix_prefix_scopes_reused_prime_slabs() {
25131        assert_eq!(
25132            active_matrix_values(40 * 4096, 29, 4096, "activation").unwrap(),
25133            29 * 4096
25134        );
25135        assert_eq!(
25136            active_matrix_values(29 * 4096, 29, 4096, "activation").unwrap(),
25137            29 * 4096
25138        );
25139        assert_eq!(
25140            active_matrix_values(29 * 4096, 24, 4096, "activation").unwrap(),
25141            24 * 4096
25142        );
25143        assert!(active_matrix_values(28 * 4096, 29, 4096, "activation").is_err());
25144        assert!(active_matrix_values(usize::MAX, usize::MAX, 2, "activation").is_err());
25145    }
25146
25147    #[test]
25148    fn step_tp_prefill_batch_refuses_before_scheduler_fallback() {
25149        assert!(validate_step_prime_batch_modes(false, false).is_ok());
25150
25151        let grouped_without_tp = validate_step_prime_batch_modes(false, true).unwrap_err();
25152        assert!(grouped_without_tp.contains("requires MEMRA_STEP_TP_PREFILL=1"));
25153
25154        for grouped in [false, true] {
25155            let err = validate_step_prime_batch_modes(true, grouped).unwrap_err();
25156            assert!(err.contains("did not clear the live-server performance gate"));
25157            assert!(err.contains("per-session grouped prefill"));
25158        }
25159    }
25160
25161    #[test]
25162    fn step_grouped_path_is_eager_single_token_only() {
25163        assert!(step_grouped_decode_shape(false, 1));
25164        assert!(!step_grouped_decode_shape(true, 1));
25165        assert!(!step_grouped_decode_shape(false, 2));
25166        assert!(!step_grouped_decode_shape(true, 2));
25167    }
25168
25169    #[test]
25170    fn step_grouped_prefill_door_is_strict_and_capacity_bounded() {
25171        assert!(!parse_step_ep_grouped_prefill(None).unwrap());
25172        assert!(!parse_step_ep_grouped_prefill(Some("")).unwrap());
25173        assert!(!parse_step_ep_grouped_prefill(Some("0")).unwrap());
25174        assert!(parse_step_ep_grouped_prefill(Some("1")).unwrap());
25175        assert!(parse_step_ep_grouped_prefill(Some("true")).is_err());
25176        assert!(parse_step_ep_grouped_prefill(Some("2")).is_err());
25177
25178        assert!(step_grouped_prefill_shape(true, true, PRIME_MIN_T));
25179        assert!(step_grouped_prefill_shape(
25180            true,
25181            true,
25182            crate::cache::PRIME_CHUNK_MAX_TOKENS,
25183        ));
25184        assert!(!step_grouped_prefill_shape(true, true, PRIME_MIN_T - 1,));
25185        assert!(!step_grouped_prefill_shape(
25186            true,
25187            true,
25188            crate::cache::PRIME_CHUNK_MAX_TOKENS + 1,
25189        ));
25190        assert!(!step_grouped_prefill_shape(false, true, PRIME_MIN_T));
25191        assert!(!step_grouped_prefill_shape(true, false, PRIME_MIN_T));
25192    }
25193
25194    #[test]
25195    fn step_tp_prefill_door_is_strict_and_default_off() {
25196        assert!(!parse_step_tp_prefill(None).unwrap());
25197        assert!(!parse_step_tp_prefill(Some("")).unwrap());
25198        assert!(!parse_step_tp_prefill(Some("0")).unwrap());
25199        assert!(parse_step_tp_prefill(Some("1")).unwrap());
25200        assert!(parse_step_tp_prefill(Some("true")).is_err());
25201        assert!(parse_step_tp_prefill(Some("2")).is_err());
25202    }
25203
25204    #[test]
25205    fn step_tp_prefill_requires_a_qualified_even_rank_shape() {
25206        assert!(step_tp_prefill_shape(
25207            true,
25208            PRIME_MIN_T,
25209            4,
25210            true,
25211            true,
25212            false,
25213        ));
25214        assert!(!step_tp_prefill_shape(
25215            false,
25216            PRIME_MIN_T,
25217            4,
25218            true,
25219            true,
25220            false,
25221        ));
25222        assert!(!step_tp_prefill_shape(
25223            true,
25224            PRIME_MIN_T - 1,
25225            4,
25226            true,
25227            true,
25228            false,
25229        ));
25230        // TP2 admits (2026-08-25); odd/1-card placements still refuse.
25231        assert!(step_tp_prefill_shape(
25232            true,
25233            PRIME_MIN_T,
25234            2,
25235            true,
25236            true,
25237            false
25238        ));
25239        assert!(!step_tp_prefill_shape(
25240            true,
25241            PRIME_MIN_T,
25242            1,
25243            true,
25244            true,
25245            false
25246        ));
25247        assert!(!step_tp_prefill_shape(
25248            true,
25249            PRIME_MIN_T,
25250            3,
25251            true,
25252            true,
25253            false
25254        ));
25255        assert!(!step_tp_prefill_shape(
25256            true,
25257            PRIME_MIN_T,
25258            4,
25259            false,
25260            true,
25261            false,
25262        ));
25263        assert!(!step_tp_prefill_shape(
25264            true,
25265            PRIME_MIN_T,
25266            4,
25267            true,
25268            false,
25269            false,
25270        ));
25271        assert!(!step_tp_prefill_shape(
25272            true,
25273            PRIME_MIN_T,
25274            4,
25275            true,
25276            true,
25277            true,
25278        ));
25279    }
25280
25281    #[test]
25282    fn fixed_schedule_retains_measured_geometry() {
25283        assert_eq!(
25284            sizes(&fixed_prime_chunk_ranges(461, 128)),
25285            vec![128, 128, 128, 77]
25286        );
25287        assert_eq!(
25288            sizes(&fixed_prime_chunk_ranges(1833, 230)),
25289            vec![230, 230, 230, 230, 230, 230, 230, 223]
25290        );
25291        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
25292        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
25293        assert_eq!(capped, vec![4096, 4088, 16]);
25294        assert!(capped.iter().all(|&rows| rows <= 4096));
25295        assert_eq!(
25296            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
25297            vec![4100],
25298            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
25299        );
25300    }
25301
25302    #[test]
25303    fn dynamic_schedule_matches_registered_shapes() {
25304        let cases = [
25305            (461, vec![64, 141, 132, 124]),
25306            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
25307            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
25308        ];
25309        for (t, expected) in cases {
25310            let chunk = auto_chunk(t);
25311            let fixed = fixed_prime_chunk_ranges(t, chunk);
25312            assert_eq!(
25313                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
25314                expected
25315            );
25316        }
25317    }
25318
25319    #[test]
25320    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
25321        for t in 256..=8192 {
25322            let chunk = auto_chunk(t);
25323            let fixed = fixed_prime_chunk_ranges(t, chunk);
25324            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
25325            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
25326            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
25327            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
25328            for pair in dynamic.windows(2) {
25329                assert_eq!(pair[0].1, pair[1].0, "T={t}");
25330            }
25331            assert!(
25332                dynamic
25333                    .iter()
25334                    .all(|(start, end)| end - start >= PRIME_MIN_T),
25335                "T={t} sizes={:?}",
25336                sizes(&dynamic)
25337            );
25338            if dynamic.len() >= 3 {
25339                let chunk_sizes = sizes(&dynamic);
25340                assert!(
25341                    chunk_sizes[0] < chunk_sizes[1],
25342                    "T={t} sizes={chunk_sizes:?}"
25343                );
25344                assert!(
25345                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
25346                    "T={t} sizes={chunk_sizes:?}"
25347                );
25348            }
25349        }
25350    }
25351}
25352
25353#[cfg(test)]
25354mod page_prefetch_tests {
25355    use super::{
25356        grouped_worker_prefetch_position, page_prefetch_positions,
25357        page_prefetch_window_from_values, worker_prefetch_positions,
25358    };
25359
25360    #[test]
25361    fn page_prefetch_window_keeps_existing_opt_in_default() {
25362        assert_eq!(page_prefetch_window_from_values(false, None), 0);
25363        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
25364        assert_eq!(page_prefetch_window_from_values(true, None), 1);
25365        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
25366        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
25367        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
25368    }
25369
25370    #[test]
25371    fn rolling_page_prefetch_advises_each_future_expert_once() {
25372        let advised: Vec<_> = (0..7)
25373            .flat_map(|position| page_prefetch_positions(position, 7, 3))
25374            .collect();
25375        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
25376
25377        let one_ahead: Vec<_> = (0..4)
25378            .flat_map(|position| page_prefetch_positions(position, 4, 1))
25379            .collect();
25380        assert_eq!(one_ahead, vec![1, 2, 3]);
25381        assert!(page_prefetch_positions(0, 4, 0).is_empty());
25382    }
25383
25384    #[test]
25385    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
25386        assert_eq!(grouped_worker_prefetch_position(0, None), None);
25387        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
25388            .chain(
25389                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
25390            )
25391            .collect();
25392        assert_eq!(positions, vec![0, 1, 2, 3]);
25393        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
25394    }
25395
25396    #[test]
25397    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
25398        let queued: Vec<_> = (0..8)
25399            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
25400            .collect();
25401        assert_eq!(queued, (0..8).collect::<Vec<_>>());
25402
25403        let one_at_a_time: Vec<_> = (0..4)
25404            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
25405            .collect();
25406        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
25407        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
25408    }
25409}
25410
25411pub struct G4DcSlots {
25412    x: CudaSlice<f32>,
25413    xn: CudaSlice<f32>,
25414    cur: CudaSlice<f32>,
25415    hq: CudaSlice<i8>,
25416    hd_: CudaSlice<f32>,
25417    q0: CudaSlice<f32>,
25418    k0: CudaSlice<f32>,
25419    v0: CudaSlice<f32>,
25420    q: CudaSlice<f32>,
25421    k: CudaSlice<f32>,
25422    v: CudaSlice<f32>,
25423    attn: CudaSlice<f32>,
25424    o: CudaSlice<f32>,
25425    attn_out: CudaSlice<f32>,
25426    zsh: CudaSlice<f32>,
25427    zq: CudaSlice<i8>,
25428    zd: CudaSlice<f32>,
25429    gate: CudaSlice<f32>,
25430    up: CudaSlice<f32>,
25431    act: CudaSlice<f32>,
25432    actq: CudaSlice<i8>,
25433    actd: CudaSlice<f32>,
25434    f0: CudaSlice<f32>,
25435    sn: CudaSlice<f32>,
25436    hn: CudaSlice<f32>,
25437    logits: CudaSlice<f32>,
25438}
25439
25440/// Whole-token decode graph state (step TP graph increment B). One stitched multi-device
25441/// parent per fa bucket, plus the persistent host->graph plumbing: the device token id the
25442/// in-graph embed gathers, the device position the ropes read (advanced in-graph), and the
25443/// fixed logits stage the head writes.
25444pub struct Step35TokenGraphState {
25445    /// (bucket_max, graph) — bucket keyed by the fa split geometry (fa_geom_eager).
25446    pub graphs: Vec<(usize, crate::tp::TokenGraph)>,
25447    pub token_d: cudarc::driver::CudaSlice<u32>,
25448    pub pos_d: cudarc::driver::CudaSlice<i32>,
25449    pub logits_stage: cudarc::driver::CudaSlice<f32>,
25450    /// Cross-child intermediates MUST live at fixed addresses (graph mem nodes remap at
25451    /// launch, so an alloc made inside one captured child is not referable from another):
25452    /// the running residual, the post-attention pair, the shared-expert row, and the
25453    /// e-context mirrors of the root-produced attention output and K/V shadow rows.
25454    pub x: cudarc::driver::CudaSlice<f32>,
25455    pub x1: cudarc::driver::CudaSlice<f32>,
25456    pub mixed_stage: cudarc::driver::CudaSlice<f32>,
25457    pub sh_stage: cudarc::driver::CudaSlice<f32>,
25458    pub k_shadow_stage: cudarc::driver::CudaSlice<f32>,
25459    pub v_shadow_stage: cudarc::driver::CudaSlice<f32>,
25460    /// Alloc-free e-section scratch (child graphs cannot contain mem nodes): router logits,
25461    /// shared-expert gate/up/act rows + sigmoid scalar, dense-FFN z/gate/up/act, head hidden.
25462    pub router_logits: cudarc::driver::CudaSlice<f32>,
25463    pub shexp_gate: cudarc::driver::CudaSlice<f32>,
25464    pub shexp_up: cudarc::driver::CudaSlice<f32>,
25465    pub shexp_act: cudarc::driver::CudaSlice<f32>,
25466    pub gate_sig: cudarc::driver::CudaSlice<f32>,
25467    pub dense_z: cudarc::driver::CudaSlice<f32>,
25468    pub dense_gate: cudarc::driver::CudaSlice<f32>,
25469    pub dense_up: cudarc::driver::CudaSlice<f32>,
25470    pub dense_act: cudarc::driver::CudaSlice<f32>,
25471    pub hn: cudarc::driver::CudaSlice<f32>,
25472    /// MEMRA_TG_PROBE_LAYER diagnostics: capture-time copies of layer K's attention output
25473    /// and post-FFN residual, dumped after replay for graph-vs-eager layer bisection.
25474    pub probe_mixed: cudarc::driver::CudaSlice<f32>,
25475    pub probe_x: cudarc::driver::CudaSlice<f32>,
25476    /// Chunk-loop (F-lite): device token history ring + its device write index, filled by
25477    /// the in-graph tail argmax chain; host reads the ring once per chunk.
25478    pub token_hist: cudarc::driver::CudaSlice<u32>,
25479    pub hist_idx: cudarc::driver::CudaSlice<i32>,
25480}
25481
25482impl HybridModel {
25483    /// Whole-token decode graph (step TP graph increment B): ONE stitched multi-device parent
25484    /// replays the entire 45-layer token — the launch-collapse the per-layer minis could not
25485    /// reach. Returns Some(logits) when the graph handled the token, None for eager fallback
25486    /// (door off, ineligible class, any layer's ring would rebase, or a bucket boundary that
25487    /// needs a rebuild this token).
25488    ///
25489    /// v1 SCOPE (diagnostic door, default OFF): the LOCAL shadow caches advance their lengths
25490    /// but not their contents under this door (the TP rank caches are fully maintained
25491    /// in-graph via the dcw counters). Sessions relying on shadow CONTENT (save/rollback)
25492    /// must not run with the door on until the local-dcw twin lands.
25493    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
25494    pub(crate) fn step35_token_graph_step(
25495        &self,
25496        e: &Engine,
25497        token: u32,
25498        cache: &mut Cache,
25499    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
25500        if !self.uses_sliding_gated_moe_program()
25501            || !crate::tp::step_tp_graph_enabled()?
25502            || !crate::tp::step_tp_dcw_enabled()?
25503            || !crate::tp::step_tp_qkv_fused_enabled()?
25504            || !crate::tp::step_tp_dev_router_enabled()?
25505            || !crate::tp::step_nvfp4_dev_routes_enabled()?
25506        {
25507            return Ok(None);
25508        }
25509        // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): the eager
25510        // token step is this route's byte-identical twin — warmup and rebase tokens
25511        // already ride it — so below the driver-free floor the token goes eager
25512        // (`Ok(None)` = the caller's eager fallback) instead of feeding cuGraphLaunch
25513        // an exhausted card (lane/graph-launch-guard-sweep-20260831).
25514        if !crate::spec::graph_launch_headroom_ok(e) {
25515            static NOTED: std::sync::Once = std::sync::Once::new();
25516            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
25517            return Ok(None);
25518        }
25519        let n_embd = self.cfg.n_embd as usize;
25520        let n_vocab = self.cfg.n_vocab as usize;
25521        let n_layers = self.layers.len();
25522        let pos = cache.pos;
25523        let staged_next = pos + 1;
25524        if staged_next < 96 {
25525            return Ok(None); // sub-vec-floor contexts keep eager (fa kernel-class boundary)
25526        }
25527
25528        // Per-layer eligibility: every TP layer contiguous-appends this token (any rebase ->
25529        // eager fallback for the whole token; the host path also updates base_d there).
25530        for il in 0..n_layers {
25531            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
25532                return Ok(None); // caches not hydrated yet — eager warms them
25533            };
25534            if tp_kv.peek_append_ring(1)?.1 {
25535                return Ok(None);
25536            }
25537        }
25538
25539        // Bucket key: the global layers' fa split geometry at this depth (SWA layers cap at
25540        // their window and share one bucket forever after ctx > window).
25541        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
25542        if !fa_vec {
25543            return Ok(None);
25544        }
25545        let sp = crate::fa_split_keys(staged_next, 8);
25546        let bucket_max = (n_splits * sp).max(staged_next);
25547
25548        let mut state_guard = self
25549            .step35_token_graph
25550            .lock()
25551            .map_err(|_| "step35 token graph lock is poisoned")?;
25552        if state_guard.is_none() {
25553            let _main = e.gpu.enter_main()?;
25554            let n_expert = self
25555                .cfg
25556                .moe
25557                .as_ref()
25558                .map(|m| m.expert_count as usize)
25559                .unwrap_or(0);
25560            let n_ff_sh = self
25561                .layers
25562                .iter()
25563                .find_map(|l| match &l.ffn {
25564                    crate::hybrid::Ffn::Moe(m) => m.gate_shexp.as_ref().map(|g| g.out_features()),
25565                    _ => None,
25566                })
25567                .unwrap_or(0);
25568            let n_ff_dense = self
25569                .layers
25570                .iter()
25571                .find_map(|l| match &l.ffn {
25572                    crate::hybrid::Ffn::Dense { ffn_gate, .. } => Some(ffn_gate.out_features()),
25573                    _ => None,
25574                })
25575                .unwrap_or(0);
25576            *state_guard = Some(Step35TokenGraphState {
25577                graphs: Vec::new(),
25578                token_d: e.stream().clone_htod(&[0u32])?,
25579                pos_d: e.htod_i32(&[pos as i32])?,
25580                logits_stage: e.htod(&vec![0.0f32; n_vocab])?,
25581                x: e.htod(&vec![0.0f32; n_embd])?,
25582                x1: e.htod(&vec![0.0f32; n_embd])?,
25583                mixed_stage: e.htod(&vec![0.0f32; n_embd])?,
25584                sh_stage: e.htod(&vec![0.0f32; n_embd])?,
25585                k_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25586                v_shadow_stage: e.htod(&vec![0.0f32; 2048])?,
25587                router_logits: e.htod(&vec![0.0f32; n_expert.max(1)])?,
25588                shexp_gate: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25589                shexp_up: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25590                shexp_act: e.htod(&vec![0.0f32; n_ff_sh.max(1)])?,
25591                gate_sig: e.htod(&[1.0f32; 1])?,
25592                dense_z: e.htod(&vec![0.0f32; n_embd])?,
25593                dense_gate: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25594                dense_up: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25595                dense_act: e.htod(&vec![0.0f32; n_ff_dense.max(1)])?,
25596                hn: e.htod(&vec![0.0f32; n_embd])?,
25597                probe_mixed: e.htod(&vec![0.0f32; n_embd])?,
25598                probe_x: e.htod(&vec![0.0f32; n_embd])?,
25599                token_hist: e.stream().clone_htod(&[0u32; 16])?,
25600                hist_idx: e.htod_i32(&[0])?,
25601            });
25602        }
25603        let state = state_guard.as_mut().expect("state armed above");
25604        // Pre-arm the argmax partials OUTSIDE any capture: the launcher allocates them on
25605        // first use, and an alloc inside a captured section is a mem node (child graphs
25606        // reject those — the tail argmax chain needs them already resident).
25607        {
25608            let _main = e.gpu.enter_main()?;
25609            let Step35TokenGraphState {
25610                logits_stage,
25611                token_d,
25612                ..
25613            } = &mut *state;
25614            e.argmax_token_device_into(logits_stage, token_d, n_vocab)?;
25615        }
25616
25617        // ONE graph, retargeted per bucket (increment C): the per-16-token whole rebuild was
25618        // ~55ms (3.4ms/token persistent); the M1 exec update path moves nsp/ski/gridDimY and
25619        // the partial-pool memset widths in ~1ms. The partial pool is pre-grown to the run
25620        // ceiling at build so the baked pointers never move.
25621        if state.graphs.is_empty() {
25622            // Build the parent at this bucket. Capture executes nothing; correctness is
25623            // pinned at replay by the token-identity gate.
25624            self.step35_token_graph_build(e, cache, state, bucket_max)?;
25625        }
25626        {
25627            let (b, g) = state.graphs.first_mut().expect("graph built above");
25628            if *b != bucket_max {
25629                g.retarget_bucket(bucket_max)?;
25630                *b = bucket_max;
25631            }
25632        }
25633        let graph = state
25634            .graphs
25635            .first()
25636            .map(|(_, g)| g)
25637            .expect("graph built above");
25638
25639        let tg_timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
25640        let t_fence = tg_timing.then(std::time::Instant::now);
25641        // Rank-stream fence: an eager token (warmup, rebase fallback) leaves len-mirror sets
25642        // queued on the rank streams, and graph children carry no ordering edge to those
25643        // streams. Graph tokens themselves enqueue nothing there (external commit), so this
25644        // sync is a no-op between consecutive replays.
25645        {
25646            let fa0 = match &self.layers[0].mixer {
25647                Mixer::Full(fa) => fa,
25648                _ => return Err("step35 token graph expects full-attention layers".into()),
25649            };
25650            let tp0 = fa0
25651                .step_tp_qkv
25652                .as_ref()
25653                .ok_or("step35 token graph lost its TP state")?;
25654            for rank in 0..tp0.runtime.devices().len() {
25655                let engine = tp0
25656                    .runtime
25657                    .rank_engine(rank)
25658                    .ok_or("step35 token graph lost a rank engine")?;
25659                let _main = engine.gpu.enter_main()?;
25660                engine.stream().synchronize()?;
25661            }
25662        }
25663
25664        // Replay: feed the token, launch, read the logits, mirror the host bookkeeping.
25665        {
25666            let _main = e.gpu.enter_main()?;
25667            e.set_u32_one(&mut state.token_d, token)?;
25668            e.set_i32_one(&mut state.pos_d, pos as i32)?;
25669        }
25670        let t_launch = tg_timing.then(std::time::Instant::now);
25671        graph.launch(e)?;
25672        let t_book = tg_timing.then(std::time::Instant::now);
25673        // Host bookkeeping OVERLAPS the replay (nsys 2026-08-21: the 45-layer txn loop was
25674        // 5.4ms/token of inter-token gap when it ran after the sync). Host-only work except
25675        // the local len_d set, which is stream-ordered AFTER the graph on e's stream. On a
25676        // replay error the counters are already advanced — acceptable: the decode aborts.
25677        for il in 0..n_layers {
25678            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
25679            let transaction = tp_kv.begin_transaction()?;
25680            let fa = match &self.layers[il].mixer {
25681                Mixer::Full(fa) => fa,
25682                _ => return Err("step35 token graph expects full-attention layers".into()),
25683            };
25684            let tp = fa
25685                .step_tp_qkv
25686                .as_ref()
25687                .ok_or("step35 token graph lost its TP state")?;
25688            // Bookkeeping-only txn (external appends; mirror sets skipped — the in-graph
25689            // incs own the counters). Shards unused.
25690            let empty: [CudaSlice<f32>; 0] = [];
25691            tp.runtime.append_tp_kv_transaction_inner(
25692                tp_kv,
25693                transaction,
25694                &empty,
25695                &empty,
25696                1,
25697                true,
25698            )?;
25699            tp.runtime
25700                .commit_tp_kv_transaction_external(tp_kv, transaction, 1)?;
25701            // Local shadow: lengths advance (v1 keeps contents stale under the door).
25702            if let Some(local) = cache.kv[il].as_mut() {
25703                local.len = pos + 1;
25704                let _main = e.gpu.enter_main()?;
25705                e.set_i32_one(&mut local.len_d, (pos + 1) as i32)?;
25706            }
25707        }
25708        cache.pos = pos + 1;
25709        let t_sync = tg_timing.then(std::time::Instant::now);
25710        let (logits, h_seed) = {
25711            let _main = e.gpu.enter_main()?;
25712            e.stream().synchronize()?;
25713            (e.dtoh(&state.logits_stage)?, e.clone_dtod(&state.x)?)
25714        };
25715        if let (Some(f), Some(l), Some(b), Some(sy)) = (t_fence, t_launch, t_book, t_sync) {
25716            use std::sync::atomic::{AtomicU64, Ordering};
25717            static NS: [AtomicU64; 5] = [
25718                AtomicU64::new(0),
25719                AtomicU64::new(0),
25720                AtomicU64::new(0),
25721                AtomicU64::new(0),
25722                AtomicU64::new(0),
25723            ];
25724            static CALLS: AtomicU64 = AtomicU64::new(0);
25725            let now = std::time::Instant::now();
25726            NS[0].fetch_add((l - f).as_nanos() as u64, Ordering::Relaxed); // fence+set
25727            NS[1].fetch_add((b - l).as_nanos() as u64, Ordering::Relaxed); // launch call
25728            NS[2].fetch_add((sy - b).as_nanos() as u64, Ordering::Relaxed); // bookkeeping
25729            NS[3].fetch_add((now - sy).as_nanos() as u64, Ordering::Relaxed); // sync+dtoh
25730            NS[4].fetch_add((now - f).as_nanos() as u64, Ordering::Relaxed); // total
25731            let calls = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
25732            if calls.is_multiple_of(100) {
25733                let avg = |i: usize| NS[i].load(Ordering::Relaxed) as f64 / calls as f64 / 1e3;
25734                eprintln!(
25735                    "[tg-timing] calls={calls} fence_us={:.0} launch_us={:.0} book_us={:.0} \
25736                     syncdtoh_us={:.0} total_us={:.0}",
25737                    avg(0),
25738                    avg(1),
25739                    avg(2),
25740                    avg(3),
25741                    avg(4)
25742                );
25743            }
25744        }
25745        // MEMRA_TG_PROBE_LAYER diagnostics: append the captured layer-K probes.
25746        if std::env::var("MEMRA_TG_PROBE_LAYER").is_ok() {
25747            use std::io::Write;
25748            let (pm, px) = {
25749                let _main = e.gpu.enter_main()?;
25750                (e.dtoh(&state.probe_mixed)?, e.dtoh(&state.probe_x)?)
25751            };
25752            for (path, data) in [
25753                ("/root/tg-probe-mixed.bin", &pm),
25754                ("/root/tg-probe-x.bin", &px),
25755            ] {
25756                let mut fo = std::fs::OpenOptions::new()
25757                    .create(true)
25758                    .append(true)
25759                    .open(path)?;
25760                for v in data {
25761                    fo.write_all(&v.to_le_bytes())?;
25762                }
25763            }
25764        }
25765        // MEMRA_DUMP_HN twin of the eager tail's dump (same format: appended raw LE f32 rows)
25766        // so a graph arm and an eager arm produce position-aligned pre-head hidden streams.
25767        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
25768            let hh = {
25769                let _main = e.gpu.enter_main()?;
25770                e.dtoh(&state.hn)?
25771            };
25772            use std::io::Write;
25773            let mut fo = std::fs::OpenOptions::new()
25774                .create(true)
25775                .append(true)
25776                .open(path)?;
25777            for v in &hh {
25778                fo.write_all(&v.to_le_bytes())?;
25779            }
25780        }
25781        // MEMRA_STEP_TP_GRAPH_DEBUG=1: per-token device-counter dump (drift hunts). One dtoh
25782        // per rank per token; diagnostics only.
25783        if std::env::var("MEMRA_STEP_TP_GRAPH_DEBUG").as_deref() == Ok("1") {
25784            for il in [0usize, 1, 44] {
25785                let tp_kv = cache.tp_kv[il].as_ref().expect("eligibility checked above");
25786                let host_len = tp_kv.staged_len();
25787                let fa = match &self.layers[il].mixer {
25788                    Mixer::Full(fa) => fa,
25789                    _ => continue,
25790                };
25791                let tp = fa
25792                    .step_tp_qkv
25793                    .as_ref()
25794                    .ok_or("step35 token graph lost its TP state")?;
25795                for rank in 0..tp.runtime.devices().len() {
25796                    let engine = tp
25797                        .runtime
25798                        .rank_engine(rank)
25799                        .ok_or("step35 token graph lost a rank engine")?;
25800                    let rank_cache = tp_kv.rank(rank).ok_or("debug rank cache missing")?;
25801                    let _main = engine.gpu.enter_main()?;
25802                    engine.stream().synchronize()?;
25803                    let len_d = engine.dtoh_i32_one(rank_cache.len_d())?;
25804                    let base_d = match rank_cache.base_d() {
25805                        Some(b) => engine.dtoh_i32_one(b)?,
25806                        None => -1,
25807                    };
25808                    eprintln!(
25809                        "[graph-debug] pos={pos} il={il} rank={rank} host_len={host_len} \
25810                         len_d={len_d} base_d={base_d}"
25811                    );
25812                }
25813            }
25814        }
25815        Ok(Some((logits, h_seed)))
25816    }
25817
25818    /// MEMRA_HEAD_SPLIT worker: dev0 computes lm-head rows [0, half), rank1 computes
25819    /// [half, n_vocab) concurrently on its otherwise-idle tail, and the halves concatenate
25820    /// on e — bit-identical per logit to the single-device matvec. Process-static workspace
25821    /// (SHEXP_WS pattern) pinned by the head tensor pointer; rank1 holds a one-time 0.5GB
25822    /// replica of its row half. Returns None when ineligible (no bf16 head / no rank1).
25823    pub(crate) fn head_split_matvec(
25824        &self,
25825        e: &Engine,
25826        hn: &CudaSlice<f32>,
25827    ) -> Result<Option<Vec<f32>>, Box<dyn std::error::Error>> {
25828        if self.head_split_fill_device(e, hn)?.is_none() {
25829            return Ok(None);
25830        }
25831        let guard = HEAD_SPLIT_WS
25832            .lock()
25833            .map_err(|_| "head split lock is poisoned")?;
25834        let ws = guard.as_ref().expect("filled above");
25835        let _main = e.gpu.enter_main()?;
25836        Ok(Some(e.dtoh(&ws.logits_e)?))
25837    }
25838
25839    /// Compute body of the split head: arms the replica + staging on first use, then fills
25840    /// the persistent full-logits row (e's half by view matvec, rank1's half by raw P2P
25841    /// push) and orders e's stream behind it. None = ineligible.
25842    fn head_split_fill_device(
25843        &self,
25844        e: &Engine,
25845        hn: &CudaSlice<f32>,
25846    ) -> Result<Option<()>, Box<dyn std::error::Error>> {
25847        use cudarc::driver::DevicePtr;
25848        let crate::model::GpuTensor::FloatBf16 { data: head, .. } = &self.output else {
25849            return Ok(None);
25850        };
25851        let Some(rank1) = self.layers.first().and_then(|l| match &l.mixer {
25852            Mixer::Full(fa) => fa
25853                .step_tp_qkv
25854                .as_ref()
25855                .and_then(|tp| tp.runtime.rank_engine(1)),
25856            _ => None,
25857        }) else {
25858            return Ok(None);
25859        };
25860        let n_embd = self.cfg.n_embd as usize;
25861        let n_vocab = self.cfg.n_vocab as usize;
25862        let half = n_vocab / 2;
25863        let mut guard = HEAD_SPLIT_WS
25864            .lock()
25865            .map_err(|_| "head split lock is poisoned")?;
25866        let pin = {
25867            let _main = e.gpu.enter_main()?;
25868            let stream = e.stream();
25869            let (ptr, _g) = head.device_ptr(&stream);
25870            ptr
25871        };
25872        if guard.as_ref().is_none_or(|ws| ws.pin != pin) {
25873            // One-time: upload rank1's row half + persistent staging.
25874            let hi_rows = n_vocab - half;
25875            let (w1, hn1, y1, ev_done) = {
25876                let _r1 = rank1.gpu.enter_main()?;
25877                (
25878                    rank1.alloc_u8_uninit(hi_rows * n_embd * 2)?,
25879                    rank1.htod(&vec![0.0f32; n_embd])?,
25880                    rank1.htod(&vec![0.0f32; hi_rows])?,
25881                    rank1.ctx().new_event(None)?,
25882                )
25883            };
25884            {
25885                use cudarc::driver::sys;
25886                let src = pin + (half * n_embd * 2) as u64;
25887                let dst = {
25888                    let _r1 = rank1.gpu.enter_main()?;
25889                    let rstream = rank1.stream();
25890                    let (d, _g) = w1.device_ptr(&rstream);
25891                    d
25892                };
25893                let _r1 = rank1.gpu.enter_main()?;
25894                let r = unsafe {
25895                    sys::cuMemcpyAsync(
25896                        dst as sys::CUdeviceptr,
25897                        src as sys::CUdeviceptr,
25898                        hi_rows * n_embd * 2,
25899                        rank1.stream().cu_stream() as sys::CUstream,
25900                    )
25901                };
25902                if r != sys::CUresult::CUDA_SUCCESS {
25903                    return Err(format!("head split replica upload: {r:?}").into());
25904                }
25905                rank1.stream().synchronize()?;
25906            }
25907            let (logits_e, ev_hn) = {
25908                let _main = e.gpu.enter_main()?;
25909                (e.htod(&vec![0.0f32; n_vocab])?, e.ctx().new_event(None)?)
25910            };
25911            let (raw_hn1, raw_y1) = {
25912                let _r1 = rank1.gpu.enter_main()?;
25913                let rstream = rank1.stream();
25914                let (a, _g0) = hn1.device_ptr(&rstream);
25915                let (b, _g1) = y1.device_ptr(&rstream);
25916                (a, b)
25917            };
25918            let raw_logits_hi = {
25919                let _main = e.gpu.enter_main()?;
25920                let stream = e.stream();
25921                let (l, _g) = logits_e.device_ptr(&stream);
25922                l + (half * 4) as u64
25923            };
25924            *guard = Some(HeadSplit {
25925                pin,
25926                w1,
25927                hn1,
25928                y1,
25929                logits_e,
25930                ev_hn,
25931                ev_done,
25932                raw_hn1,
25933                raw_y1,
25934                raw_logits_hi,
25935                samp: None,
25936            });
25937        }
25938        let ws = guard.as_mut().expect("armed above");
25939        let hi_rows = n_vocab - half;
25940        // e: signal hn ready; rank1: pull hn, matvec its half, push the logits half back.
25941        let raw_hn = {
25942            let _main = e.gpu.enter_main()?;
25943            let stream = e.stream();
25944            let (h, _g) = hn.device_ptr(&stream);
25945            ws.ev_hn.record(&stream)?;
25946            h
25947        };
25948        {
25949            let _r1 = rank1.gpu.enter_main()?;
25950            rank1.stream().wait(&ws.ev_hn)?;
25951            crate::tp::raw_copy_bytes(ws.raw_hn1, raw_hn, n_embd * 4, rank1)?;
25952            let HeadSplit { w1, hn1, y1, .. } = &mut *ws;
25953            rank1.matvec_bf16_into(w1, hn1, y1, n_embd, hi_rows)?;
25954            crate::tp::raw_copy_bytes(ws.raw_logits_hi, ws.raw_y1, hi_rows * 4, rank1)?;
25955            ws.ev_done.record(&rank1.stream())?;
25956        }
25957        {
25958            let _main = e.gpu.enter_main()?;
25959            let head_lo = head.slice(0..half * n_embd * 2);
25960            let HeadSplit { logits_e, .. } = &mut *ws;
25961            // Writes rows [0, half) of logits_e; rank1's raw push fills [half, n_vocab).
25962            e.matvec_bf16_view_into(&head_lo, hn, logits_e, n_embd, half)?;
25963            e.stream().wait(&ws.ev_done)?;
25964            Ok(Some(()))
25965        }
25966    }
25967
25968    /// Device twin of `head_split_matvec` for the chain: fills the persistent full-logits
25969    /// row exactly like the host variant (identical halves, identical concat) and runs the
25970    /// device argmax into `token_d` — NO host readback. Returns false when the split is
25971    /// ineligible (caller falls back to the plain matmul head).
25972    pub(crate) fn head_split_argmax_device(
25973        &self,
25974        e: &Engine,
25975        hn: &CudaSlice<f32>,
25976        token_d: &mut CudaSlice<u32>,
25977    ) -> Result<bool, Box<dyn std::error::Error>> {
25978        if self.head_split_fill_device(e, hn)?.is_none() {
25979            return Ok(false);
25980        }
25981        let n_vocab = self.cfg.n_vocab as usize;
25982        let guard = HEAD_SPLIT_WS
25983            .lock()
25984            .map_err(|_| "head split lock is poisoned")?;
25985        let ws = guard.as_ref().expect("filled above");
25986        let _main = e.gpu.enter_main()?;
25987        e.argmax_token_device_into(&ws.logits_e, token_d, n_vocab)?;
25988        Ok(true)
25989    }
25990
25991    /// SAMPLED twin of `head_split_argmax_device`. The split head already materializes the
25992    /// full concatenated row in `ws.logits_e`, so sampling does NOT have to give up HEAD_SPLIT
25993    /// — it draws from that row on device (filter thresholds, Gumbel perturbation, argmax)
25994    /// exactly as the serve tick does. Worth ~0.2 ms/token: the post-W8 census had the
25995    /// unsplit q8 head at ~364 us against ~82 us per half.
25996    pub(crate) fn head_split_sample_device(
25997        &self,
25998        e: &Engine,
25999        hn: &CudaSlice<f32>,
26000        token_d: &mut CudaSlice<u32>,
26001        samp: &crate::decode_batch::DevSamp,
26002        ctr: u32,
26003    ) -> Result<bool, Box<dyn std::error::Error>> {
26004        if self.head_split_fill_device(e, hn)?.is_none() {
26005            return Ok(false);
26006        }
26007        let n_vocab = self.cfg.n_vocab as usize;
26008        let guard = HEAD_SPLIT_WS
26009            .lock()
26010            .map_err(|_| "head split lock is poisoned")?;
26011        let mut guard = guard;
26012        let ws = guard.as_mut().expect("filled above");
26013        let _main = e.gpu.enter_main()?;
26014        if ws.samp.is_none() {
26015            ws.samp = Some(SampScratch {
26016                pb: e.zeros(n_vocab)?,
26017                th: e.zeros(1)?,
26018                z: e.zeros(1)?,
26019                mx: e.zeros(1)?,
26020                rows: e.htod_i32(&[0i32])?,
26021            });
26022        }
26023        let filtered = samp.top_k > 0 || samp.top_p < 1.0 || samp.min_p > 0.0;
26024        let HeadSplit {
26025            logits_e,
26026            samp: scratch,
26027            ..
26028        } = &mut *ws;
26029        let sc = scratch.as_mut().expect("armed above");
26030        if filtered {
26031            e.filter_stats(
26032                logits_e, n_vocab, &sc.rows, &mut sc.th, &mut sc.z, &mut sc.mx, n_vocab, 1,
26033                samp.temp, samp.top_k, samp.top_p, samp.min_p,
26034            )?;
26035            let SampScratch { pb, th, mx, .. } = sc;
26036            e.gumbel_perturb_filtered_col(
26037                logits_e, 0, pb, n_vocab, samp.seed, ctr, samp.temp, mx, th, 0,
26038            )?;
26039        } else {
26040            e.gumbel_perturb_col(logits_e, 0, &mut sc.pb, n_vocab, samp.seed, ctr, samp.temp)?;
26041        }
26042        e.argmax_token_device_col(&sc.pb, 0, n_vocab, token_d, 0)?;
26043        Ok(true)
26044    }
26045
26046    /// One-per-chunk host readback of the persistent full-logits row (the LAST chain
26047    /// token's row).
26048    pub(crate) fn head_split_logits_dtoh(
26049        &self,
26050        e: &Engine,
26051    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
26052        let guard = HEAD_SPLIT_WS
26053            .lock()
26054            .map_err(|_| "head split lock is poisoned")?;
26055        let ws = guard.as_ref().ok_or("head split logits not armed")?;
26056        let _main = e.gpu.enter_main()?;
26057        e.dtoh(&ws.logits_e)
26058    }
26059
26060    /// Chunk-loop replay (F-lite, MEMRA_STEP_TP_GRAPH_LOOP): run up to `k_target` greedy
26061    /// tokens as back-to-back graph launches chained through the in-graph tail argmax —
26062    /// ONE host sync, ONE history readback, and ONE bulk KV transaction per layer per
26063    /// chunk. Returns None when ineligible (caller falls back to the per-token path).
26064    /// The chunk consumes `token` (already emitted by the caller) as launch 0's input and
26065    /// returns the ids the chain argmax'd (hist[0..k]) plus the LAST launch's logits row —
26066    /// hist[k-1] is exactly argmax(logits), so the caller emits hist[..k-1] and lets its
26067    /// own loop re-derive hist[k-1] from the returned row.
26068    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
26069    pub fn step35_token_graph_chunk(
26070        &self,
26071        e: &Engine,
26072        token: u32,
26073        k_target: usize,
26074        cache: &mut Cache,
26075    ) -> Result<Option<(Vec<u32>, Vec<f32>)>, Box<dyn std::error::Error>> {
26076        if !self.uses_sliding_gated_moe_program()
26077            || !crate::tp::step_tp_graph_enabled()?
26078            || !crate::tp::step_tp_dcw_enabled()?
26079            || !crate::tp::step_tp_qkv_fused_enabled()?
26080            || !crate::tp::step_tp_dev_router_enabled()?
26081            || !crate::tp::step_nvfp4_dev_routes_enabled()?
26082        {
26083            return Ok(None);
26084        }
26085        // GRAPH-LAUNCH HEADROOM GUARD: same guard, same eager twin as
26086        // `step35_token_graph_step` (the chunk is that step replayed k times).
26087        if !crate::spec::graph_launch_headroom_ok(e) {
26088            static NOTED: std::sync::Once = std::sync::Once::new();
26089            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("step-tp-token"));
26090            return Ok(None);
26091        }
26092        let n_layers = self.layers.len();
26093        let pos = cache.pos;
26094        let staged_next = pos + 1;
26095        if staged_next < 96 {
26096            return Ok(None);
26097        }
26098        // Bucket for the FIRST token; the chunk must not cross the bucket boundary (the
26099        // exec's n_splits ladder must match eager per depth).
26100        let (fa_vec, n_splits) = e.fa_geom_eager(staged_next, 128, 8, false);
26101        if !fa_vec {
26102            return Ok(None);
26103        }
26104        let sp = crate::fa_split_keys(staged_next, 8);
26105        let bucket_max = (n_splits * sp).max(staged_next);
26106        let to_boundary = bucket_max.saturating_sub(staged_next) + 1;
26107        let mut k = k_target.min(to_boundary).min(16);
26108        if k < 2 {
26109            return Ok(None);
26110        }
26111        // Every layer must contiguous-append all k rows (no rebase inside the chunk).
26112        for il in 0..n_layers {
26113            let Some(tp_kv) = cache.tp_kv[il].as_ref() else {
26114                return Ok(None);
26115            };
26116            while k >= 2 && tp_kv.peek_append_ring(k)?.1 {
26117                k -= 1;
26118            }
26119            if k < 2 {
26120                return Ok(None);
26121            }
26122        }
26123
26124        let mut state_guard = self
26125            .step35_token_graph
26126            .lock()
26127            .map_err(|_| "step35 token graph lock is poisoned")?;
26128        let Some(state) = state_guard.as_mut() else {
26129            return Ok(None); // per-token path arms the state + stages first
26130        };
26131        if state.graphs.is_empty() {
26132            return Ok(None);
26133        }
26134        {
26135            let (b, g) = state.graphs.first_mut().expect("checked above");
26136            if *b != bucket_max {
26137                g.retarget_bucket(bucket_max)?;
26138                *b = bucket_max;
26139            }
26140        }
26141        let graph = state.graphs.first().map(|(_, g)| g).expect("checked above");
26142
26143        // Rank-stream fence (eager stragglers; see the per-token path).
26144        {
26145            let fa0 = match &self.layers[0].mixer {
26146                Mixer::Full(fa) => fa,
26147                _ => return Err("step35 token graph expects full-attention layers".into()),
26148            };
26149            let tp0 = fa0
26150                .step_tp_qkv
26151                .as_ref()
26152                .ok_or("step35 token graph lost its TP state")?;
26153            for rank in 0..tp0.runtime.devices().len() {
26154                let engine = tp0
26155                    .runtime
26156                    .rank_engine(rank)
26157                    .ok_or("step35 token graph lost a rank engine")?;
26158                let _main = engine.gpu.enter_main()?;
26159                engine.stream().synchronize()?;
26160            }
26161        }
26162
26163        // Seed the chain and fire k launches back-to-back: launch i embeds the token the
26164        // PREVIOUS launch's tail argmax wrote (launch 0 embeds the host-seeded `token`).
26165        {
26166            let _main = e.gpu.enter_main()?;
26167            e.set_u32_one(&mut state.token_d, token)?;
26168            e.set_i32_one(&mut state.pos_d, pos as i32)?;
26169            e.set_i32_one(&mut state.hist_idx, 0)?;
26170        }
26171        for _ in 0..k {
26172            graph.launch(e)?;
26173        }
26174        // Bulk host bookkeeping overlaps the replays: one k-row txn per layer.
26175        for il in 0..n_layers {
26176            let tp_kv = cache.tp_kv[il].as_mut().expect("eligibility checked above");
26177            let transaction = tp_kv.begin_transaction()?;
26178            let fa = match &self.layers[il].mixer {
26179                Mixer::Full(fa) => fa,
26180                _ => return Err("step35 token graph expects full-attention layers".into()),
26181            };
26182            let tp = fa
26183                .step_tp_qkv
26184                .as_ref()
26185                .ok_or("step35 token graph lost its TP state")?;
26186            let empty: [CudaSlice<f32>; 0] = [];
26187            tp.runtime.append_tp_kv_transaction_inner(
26188                tp_kv,
26189                transaction,
26190                &empty,
26191                &empty,
26192                k,
26193                true,
26194            )?;
26195            tp.runtime
26196                .commit_tp_kv_transaction_external(tp_kv, transaction, k)?;
26197            if let Some(local) = cache.kv[il].as_mut() {
26198                local.len = pos + k;
26199                let _main = e.gpu.enter_main()?;
26200                e.set_i32_one(&mut local.len_d, (pos + k) as i32)?;
26201            }
26202        }
26203        cache.pos = pos + k;
26204        let (hist, logits) = {
26205            let _main = e.gpu.enter_main()?;
26206            e.stream().synchronize()?;
26207            (e.dtoh_u32(&state.token_hist)?, e.dtoh(&state.logits_stage)?)
26208        };
26209        Ok(Some((hist[..k].to_vec(), logits)))
26210    }
26211}
26212
26213impl HybridModel {
26214    /// Capture the whole-token parent for one fa bucket. Capture executes nothing; the
26215    /// section closures issue the SAME calls the eager dcw path runs (bit-proven), with the
26216    /// per-token operands living at the fixed stage addresses. Groups: the two rank sections
26217    /// of each phase fork in parallel and merge into the following root section.
26218    #[allow(clippy::too_many_arguments)]
26219    fn step35_token_graph_build(
26220        &self,
26221        e: &Engine,
26222        cache: &mut Cache,
26223        state: &mut Step35TokenGraphState,
26224        bucket_max: usize,
26225    ) -> Result<(), Box<dyn std::error::Error>> {
26226        use cudarc::driver::DevicePtr;
26227        let n_embd = self.cfg.n_embd as usize;
26228        let eps = self.cfg.rms_eps;
26229        let n_layers = self.layers.len();
26230        let started = std::time::Instant::now();
26231        if !crate::router_kernel_on() {
26232            return Err(
26233                "step35 token graph requires the router kernel (MEMRA_ROUTER_KERNEL=0)".into(),
26234            );
26235        }
26236        if !Engine::bf16_mmv_on() || !n_embd.is_multiple_of(8) {
26237            return Err("step35 token graph requires MEMRA_BF16_MMV bf16-resident matvecs".into());
26238        }
26239
26240        // Device embed table (the spec/graph lanes' lazily-uploaded copy).
26241        let embd_gpu = self
26242            .embd_gpu_try(e)
26243            .ok_or("step35 token graph could not upload the device embed table")?;
26244        let embd_qtype = match self.embd.ggml_type {
26245            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
26246            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
26247            other => return Err(format!("token graph embed dtype {other:?} unhandled").into()),
26248        };
26249        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
26250
26251        // Fixed-stage pointers the sections reference.
26252        let (p_mixed, p_kshadow, p_vshadow) = {
26253            let _main = e.gpu.enter_main()?;
26254            let stream = e.stream();
26255            let (a, _g) = state.mixed_stage.device_ptr(&stream);
26256            let (b, _g) = state.k_shadow_stage.device_ptr(&stream);
26257            let (c, _g) = state.v_shadow_stage.device_ptr(&stream);
26258            (a, b, c)
26259        };
26260
26261        crate::tp::token_graph_build_begin()?;
26262        let mut group_id: u32 = 0;
26263        for il in 0..n_layers {
26264            let layer = &self.layers[il];
26265            let fa = match &layer.mixer {
26266                Mixer::Full(fa) => fa,
26267                _ => return Err("step35 token graph expects full-attention layers".into()),
26268            };
26269            let tp = fa
26270                .step_tp_qkv
26271                .as_ref()
26272                .ok_or("step35 token graph lost its TP state")?;
26273            let attention = tp
26274                .attention
26275                .as_ref()
26276                .ok_or("step35 token graph lost its attention aux")?;
26277            let geometry = self.step35_geom(il);
26278            let window = geometry.window.map(|w| w as usize);
26279            let head_dim = geometry.head_dim_k as usize;
26280            let heads = geometry.n_head as usize;
26281            let kv_heads = geometry.n_head_kv as usize;
26282            let ranks = tp.runtime.devices().len();
26283            let local_heads = heads / ranks;
26284            let local_kv_heads = kv_heads / ranks;
26285            let layer_bucket = window.map(|w| bucket_max.min(w)).unwrap_or(bucket_max);
26286            let use_gate_shards =
26287                attention.gate_shards.is_some() || attention.gate_shards_bf16.is_some();
26288            if !use_gate_shards {
26289                return Err("step35 token graph requires the fused gate shards".into());
26290            }
26291
26292            let ws_index = tp
26293                .runtime
26294                .decode_v2_ensure(e, &tp.q, &tp.k, &tp.v, &tp.o, heads)?;
26295            let ws_mutex = tp.runtime.decode_v2_workspace();
26296            let mut ws_guard = ws_mutex
26297                .lock()
26298                .map_err(|_| "step TP decode v2 workspace lock is poisoned")?;
26299            let ws = ws_guard
26300                .get_mut(ws_index)
26301                .ok_or("step TP decode v2 workspace missing after ensure")?;
26302            tp.runtime
26303                .decode_v2_arm_token_mirrors(ws, p_mixed, (p_kshadow, p_vshadow))?;
26304            let mut rope_freqs = Vec::with_capacity(ranks);
26305            for rank in 0..ranks {
26306                let engine = tp
26307                    .runtime
26308                    .rank_engine(rank)
26309                    .ok_or("step35 token graph lost a rank engine")?;
26310                rope_freqs.push(if geometry.rope_factors {
26311                    self.step35_aux
26312                        .as_ref()
26313                        .and_then(|aux| aux.rope_freqs(engine))
26314                } else {
26315                    None
26316                });
26317            }
26318            let gate_shards_arg = if let Some(shards) = attention.gate_shards.as_deref() {
26319                Some(crate::tp::StepTpGateShards::F32(shards))
26320            } else {
26321                attention
26322                    .gate_shards_bf16
26323                    .as_deref()
26324                    .map(crate::tp::StepTpGateShards::Bf16)
26325            };
26326
26327            // ---- E1: embed (layer 0) / attn norm into h_stage + pos copy ----
26328            let decode_input = attention
26329                .decode_input
26330                .as_ref()
26331                .ok_or("step35 token graph requires the replicated decode input")?;
26332            let mut decode_input = decode_input
26333                .lock()
26334                .map_err(|_| "replicated decode input lock is poisoned")?;
26335            // Stage arming happens through the eager stage flow once; require it here.
26336            if ws.h_stage.is_none() {
26337                return Err(
26338                    "step35 token graph requires the stage flow armed (run eager dcw first)".into(),
26339                );
26340            }
26341            {
26342                let state_x = &mut state.x;
26343                let token_d = &state.token_d;
26344                let pos_d = &state.pos_d;
26345                crate::tp::graph_section(e, None, || {
26346                    let _main = e.gpu.enter_main()?;
26347                    if il == 0 {
26348                        e.embed_gather_device_into(
26349                            embd_gpu,
26350                            token_d,
26351                            state_x,
26352                            n_embd,
26353                            embd_qtype,
26354                            embd_row_bytes,
26355                        )?;
26356                    }
26357                    {
26358                        let h_stage = ws.h_stage.as_mut().expect("stage armed checked above");
26359                        e.rms_norm(
26360                            state_x,
26361                            layer.attn_norm.float_data(),
26362                            h_stage,
26363                            n_embd,
26364                            1,
26365                            eps,
26366                        )?;
26367                    }
26368                    {
26369                        let pos_stage = ws.pos_stage.as_mut().expect("stage armed above");
26370                        let mut dst = pos_stage.slice_mut(0..1);
26371                        e.stream().memcpy_dtod(&pos_d.slice(0..1), &mut dst)?;
26372                    }
26373                    Ok(())
26374                })?;
26375            }
26376
26377            // ---- R0/R1 (parallel): projections + dcw attention interior ----
26378            group_id += 1;
26379            for rank in 0..ranks {
26380                let engine = tp
26381                    .runtime
26382                    .rank_engine(rank)
26383                    .ok_or("step35 token graph lost a rank engine")?;
26384                {
26385                    // fa partial pool must reach the RUN CEILING before capture — an
26386                    // in-capture grow is a mem node (child graphs reject those), and the
26387                    // retarget path (increment C) widens the baked memsets up to the ceiling
26388                    // without moving the pool pointers. Two ensures cover both sp rungs.
26389                    let ceiling = window
26390                        .map(|w| cache.max_ctx.min(w))
26391                        .unwrap_or(cache.max_ctx);
26392                    let _main = engine.gpu.enter_main()?;
26393                    engine.fa_dcw_pool_ensure(
26394                        head_dim,
26395                        local_heads,
26396                        local_kv_heads,
26397                        ceiling.min(2048),
26398                    )?;
26399                    engine.fa_dcw_pool_ensure(head_dim, local_heads, local_kv_heads, ceiling)?;
26400                    engine.fa_dcw_pool_ensure(
26401                        head_dim,
26402                        local_heads,
26403                        local_kv_heads,
26404                        layer_bucket,
26405                    )?;
26406                }
26407                let runtime = &tp.runtime;
26408                let q_norm = &attention.q_norm;
26409                let k_norm = &attention.k_norm;
26410                let gate_ref = gate_shards_arg.as_ref();
26411                crate::tp::graph_section(engine, Some(group_id), || {
26412                    runtime.decode_v2_input_qkv_rank(
26413                        ws,
26414                        &state.pos_d,
26415                        &mut decode_input,
26416                        &tp.q,
26417                        &tp.k,
26418                        &tp.v,
26419                        q_norm,
26420                        k_norm,
26421                        head_dim,
26422                        geometry.n_rot as usize,
26423                        geometry.rope_base,
26424                        &rope_freqs,
26425                        eps,
26426                        gate_ref,
26427                        true,
26428                        true,
26429                        false,
26430                        rank,
26431                        None,
26432                    )?;
26433                    // Merged dcw interior at the BUCKET geometry (one-partition law makes the
26434                    // replayed values track the live counters).
26435                    let distributed = cache.tp_kv[il]
26436                        .as_mut()
26437                        .ok_or("step35 token graph lost a TP cache")?;
26438                    let (kv_dim_k, kv_dim_v) = (distributed.kv_dim_k(), distributed.kv_dim_v());
26439                    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
26440                    let capacity = distributed.physical_capacity();
26441                    {
26442                        let rank_cache = distributed
26443                            .rank_mut(rank)
26444                            .ok_or("step35 token graph lost a rank cache")?;
26445                        let (k_plane, v_plane, len_d, base_d) =
26446                            rank_cache.planes_and_counters_mut();
26447                        engine.append_kv_quantized_dcw(
26448                            &ws.k[rank],
26449                            &ws.v_raw[rank],
26450                            k_plane,
26451                            v_plane,
26452                            len_d,
26453                            base_d,
26454                            kv_dim_k,
26455                            kv_dim_v,
26456                            ktb,
26457                            vtb,
26458                        )?;
26459                    }
26460                    {
26461                        let rank_cache = distributed
26462                            .rank_mut(rank)
26463                            .ok_or("step35 token graph lost a rank cache")?;
26464                        engine.inc_i32(rank_cache.len_d_mut())?;
26465                    }
26466                    let rank_cache = distributed
26467                        .rank(rank)
26468                        .ok_or("step35 token graph lost a rank cache")?;
26469                    let k_ring = engine.view_u8_range(rank_cache.k(), 0, capacity * ktb);
26470                    let v_ring = engine.view_u8_range(rank_cache.v(), 0, capacity * vtb);
26471                    // Graph build keeps the UNFUSED combine + gate pair: the bucket
26472                    // retarget addresses combine's nsp at arg slot 6, and the fused
26473                    // combine_gate kernel shifts it to 7 — bit-identical either way, so
26474                    // only the eager arm takes FUSION #2d.
26475                    engine.fa_decode_dcw(
26476                        &ws.q[rank],
26477                        &k_ring,
26478                        &v_ring,
26479                        &mut ws.attn_out[rank],
26480                        head_dim,
26481                        local_heads,
26482                        local_kv_heads,
26483                        rank_cache.len_d(),
26484                        rank_cache.base_d(),
26485                        window.unwrap_or(0),
26486                        layer_bucket,
26487                        geometry.attention_scale(),
26488                        ktb,
26489                        vtb,
26490                        None,
26491                    )?;
26492                    engine.attn_head_gate(
26493                        &ws.attn_out[rank],
26494                        &ws.gate[rank],
26495                        &mut ws.gated[rank],
26496                        None,
26497                        head_dim,
26498                        local_heads,
26499                        1,
26500                    )?;
26501                    runtime.decode_v2_finish_rank_partial(ws, &tp.o, true, rank)?;
26502                    Ok(())
26503                })?;
26504            }
26505
26506            // ---- ROOT: combine + shadows + e-mirrors ----
26507            {
26508                let root = tp
26509                    .runtime
26510                    .rank_engine(0)
26511                    .ok_or("step35 token graph lost the root engine")?;
26512                let runtime = &tp.runtime;
26513                crate::tp::graph_section(root, None, || runtime.decode_v2_finish_root_fused(ws))?;
26514            }
26515            drop(ws_guard);
26516            drop(decode_input);
26517
26518            let probe_layer: Option<usize> = std::env::var("MEMRA_TG_PROBE_LAYER")
26519                .ok()
26520                .and_then(|v| v.parse().ok());
26521            if probe_layer == Some(il) {
26522                let Step35TokenGraphState {
26523                    mixed_stage,
26524                    probe_mixed,
26525                    ..
26526                } = &mut *state;
26527                crate::tp::graph_section(e, None, || {
26528                    let _main = e.gpu.enter_main()?;
26529                    let mut dst = probe_mixed.slice_mut(0..n_embd);
26530                    e.stream()
26531                        .memcpy_dtod(&mixed_stage.slice(0..n_embd), &mut dst)?;
26532                    Ok(())
26533                })?;
26534            }
26535
26536            // ---- FFN half ----
26537            match &layer.ffn {
26538                crate::hybrid::Ffn::Dense {
26539                    ffn_gate,
26540                    ffn_up,
26541                    ffn_down,
26542                } => {
26543                    let n_ff = ffn_gate.out_features();
26544                    let lim = self.cfg.clamp_shexp_at(il as u32);
26545                    // Alloc-free inline of ffn_swiglu_decode's bf16 tail: dual gate/up matvec is
26546                    // bit-identical per row to the two matmul-dispatched matvec_bf16 launches.
26547                    // A clamped dense layer would take eager's q8_1 branch instead -- refuse.
26548                    if lim.is_some() {
26549                        return Err("step35 token graph dense FFN with clamp unsupported".into());
26550                    }
26551                    let (wg_d, wu_d, wd_d) = match (ffn_gate, ffn_up, ffn_down) {
26552                        (
26553                            crate::model::GpuTensor::FloatBf16 { data: wg, .. },
26554                            crate::model::GpuTensor::FloatBf16 { data: wu, .. },
26555                            crate::model::GpuTensor::FloatBf16 { data: wd, .. },
26556                        ) => (wg, wu, wd),
26557                        _ => {
26558                            return Err(
26559                                "step35 token graph dense FFN requires bf16-resident weights"
26560                                    .into(),
26561                            );
26562                        }
26563                    };
26564                    crate::tp::graph_section(e, None, || {
26565                        let _main = e.gpu.enter_main()?;
26566                        let Step35TokenGraphState {
26567                            x,
26568                            x1,
26569                            mixed_stage,
26570                            dense_z,
26571                            dense_gate,
26572                            dense_up,
26573                            dense_act,
26574                            sh_stage,
26575                            ..
26576                        } = &mut *state;
26577                        e.add_rms_norm(
26578                            x,
26579                            mixed_stage,
26580                            layer.post_attn_norm.float_data(),
26581                            x1,
26582                            dense_z,
26583                            n_embd,
26584                            1,
26585                            eps,
26586                        )?;
26587                        // TWO SINGLE matvecs, not the dual: eager dense rides two
26588                        // matmul-dispatched matvec_bf16 launches; the dual twin measured a
26589                        // ~2e-9 residual difference here (token-graph bisection, 2026-08-21).
26590                        e.matvec_bf16_into(wg_d, dense_z, dense_gate, n_embd, n_ff)?;
26591                        e.matvec_bf16_into(wu_d, dense_z, dense_up, n_embd, n_ff)?;
26592                        Self::ffn_act_lim(
26593                            e, &self.cfg, dense_gate, dense_up, 1.0, 1.0, lim, dense_act, n_ff,
26594                        )?;
26595                        e.matvec_bf16_into(wd_d, dense_act, sh_stage, n_ff, n_embd)?;
26596                        e.add(x1, sh_stage, x, n_embd)?;
26597                        Ok(())
26598                    })?;
26599                }
26600                crate::hybrid::Ffn::Moe(m) => {
26601                    let moe = self
26602                        .cfg
26603                        .moe
26604                        .as_ref()
26605                        .ok_or("step35 token graph needs moe cfg")?;
26606                    let n_expert = moe.expert_count as usize;
26607                    let n_used = moe.expert_used_count as usize;
26608                    let sigmoid = self
26609                        .cfg
26610                        .sigmoid_router()
26611                        .ok_or("step35 token graph needs the sigmoid router")?;
26612                    let step_tp = m
26613                        .step_tp
26614                        .as_ref()
26615                        .ok_or("step35 token graph needs TP experts")?;
26616                    let bank = match &step_tp.experts {
26617                        crate::hybrid::StepTpExpertBank::Nvfp4(bank) => bank,
26618                        _ => return Err("step35 token graph needs the NVFP4 bank".into()),
26619                    };
26620                    let routes_ws_mutex = bank.device_workspace_handle();
26621                    let mut routes_guard = routes_ws_mutex
26622                        .lock()
26623                        .map_err(|_| "routes workspace lock is poisoned")?;
26624                    let routes_ws = routes_guard
26625                        .as_mut()
26626                        .ok_or("step35 token graph requires the routes workspace warmed")?;
26627                    routes_ws.arm_stages(e, bank.input_width, n_used)?;
26628                    step_tp.runtime.routes_arm_raw(bank, routes_ws)?;
26629                    let p_z = {
26630                        let root = step_tp
26631                            .runtime
26632                            .rank_engine(0)
26633                            .ok_or("routes root engine missing")?;
26634                        let _main = root.gpu.enter_main()?;
26635                        let stream = root.stream();
26636                        let in_stage = routes_ws
26637                            .in_stage_handle()
26638                            .ok_or("routes in stage not armed")?;
26639                        let (a, _g) = in_stage.device_ptr(&stream);
26640                        a
26641                    };
26642                    let local_out = bank.expert_width / ranks;
26643
26644                    // ---- E2: post-attn norm into the routes in-stage + router + staging ----
26645                    crate::tp::graph_section(e, None, || {
26646                        let _main = e.gpu.enter_main()?;
26647                        {
26648                            let in_stage = routes_ws
26649                                .in_stage_mut()
26650                                .ok_or("routes in stage not armed")?;
26651                            let Step35TokenGraphState {
26652                                x, x1, mixed_stage, ..
26653                            } = &mut *state;
26654                            e.add_rms_norm(
26655                                x,
26656                                mixed_stage,
26657                                layer.post_attn_norm.float_data(),
26658                                x1,
26659                                in_stage,
26660                                n_embd,
26661                                1,
26662                                eps,
26663                            )?;
26664                        }
26665                        {
26666                            let z_ref = routes_ws
26667                                .in_stage_handle()
26668                                .ok_or("routes in stage not armed")?;
26669                            e.router_gemv_into(
26670                                m.gate_inp.float_data(),
26671                                z_ref,
26672                                &mut state.router_logits,
26673                                n_embd,
26674                                n_expert,
26675                                1,
26676                            )?;
26677                        }
26678                        let (sel_e, w_e) = routes_ws
26679                            .dev_route_e_mut()
26680                            .ok_or("routes staging not armed")?;
26681                        e.moe_router_sigmoid_topk_into(
26682                            &state.router_logits,
26683                            1,
26684                            n_expert,
26685                            n_used,
26686                            m.active_count(),
26687                            &m.exp_probs_b_dev,
26688                            &m.active_experts_dev,
26689                            sigmoid.0,
26690                            sigmoid.1,
26691                            sel_e,
26692                            w_e,
26693                        )?;
26694                        Ok(())
26695                    })?;
26696
26697                    // ---- R0r/R1r (parallel): routes sweeps ----
26698                    group_id += 1;
26699                    for rank in 0..ranks {
26700                        let engine = step_tp
26701                            .runtime
26702                            .rank_engine(rank)
26703                            .ok_or("routes rank engine missing")?;
26704                        let runtime = &step_tp.runtime;
26705                        crate::tp::graph_section(engine, Some(group_id), || {
26706                            runtime.routes_rank_section(
26707                                bank,
26708                                routes_ws,
26709                                p_z,
26710                                local_out,
26711                                n_used,
26712                                step_tp.activation_limit,
26713                                rank,
26714                            )
26715                        })?;
26716                    }
26717
26718                    // ---- ROOTr: combine into the out stage ----
26719                    {
26720                        let root = step_tp
26721                            .runtime
26722                            .rank_engine(0)
26723                            .ok_or("routes root engine missing")?;
26724                        let runtime = &step_tp.runtime;
26725                        crate::tp::graph_section(root, None, || {
26726                            runtime.routes_root_section(bank, routes_ws)
26727                        })?;
26728                    }
26729
26730                    // ---- E3: shexp + add_shared onto the out stage + residual ----
26731                    // Alloc-free inline of moe_ffn_grouped_add_shared's bf16_dual arm (the arm
26732                    // eager takes under MEMRA_BF16_MMV, guarded at fn entry).
26733                    let lim_sh = self.cfg.clamp_shexp_at(il as u32);
26734                    let (wg_sh, wu_sh, wd_sh) = match (&m.gate_shexp, &m.up_shexp, &m.down_shexp) {
26735                        (
26736                            Some(crate::model::GpuTensor::FloatBf16 { data: wg, .. }),
26737                            Some(crate::model::GpuTensor::FloatBf16 { data: wu, .. }),
26738                            Some(crate::model::GpuTensor::FloatBf16 { data: wd, .. }),
26739                        ) => (wg, wu, wd),
26740                        _ => {
26741                            return Err(
26742                                "step35 token graph shexp requires bf16-resident weights".into()
26743                            );
26744                        }
26745                    };
26746                    let n_ff_sh = m
26747                        .gate_shexp
26748                        .as_ref()
26749                        .expect("matched Some above")
26750                        .out_features();
26751                    // No sigmoid gate on this model's shexp -> gate_sig stays at its 1.0
26752                    // init, reproducing eager's ones vector without a launch.
26753                    let gate_inp_shexp = m.gate_inp_shexp.as_ref();
26754                    crate::tp::graph_section(e, None, || {
26755                        let _main = e.gpu.enter_main()?;
26756                        let (z_ref, out_stage) = routes_ws
26757                            .in_and_out_stages_mut()
26758                            .ok_or("routes stages not armed")?;
26759                        let Step35TokenGraphState {
26760                            x,
26761                            x1,
26762                            sh_stage,
26763                            shexp_gate,
26764                            shexp_up,
26765                            shexp_act,
26766                            gate_sig,
26767                            ..
26768                        } = &mut *state;
26769                        e.matvec_bf16_dual_into(
26770                            wg_sh, wu_sh, z_ref, shexp_gate, shexp_up, n_embd, n_ff_sh,
26771                        )?;
26772                        Self::ffn_act_lim(
26773                            e, &self.cfg, shexp_gate, shexp_up, 1.0, 1.0, lim_sh, shexp_act,
26774                            n_ff_sh,
26775                        )?;
26776                        e.matvec_bf16_into(wd_sh, shexp_act, sh_stage, n_ff_sh, n_embd)?;
26777                        if let Some(gate_w) = gate_inp_shexp {
26778                            e.sigmoid_dot_rows_into(
26779                                z_ref,
26780                                gate_w.float_data(),
26781                                gate_sig,
26782                                n_embd,
26783                                1,
26784                            )?;
26785                        }
26786                        e.add_scaled_rows(sh_stage, gate_sig, out_stage, n_embd, 1)?;
26787                        e.add(x1, out_stage, x, n_embd)?;
26788                        Ok(())
26789                    })?;
26790                }
26791            }
26792            if probe_layer == Some(il) {
26793                let Step35TokenGraphState { x, probe_x, .. } = &mut *state;
26794                crate::tp::graph_section(e, None, || {
26795                    let _main = e.gpu.enter_main()?;
26796                    let mut dst = probe_x.slice_mut(0..n_embd);
26797                    e.stream().memcpy_dtod(&x.slice(0..n_embd), &mut dst)?;
26798                    Ok(())
26799                })?;
26800            }
26801        }
26802
26803        // ---- Tail: output norm + head into the logits stage ----
26804        let head = match &self.output {
26805            crate::model::GpuTensor::FloatBf16 { data, .. } => data,
26806            _ => return Err("step35 token graph head requires the bf16-resident output".into()),
26807        };
26808        crate::tp::graph_section(e, None, || {
26809            let _main = e.gpu.enter_main()?;
26810            let Step35TokenGraphState {
26811                x,
26812                hn,
26813                logits_stage,
26814                token_d,
26815                pos_d,
26816                token_hist,
26817                hist_idx,
26818                ..
26819            } = &mut *state;
26820            e.rms_norm(x, self.output_norm.float_data(), hn, n_embd, 1, eps)?;
26821            e.matvec_bf16_into(head, hn, logits_stage, n_embd, self.cfg.n_vocab as usize)?;
26822            // Chunk-loop tail: greedy argmax feeds token_d (host-identical tie-break,
26823            // argmax_gate-validated), the id lands in the history ring, and pos advances on
26824            // device — consecutive launches chain with NO host sync. Single-token mode
26825            // overwrites token_d/pos_d from the host before each launch, so these nodes are
26826            // harmless there.
26827            e.argmax_token_device_into(logits_stage, token_d, self.cfg.n_vocab as usize)?;
26828            e.u32_hist_append(token_d, token_hist, hist_idx)?;
26829            e.inc_i32(pos_d)?;
26830            Ok(())
26831        })?;
26832
26833        let graph = crate::tp::token_graph_build_finish()?;
26834        state.graphs.push((bucket_max, graph));
26835        eprintln!(
26836            "[step35-token-graph] built bucket={bucket_max} layers={n_layers} \
26837             build_ms={:.0} performance_claim=false",
26838            started.elapsed().as_secs_f64() * 1e3
26839        );
26840        Ok(())
26841    }
26842}