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;
9
10/// Resident trunk transients for the eager prime (piecewise-graph foundation; see
11/// HybridModel::prime_slabs). Every buffer is fully overwritten before use per prime.
12pub struct PrimeSlabs {
13    pub t_cap: usize,
14    pub h: CudaSlice<f32>,
15    pub x1: CudaSlice<f32>,
16    pub z: CudaSlice<f32>,
17    pub act: CudaSlice<f32>,
18    pub xa: CudaSlice<f32>,
19    pub xb: CudaSlice<f32>,
20    pub h16: CudaSlice<u8>,
21    pub z16: CudaSlice<u8>,
22    /// piecewise boundary slabs (increment 2): GEMM outputs land here so the
23    /// downstream captured segments see fixed addresses.
24    pub gate: CudaSlice<f32>, // t * n_ff_max
25    pub up: CudaSlice<f32>,      // t * n_ff_max
26    pub ffn_out: CudaSlice<f32>, // t * n_embd
27    /// piecewise increment 3: per-layer S-glue segment graphs (down-add + next
28    /// attn-norm, ALL-slab IO, zero in-graph allocations -> keeperless capture is
29    /// clean). Baked at this t_cap; replay only when t == t_cap. seg_glue[il] fires
30    /// between layer il and il+1 (ping-pong parity is deterministic per il).
31    pub seg_glue: Vec<Option<cudarc::driver::CudaGraph>>,
32    /// increment 5 (core-split edition): the mixer out-GEMM writes _into_ `mixed`
33    /// directly (no staging copy — the increment-4 copy route was refuted), making
34    /// S-mid [add + post-norm] all-slab and capturable.
35    pub mixed: CudaSlice<f32>,
36    pub seg_mid: Vec<Option<cudarc::driver::CudaGraph>>,
37    pub seg_t: usize,
38}
39
40// Split prime ranges cannot enter the full-range segment-graph arm, and every slab access
41// is serialized by its device mutex after binding that device's CUDA context on the thread.
42unsafe impl Send for PrimeSlabs {}
43
44fn empty_cache_layers<T>(n: usize) -> Vec<Option<T>> {
45    std::iter::repeat_with(|| None).take(n).collect()
46}
47
48/// Temporarily move a PP-2 cache's layer state into two independently-owned cache shells.
49/// The stage walkers then receive disjoint `&mut Cache` values and can run on separate host
50/// threads without aliasing. GPU buffers are moved, not copied; Drop restores every layer
51/// and publishes the last position completed by both stages.
52struct PrimeCacheStages<'a> {
53    parent: &'a mut Cache,
54    cut: usize,
55    stage0: Cache,
56    stage1: Cache,
57}
58
59impl<'a> PrimeCacheStages<'a> {
60    fn new(parent: &'a mut Cache, cut: usize) -> Self {
61        let n = parent.kv.len();
62        assert_eq!(parent.recur.len(), n, "cache layer vectors disagree");
63        assert!(cut <= n, "PP-2 cache cut {cut} exceeds {n} layers");
64        let mut kv0 = empty_cache_layers(n);
65        let mut kv1 = empty_cache_layers(n);
66        let mut recur0 = empty_cache_layers(n);
67        let mut recur1 = empty_cache_layers(n);
68        for i in 0..cut {
69            kv0[i] = parent.kv[i].take();
70            recur0[i] = parent.recur[i].take();
71        }
72        for i in cut..n {
73            kv1[i] = parent.kv[i].take();
74            recur1[i] = parent.recur[i].take();
75        }
76        let pos = parent.pos;
77        let max_ctx = parent.max_ctx;
78        Self {
79            parent,
80            cut,
81            stage0: Cache {
82                kv: kv0,
83                recur: recur0,
84                pos,
85                max_ctx,
86                last_logits_dev: None,
87                dflash_taps: None,
88            },
89            stage1: Cache {
90                kv: kv1,
91                recur: recur1,
92                pos,
93                max_ctx,
94                last_logits_dev: None,
95                dflash_taps: None,
96            },
97        }
98    }
99
100    fn parts(&mut self) -> (&mut Cache, &mut Cache) {
101        (&mut self.stage0, &mut self.stage1)
102    }
103}
104
105impl Drop for PrimeCacheStages<'_> {
106    fn drop(&mut self) {
107        let n = self.parent.kv.len();
108        for i in 0..n {
109            let source = if i < self.cut {
110                &mut self.stage0
111            } else {
112                &mut self.stage1
113            };
114            debug_assert!(self.parent.kv[i].is_none());
115            debug_assert!(self.parent.recur[i].is_none());
116            self.parent.kv[i] = source.kv[i].take();
117            self.parent.recur[i] = source.recur[i].take();
118        }
119        self.parent.pos = self.stage0.pos.min(self.stage1.pos);
120    }
121}
122
123/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
124pub(crate) struct AttnPre {
125    pub q: cudarc::driver::CudaSlice<f32>,
126    pub k: cudarc::driver::CudaSlice<f32>,
127    pub v: cudarc::driver::CudaSlice<f32>,
128    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
129}
130
131/// task #18: one sequence's GDN prep outputs (the scan inputs).
132pub(crate) struct GdnPrep {
133    pub hk: usize,
134    pub q_l2: cudarc::driver::CudaSlice<f32>,
135    pub k_l2: cudarc::driver::CudaSlice<f32>,
136    pub v_g: cudarc::driver::CudaSlice<f32>,
137    pub beta: cudarc::driver::CudaSlice<f32>,
138    pub g_log: cudarc::driver::CudaSlice<f32>,
139    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
140    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
141}
142
143/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
144pub(crate) struct VerifyStreamScratch {
145    pub pos_d: CudaSlice<i32>,
146    pub row_ctrs: Vec<CudaSlice<i32>>,
147}
148use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MoeWeights};
149
150struct MoeInputTraceWriter {
151    dir: std::path::PathBuf,
152    index: std::fs::File,
153    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
154}
155
156static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<std::sync::Mutex<Option<MoeInputTraceWriter>>> =
157    std::sync::OnceLock::new();
158
159/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
160/// per-expert launch chain). See `moe_gdec_token`.
161fn gdec_enabled() -> bool {
162    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *E.get_or_init(|| {
164        std::env::var("MEMRA_MOE_GDEC")
165            .map(|v| v != "0")
166            .unwrap_or(true)
167    })
168}
169
170/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
171/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
172/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
173/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
174/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
175/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
176/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
177/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
178/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
179/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
180fn moe_slab_enabled() -> bool {
181    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
182}
183
184/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
185/// default flip. `=0` selects the established path, while any other explicit value enables the
186/// grouped research arm for the current call.
187fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
188    std::env::var("MEMRA_MOE_GROUPED")
189        .map(|value| value != "0")
190        .unwrap_or(false)
191}
192
193/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
194/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
195fn moe_prefetch_enabled() -> bool {
196    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
197    *E.get_or_init(|| {
198        std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
199            || crate::spill_pread::worker_enabled()
200    })
201}
202
203/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
204/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
205/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
206/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
207fn moe_page_prefetch_window() -> usize {
208    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
209    *W.get_or_init(|| {
210        page_prefetch_window_from_values(
211            std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
212            std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW")
213                .ok()
214                .as_deref(),
215        )
216    })
217}
218
219fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
220    if !enabled {
221        return 0;
222    }
223    raw_window.and_then(|value| value.parse().ok()).unwrap_or(1)
224}
225
226/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
227/// full window; each later position adds one expert at the far edge. Thus widening the window does
228/// not repeatedly issue `MADV_WILLNEED` for the same range.
229fn page_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
230    if window == 0 || position >= len {
231        return len..len;
232    }
233    let (start, count) = if position == 0 {
234        (1, window)
235    } else {
236        (position.saturating_add(window), 1)
237    };
238    let start = start.min(len);
239    start..start.saturating_add(count).min(len)
240}
241
242/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
243/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
244fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
245    let position = current.map_or(0, |position| position.saturating_add(1));
246    (position < order_len).then_some(position)
247}
248
249/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
250/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
251/// window. Position zero primes the current expert too: its three independent reads can run in
252/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
253fn worker_prefetch_window() -> usize {
254    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
255    *WINDOW.get_or_init(|| {
256        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
257        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
258            .ok()
259            .and_then(|value| value.parse::<usize>().ok())
260            .unwrap_or(automatic.max(1))
261    })
262}
263
264/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
265/// this includes the current expert when the window is seeded so all three current projections
266/// enter the CPU pool together.
267fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
268    if window == 0 || position >= len {
269        return len..len;
270    }
271    let (start, count) = if position == 0 {
272        (0, window)
273    } else {
274        (position.saturating_add(window).saturating_sub(1), 1)
275    };
276    let start = start.min(len);
277    start..start.saturating_add(count).min(len)
278}
279
280/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
281/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
282/// expert weight pointers come from the per-layer device table. Requires the fused router (the
283/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
284fn moe_dev_enabled() -> bool {
285    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
286    *E.get_or_init(|| {
287        std::env::var("MEMRA_MOE_DEV")
288            .map(|v| v != "0")
289            .unwrap_or(true)
290            && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0"))
291    })
292}
293
294/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
295/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
296fn sigmoid_router_enabled() -> bool {
297    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
298    *E.get_or_init(|| {
299        std::env::var("MEMRA_SIG_ROUTER")
300            .map(|v| v != "0")
301            .unwrap_or(true)
302    })
303}
304
305/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
306/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
307/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
308/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
309fn moe_q8_enabled() -> bool {
310    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
311    *E.get_or_init(|| {
312        std::env::var("MEMRA_MOE_Q8")
313            .map(|v| v != "0")
314            .unwrap_or(true)
315    })
316}
317
318/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
319/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
320fn expert_dp4a_supported(qt: i32) -> bool {
321    qt == crate::QT_Q4_0
322        || qt == crate::QT_IQ3_S
323        || qt == crate::QT_IQ4_XS
324        || qt == crate::QT_Q3_K
325        || qt == crate::QT_Q4_K
326        || qt == crate::QT_Q6_K
327}
328
329fn q8_expert_supported(qt: i32) -> bool {
330    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
331    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
332    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
333    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
334    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
335    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
336    let kq = *KQ.get_or_init(|| {
337        std::env::var("MEMRA_MOE_Q8_KQ")
338            .map(|v| v != "0")
339            .unwrap_or(true)
340    });
341    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
342    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
343    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
344    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
345    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
346    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
347    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4")
348        .map(|v| v != "0")
349        .unwrap_or(true);
350    qt == crate::QT_IQ3_S
351        || qt == crate::QT_IQ4_XS
352        || (nvfp4_q8 && qt == crate::QT_NVFP4)
353        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
354}
355
356/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
357/// k-quant tensors must fall to the _em dot path instead.
358fn q8_expert_dec_supported(qt: i32) -> bool {
359    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
360}
361
362/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
363/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
364/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
365/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
366/// q35 layers, which is why that cell measured FLAT.
367fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
368    match qt {
369        crate::QT_Q4_0 => in_f % 32 == 0,
370        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K | crate::QT_Q6_K => {
371            in_f % 256 == 0
372        }
373        // NVFP4 (block 64) added lane/moebatch-q35moe 2026-08-21: the ornith15 expert bank is
374        // uniform NVFP4, which passed the pairs q8 gate but missed BOTH batched doors
375        // (use_mma's dec set and this table), so 14.7k-token prefill rode the per-pair _em
376        // fallback — 88.6% of the prime wall (prime-anatomy receipt).
377        crate::QT_NVFP4 => in_f % 64 == 0,
378        _ => false,
379    }
380}
381
382/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
383/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
384fn moe_prewarm_enabled() -> bool {
385    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386    *E.get_or_init(|| {
387        std::env::var("MEMRA_MOE_PREWARM")
388            .map(|v| v != "0")
389            .unwrap_or(true)
390    })
391}
392
393/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
394/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
395/// can then vote for and exercise those experts on GPU before the cache is frozen.
396fn cpu_expert_profile_admit_enabled() -> bool {
397    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
398    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
399}
400
401/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
402/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
403/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
404pub const PRIME_MIN_T: usize = 16;
405const PRIME_PIPE_MICROBATCHES: usize = 8;
406const PRIME_PIPE_MIN_CHUNK: usize = 128;
407const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
408const PRIME_PIPE_LINEAR_WORK: usize = 8;
409
410fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
411    crate::pp::prime_pp_on()
412        && !crate::pp::pp2_streams_off()
413        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
414}
415
416/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative.
417/// Naked PP-2 primes use the measured pipeline geometry: up to eight microchunks, never
418/// below 128 tokens, while the legacy 4096-token cap remains the long-context bound.
419pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
420    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
421        let parsed = value
422            .parse::<usize>()
423            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
424        return if crate::cache::swa_ring_on() {
425            if parsed == 0 {
426                crate::cache::PRIME_CHUNK_MAX_TOKENS
427            } else {
428                parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
429            }
430        } else {
431            parsed
432        };
433    }
434    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
435    if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
436        chunk.min(
437            t.div_ceil(PRIME_PIPE_MICROBATCHES)
438                .max(PRIME_PIPE_MIN_CHUNK),
439        )
440    } else {
441        chunk
442    }
443}
444
445fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
446    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
447}
448
449fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
450    if chunk == 0 || t <= chunk {
451        return vec![(0, t)];
452    }
453    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
454    let mut start = 0usize;
455    while start < t {
456        let mut end = (start + chunk).min(t);
457        if t - end > 0 && t - end < PRIME_MIN_T {
458            if ring_on {
459                let shifted = t - PRIME_MIN_T;
460                end = if shifted > start { shifted } else { t };
461            } else {
462                end = t;
463            }
464        }
465        ranges.push((start, end));
466        start = end;
467    }
468    ranges
469}
470
471fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
472    let prefix = prefix as u128;
473    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
474}
475
476fn dynamic_prime_chunk_ranges(
477    t: usize,
478    fixed_chunk: usize,
479    fixed: &[(usize, usize)],
480) -> Vec<(usize, usize)> {
481    let n = fixed.len();
482    if n < 3 {
483        return fixed.to_vec();
484    }
485
486    let max_first = t - (n - 1) * PRIME_MIN_T;
487    let first = fixed_chunk
488        .div_ceil(2)
489        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
490        .min(max_first);
491    let mut ranges = Vec::with_capacity(n);
492    ranges.push((0, first));
493
494    let first_work = prime_chunk_work(first, t);
495    let work_span = prime_chunk_work(t, t) - first_work;
496    let denominator = (n - 1) as u128;
497    let mut previous = first;
498    for boundary in 1..n - 1 {
499        let target = first_work * denominator + work_span * (boundary as u128);
500        let remaining = n - 1 - boundary;
501        let mut low = previous + PRIME_MIN_T;
502        let mut high = t - remaining * PRIME_MIN_T;
503        while low < high {
504            let mid = low + (high - low) / 2;
505            if prime_chunk_work(mid, t) * denominator >= target {
506                high = mid;
507            } else {
508                low = mid + 1;
509            }
510        }
511        ranges.push((previous, low));
512        previous = low;
513    }
514    ranges.push((previous, t));
515    ranges
516}
517
518/// Internal prime ranges. The naked PP-2 pipeline defaults to a short-fill,
519/// equal-modeled-time schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
520/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
521pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
522    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
523    let chunk = prime_chunk_tokens(t, n_layers);
524    let fixed = fixed_prime_chunk_ranges(t, chunk);
525    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
526        Ok(value) => value == "dynamic",
527        Err(_) => true,
528    };
529    if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
530        fixed
531    } else {
532        dynamic_prime_chunk_ranges(t, chunk, &fixed)
533    }
534}
535
536impl HybridModel {
537    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
538    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
539    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
540    /// (it forces a dtoh + host hash per layer).
541    fn prime_trace_path() -> Option<&'static str> {
542        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
543        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
544            .as_deref()
545    }
546
547    /// PRIME ANATOMY (diagnostic): `MEMRA_PRIME_ANATOMY=1` synchronizes the stream around
548    /// each prime_layers stage and accumulates wall time per stage class, printed after
549    /// every prime_layers call (cumulative across chunks/reps). The per-stage syncs
550    /// serialize launch/execute overlap, so the summed total exceeds the naked prime wall —
551    /// attribution ratios only, never a measured default run. Non-seg serial arm only.
552    fn prime_anatomy_on() -> bool {
553        static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
554        *E.get_or_init(|| std::env::var("MEMRA_PRIME_ANATOMY").as_deref() == Ok("1"))
555    }
556
557    fn prime_anatomy_slots() -> &'static [std::sync::atomic::AtomicU64; 5] {
558        static S: [std::sync::atomic::AtomicU64; 5] = [
559            std::sync::atomic::AtomicU64::new(0), // 0 mixer full-attn
560            std::sync::atomic::AtomicU64::new(0), // 1 mixer linear-attn (GDN)
561            std::sync::atomic::AtomicU64::new(0), // 2 ffn MoE (router + experts + shexp)
562            std::sync::atomic::AtomicU64::new(0), // 3 ffn dense
563            std::sync::atomic::AtomicU64::new(0), // 4 norms/adds/glue
564        ];
565        &S
566    }
567
568    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
569    pub fn forward(
570        &self,
571        e: &Engine,
572        tokens: &[u32],
573    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
574        if self.is_gemma4_e4b() {
575            return self.gemma4_e4b_forward(e, tokens, false);
576        }
577        if self.cfg.gemma4.is_some() {
578            return self.gemma4_forward(e, tokens, false);
579        }
580        let cfg = &self.cfg;
581        let n_embd = cfg.n_embd as usize;
582        let t = tokens.len();
583        let eps = cfg.rms_eps;
584        let pos: Vec<i32> = (0..t as i32).collect();
585        let pos_d = e.htod_i32(&pos)?;
586
587        let mut x = self.embed(e, tokens)?; // [T, n_embd]
588
589        for (il, layer) in self.layers.iter().enumerate() {
590            // attn_norm
591            let mut h = e.uninit(t * n_embd)?;
592            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
593
594            let mixed = match &layer.mixer {
595                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
596                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
597                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
598            };
599
600            // residual 1
601            let mut x1 = e.uninit(t * n_embd)?;
602            e.add(&x, &mixed, &mut x1, t * n_embd)?;
603
604            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
605            let mut z = e.uninit(t * n_embd)?;
606            e.rms_norm(
607                &x1,
608                layer.post_attn_norm.float_data(),
609                &mut z,
610                n_embd,
611                t,
612                eps,
613            )?;
614            let ffn_out = match &layer.ffn {
615                crate::hybrid::Ffn::Dense {
616                    ffn_gate,
617                    ffn_up,
618                    ffn_down,
619                } => {
620                    let n_ff = ffn_gate.out_features();
621                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
622                    let up = g2.pop().unwrap();
623                    let gate = g2.pop().unwrap();
624                    let mut act = e.uninit(t * n_ff)?;
625                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
626                    // both the dense MLP and the shared expert, and its limit is
627                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
628                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
629                    Self::ffn_act_lim(
630                        e,
631                        &self.cfg,
632                        &gate,
633                        &up,
634                        1.0,
635                        1.0,
636                        self.cfg.clamp_shexp_at(il as u32),
637                        &mut act,
638                        t * n_ff,
639                    )?;
640                    e.matmul(ffn_down, &act, t)?
641                }
642                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
643            };
644            let mut x2 = e.uninit(t * n_embd)?;
645            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
646            x = x2;
647        }
648
649        let mut hn = e.uninit(t * n_embd)?;
650        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
651        let logits = e.matmul(&self.output, &hn, t)?;
652        Ok(e.dtoh(&logits)?)
653    }
654
655    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
656    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
657    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
658    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
659    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
660    pub fn forward_last(
661        &self,
662        e: &Engine,
663        tokens: &[u32],
664    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
665        if self.cfg.gemma4.is_some() {
666            return self.gemma4_forward(e, tokens, true);
667        }
668        let cfg = &self.cfg;
669        let n_embd = cfg.n_embd as usize;
670        let t = tokens.len();
671        let eps = cfg.rms_eps;
672        let pos: Vec<i32> = (0..t as i32).collect();
673        let pos_d = e.htod_i32(&pos)?;
674
675        let mut x = self.embed(e, tokens)?; // [T, n_embd]
676        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
677        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
678        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
679        let anat = Self::prime_anatomy_on();
680        let mut anat_last = if anat {
681            e.stream().synchronize()?;
682            Some(std::time::Instant::now())
683        } else {
684            None
685        };
686        macro_rules! anat_mark {
687            ($slot:expr) => {
688                if let Some(ts) = anat_last.as_mut() {
689                    e.stream().synchronize()?;
690                    Self::prime_anatomy_slots()[$slot].fetch_add(
691                        ts.elapsed().as_nanos() as u64,
692                        std::sync::atomic::Ordering::Relaxed,
693                    );
694                    *ts = std::time::Instant::now();
695                }
696            };
697        }
698        for (il, layer) in self.layers.iter().enumerate() {
699            let mut h = e.uninit(t * n_embd)?;
700            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
701            if probe {
702                e.stream().synchronize()?;
703                eprintln!("[probe] L{il} norm ok");
704            }
705            anat_mark!(4);
706            let mixed = match &layer.mixer {
707                Mixer::Full(fa) => {
708                    let y = self.full_attn(e, fa, &h, &pos_d, t, il)?;
709                    anat_mark!(0);
710                    y
711                }
712                Mixer::Linear(la) => {
713                    let y = self.linear_attn(e, la, &h, t)?;
714                    anat_mark!(1);
715                    y
716                }
717                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
718            };
719            if probe {
720                e.stream().synchronize()?;
721                eprintln!("[probe] L{il} mixer ok");
722            }
723            let mut x1 = e.uninit(t * n_embd)?;
724            e.add(&x, &mixed, &mut x1, t * n_embd)?;
725            let mut z = e.uninit(t * n_embd)?;
726            e.rms_norm(
727                &x1,
728                layer.post_attn_norm.float_data(),
729                &mut z,
730                n_embd,
731                t,
732                eps,
733            )?;
734            anat_mark!(4);
735            let ffn_out = match &layer.ffn {
736                crate::hybrid::Ffn::Dense {
737                    ffn_gate,
738                    ffn_up,
739                    ffn_down,
740                } => {
741                    let n_ff = ffn_gate.out_features();
742                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
743                    let up = g2.pop().unwrap();
744                    let gate = g2.pop().unwrap();
745                    let mut act = e.uninit(t * n_ff)?;
746                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
747                    Self::ffn_act_lim(
748                        e,
749                        &self.cfg,
750                        &gate,
751                        &up,
752                        1.0,
753                        1.0,
754                        self.cfg.clamp_shexp_at(il as u32),
755                        &mut act,
756                        t * n_ff,
757                    )?;
758                    let y = e.matmul(ffn_down, &act, t)?;
759                    anat_mark!(3);
760                    y
761                }
762                crate::hybrid::Ffn::Moe(m) => {
763                    let y = self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?;
764                    anat_mark!(2);
765                    y
766                }
767            };
768            if probe {
769                e.stream().synchronize()?;
770                eprintln!("[probe] L{il} ffn ok");
771            }
772            let mut x2 = e.uninit(t * n_embd)?;
773            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
774            x = x2;
775        }
776        if anat {
777            let s = Self::prime_anatomy_slots();
778            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
779            eprintln!(
780                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
781                 dense={:.1} norms_adds={:.1} (t={t}, forward_last)",
782                ms(0),
783                ms(1),
784                ms(2),
785                ms(3),
786                ms(4)
787            );
788        }
789        // norm over all T, then slice the LAST row and run lm_head on that single row.
790        let mut hn = e.uninit(t * n_embd)?;
791        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
792        let last = e.view(&hn, t * n_embd); // [T, n_embd]
793        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
794        let mut hlast = e.uninit(n_embd)?;
795        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
796        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
797        Ok(e.dtoh(&logits)?)
798    }
799
800    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
801    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
802    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
803    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
804    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
805    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
806    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
807    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
808    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
809    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
810    ///       argmax gate is the accuracy authority, exactly as for forward_last);
811    ///   (c) `cache.pos`/KV len/len_d advance by T.
812    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
813    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
814    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
815    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
816    ///
817    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
818    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
819    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
820    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
821    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
822    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
823    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
824    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
825    /// differently under load — research/tick-seg-20260807, receipt in
826    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
827    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
828    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
829    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
830    /// caller that SPLITS one request across calls passes the remainder.
831    pub fn prime_cache(
832        &self,
833        e: &Engine,
834        tokens: &[u32],
835        cache: &mut Cache,
836        queued_after: usize,
837    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
838        self.prime_cache_overlaid(e, tokens, cache, queued_after, None)
839    }
840
841    /// `prime_cache` with a vision embedding overlay (lane/vision): image merger outputs
842    /// replace the `<|image_pad|>` token embeddings at prompt-relative positions before the
843    /// trunk walk — the mixed-embedding prime. Text-only callers use `prime_cache` (overlay
844    /// None, byte-identical path). v1 scope: the serial chunk walk only — PP prime arms and
845    /// gemma4 refuse loudly (the vision serving box is single-GPU).
846    pub fn prime_cache_overlaid(
847        &self,
848        e: &Engine,
849        tokens: &[u32],
850        cache: &mut Cache,
851        queued_after: usize,
852        overlay: Option<&crate::vision::EmbedOverlay>,
853    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
854        let n_embd = self.cfg.n_embd as usize;
855        let t = tokens.len();
856        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
857        // session cache — every chunk (including the first) takes the continuation arm
858        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
859        assert!(
860            t >= PRIME_MIN_T,
861            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
862        );
863        assert!(
864            cache.pos + t <= cache.max_ctx,
865            "prime_cache: prompt exceeds cache max_ctx"
866        );
867
868        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
869        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
870        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
871        // each chunk runs the full layer stack with transients sized to the chunk, appending its
872        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
873        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
874        // exactly the state carry it was built for). Full-attn chunks after the first attend to
875        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
876        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
877        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
878        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
879        if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
880            if self.is_gemma4_e4b() {
881                if overlay.is_some() {
882                    return Err(
883                        "vision embedding overlay is unsupported on gemma4 E4B (PLE prime)".into(),
884                    );
885                }
886                return self.gemma4_e4b_prime(e, tokens, cache);
887            }
888            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
889            // An overlay takes the masked-prefill arm: image rows splice in unscaled
890            // (gemma4.cpp:182 — embd batches skip the sqrt(n_embd) scale) and the image
891            // spans become bidirectional attention islands (lane/gemma-vision).
892            return self.gemma4_prime(e, tokens, cache, overlay);
893        }
894        let ranges = prime_chunk_ranges(t, self.layers.len());
895        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
896        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
897        // the prefill's ARITHMETIC, so two rigs with different values produced different
898        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
899        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
900        // (VERDICT.md) — and it is NOT what docs originally said:
901        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
902        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
903        //     output head), so growing a chunk cannot move an existing row's value.
904        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
905        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
906        //     not describe our leak.
907        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
908        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
909        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
910        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
911        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
912        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
913        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
914        // the source — every row is in one numeric class, so the chunk size no longer steers
915        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
916        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
917        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
918        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
919        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
920        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
921        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
922        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
923        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
924        // across calls, the request still ends at the same absolute position, whatever the tick
925        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
926        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
927        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
928        // default. Read per call, not cached (the probe flips it in-process between arms). Never
929        // on in a measured default run.
930        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
931        let seq_end = if legacy_calllocal {
932            cache.pos + t
933        } else {
934            cache.pos + t + queued_after
935        };
936        if ranges.len() == 1 {
937            return self.prime_chunk(e, tokens, cache, seq_end, 0, overlay);
938        }
939        // PIPELINED PP-2 PRIME (lane/cx-pipeline-prime, 2026-08-08): overlap stage 0 of
940        // chunk N+1 with stage 1 of chunk N. The serial split stays reachable through
941        // MEMRA_PRIME_PIPE=0 and is the exactness oracle. N>2 keeps the serial walker;
942        // this lane owns the balanced two-stage schedule only.
943        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
944            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
945                if overlay.is_some() {
946                    return Err(
947                        "vision embedding overlay + pipelined PP prime unsupported (v1); \
948                         run the serial prime (single device or MEMRA_PRIME_PIPE=0)"
949                            .into(),
950                    );
951                }
952                if crate::pp::pp_multi_stream_same_device() {
953                    return Err(
954                        "prime chunk pipeline refused with 2 stage streams on one device — \
955                         that concurrent-stream placement remains quarantined by the deferred \
956                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
957                         the serial split."
958                            .into(),
959                    );
960                }
961                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
962            }
963        }
964        let mut hiddens = e.uninit(t * n_embd)?;
965        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
966        for &(start, end) in &ranges {
967            // chunked prime writes tap rows at the chunk's absolute offset
968            if let Some(taps) = cache.dflash_taps.as_mut() {
969                taps.base = start;
970            }
971            let (l, hs, x) =
972                self.prime_chunk(e, &tokens[start..end], cache, seq_end, start, overlay)?;
973            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
974            last = Some((l, hs));
975        }
976        let (logits, h_seed) = last.unwrap();
977        Ok((logits, h_seed, hiddens))
978    }
979
980    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
981    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
982    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
983    /// norm, lm head, and caller hidden-stack copy as the serial split.
984    fn prime_cache_pp2_pipelined(
985        &self,
986        e: &Engine,
987        tokens: &[u32],
988        cache: &mut Cache,
989        seq_end: usize,
990        ranges: &[(usize, usize)],
991        fence: &[usize],
992    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
993        debug_assert_eq!(fence.len(), 3);
994        debug_assert!(ranges.len() >= 2);
995        let rt = crate::pp::PpNRt::get(e)?;
996        assert_eq!(
997            rt.n_stages(),
998            2,
999            "prime pipeline requires exactly two PP stages"
1000        );
1001        let n_embd = self.cfg.n_embd as usize;
1002        let t = tokens.len();
1003        let initial_base = cache.pos;
1004        let caller_stream = e.stream();
1005
1006        // #87 reverse publication before any new stage allocation, then prewarm both
1007        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
1008        // after stage 1(N) is queued would synchronize that stream and erase the first
1009        // overlap on a two-chunk prompt.
1010        rt.fence_stages_behind(&caller_stream)?;
1011        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
1012        rt.prepare_overlap_slots(0, max_payload)?;
1013
1014        let mut hiddens = e.uninit(t * n_embd)?;
1015        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
1016        let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
1017        let (cache0, cache1) = stage_caches.parts();
1018        let (first_start, first_end) = ranges[0];
1019        let mut slot = self.prime_pp2_stage0_enqueue(
1020            e,
1021            rt,
1022            &tokens[first_start..first_end],
1023            cache0,
1024            seq_end,
1025            fence,
1026            initial_base + first_start,
1027            true,
1028        )?;
1029        cache0.pos = initial_base + first_end;
1030
1031        for (i, &(start, end)) in ranges.iter().enumerate() {
1032            let base = initial_base + start;
1033            debug_assert_eq!(
1034                cache1.pos, base,
1035                "stage 1 must drain chunks in original position order"
1036            );
1037            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
1038                let next_base = initial_base + next_start;
1039                debug_assert_eq!(
1040                    cache0.pos, next_base,
1041                    "stage 0 must issue chunks in original position order"
1042                );
1043                let cache0_stage = &mut *cache0;
1044                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
1045                // on one host thread therefore serialize even if the calls are ordered as
1046                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
1047                // stage 1 consumes slot N while stage 0 produces slot N+1.
1048                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
1049                    let stage0 = scope.spawn(move || -> Result<usize, String> {
1050                        let next = self
1051                            .prime_pp2_stage0_enqueue(
1052                                e,
1053                                rt,
1054                                &tokens[next_start..next_end],
1055                                cache0_stage,
1056                                seq_end,
1057                                fence,
1058                                next_base,
1059                                true,
1060                            )
1061                            .map_err(|err| err.to_string())?;
1062                        cache0_stage.pos = initial_base + next_end;
1063                        Ok(next)
1064                    });
1065                    let x = self.prime_pp2_stage1_enqueue(
1066                        e,
1067                        rt,
1068                        slot,
1069                        end - start,
1070                        cache1,
1071                        seq_end,
1072                        fence,
1073                        base,
1074                        true,
1075                    )?;
1076                    let out = {
1077                        rt.bind_stage(1)?;
1078                        let _st1 = rt.enter(1);
1079                        let e1 = rt.engine(1, e);
1080                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1081                    };
1082                    let next = stage0
1083                        .join()
1084                        .map_err(|_| "pipeprime stage-0 host walker panicked")?
1085                        .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1086                    Ok((out, Some(next)))
1087                })?
1088            } else {
1089                let x = self.prime_pp2_stage1_enqueue(
1090                    e,
1091                    rt,
1092                    slot,
1093                    end - start,
1094                    cache1,
1095                    seq_end,
1096                    fence,
1097                    base,
1098                    true,
1099                )?;
1100                let out = {
1101                    rt.bind_stage(1)?;
1102                    let _st1 = rt.enter(1);
1103                    let e1 = rt.engine(1, e);
1104                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
1105                };
1106                (out, None)
1107            };
1108
1109            rt.publish_to(1, &caller_stream)?;
1110            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1111            last = Some((out.0, out.1));
1112            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1113
1114            if let Some(next) = next_slot {
1115                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
1116                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
1117                // Stage 0(N+1) is already queued before this wait is appended, so its
1118                // overlap with stage 1(N) is preserved.
1119                rt.fence_stages_behind(&caller_stream)?;
1120                slot = next;
1121            }
1122        }
1123
1124        debug_assert_eq!(cache0.pos, initial_base + t);
1125        debug_assert_eq!(cache1.pos, initial_base + t);
1126        let (logits, h_seed) = last.unwrap();
1127        Ok((logits, h_seed, hiddens))
1128    }
1129
1130    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
1131    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
1132    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
1133    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
1134    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
1135    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
1136    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1137        if Engine::gdn_db_on()
1138            && Engine::gdn_chunked_enabled()
1139            && t >= 16
1140            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1141            && num_k * 2 == num_v
1142        {
1143            num_k
1144        } else {
1145            num_v
1146        }
1147    }
1148
1149    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
1150    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
1151    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
1152    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
1153    fn f16out_on(e: &Engine, t: usize) -> bool {
1154        crate::f16_ffi::pp_f16_enabled()
1155            && t >= 16
1156            && !e.verify_exact_on()
1157            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1158    }
1159
1160    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
1161    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
1162    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
1163    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
1164    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
1165    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
1166    /// see one entry, byte-identical behavior.
1167    pub fn prime_slabs_get(
1168        &self,
1169        e: &Engine,
1170        t: usize,
1171        n_embd: usize,
1172        n_ff_max: usize,
1173    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1174        let mut slabs = self.prime_slabs.lock().unwrap();
1175        let dev = e.ctx().ordinal();
1176        let need_new = match slabs.get(&dev) {
1177            None => true,
1178            Some(sl) => sl.lock().unwrap().t_cap < t,
1179        };
1180        if need_new {
1181            slabs.insert(
1182                dev,
1183                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1184                    t_cap: t,
1185                    h: e.uninit(t * n_embd)?,
1186                    x1: e.uninit(t * n_embd)?,
1187                    z: e.uninit(t * n_embd)?,
1188                    act: e.uninit(t * n_ff_max)?,
1189                    xa: e.uninit(t * n_embd)?,
1190                    xb: e.uninit(t * n_embd)?,
1191                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1192                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1193                    gate: e.uninit(t * n_ff_max)?,
1194                    up: e.uninit(t * n_ff_max)?,
1195                    ffn_out: e.uninit(t * n_embd)?,
1196                    seg_glue: Vec::new(),
1197                    mixed: e.uninit(t * n_embd)?,
1198                    seg_mid: Vec::new(),
1199                    seg_t: 0,
1200                })),
1201            );
1202        }
1203        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1204    }
1205
1206    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
1207    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
1208    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
1209    fn prime_chunk(
1210        &self,
1211        e: &Engine,
1212        tokens: &[u32],
1213        cache: &mut Cache,
1214        seq_end: usize,
1215        chunk_off: usize,
1216        overlay: Option<&crate::vision::EmbedOverlay>,
1217    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1218        if crate::pp::pp_host_bounce_active()
1219            && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1220        {
1221            return Err(
1222                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1223                 has no active prime stage split and would peer-read remote weights; keep \
1224                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1225                    .into(),
1226            );
1227        }
1228        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
1229        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
1230        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
1231        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
1232        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
1233        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
1234        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
1235        // loader is off and there is nothing remote to split for.
1236        if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1237            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1238                if overlay.is_some() {
1239                    return Err("vision embedding overlay + PP prime unsupported (v1); \
1240                         run single-device or MEMRA_PRIME_PP=0"
1241                        .into());
1242                }
1243                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1244            }
1245        }
1246        if crate::pp::pp_host_bounce_active() {
1247            return Err(
1248                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1249                 refusing an unsplit remote-weight walk"
1250                    .into(),
1251            );
1252        }
1253        let t = tokens.len();
1254        let base = cache.pos;
1255        debug_assert!(
1256            seq_end >= base + t,
1257            "prime_chunk: seq_end must cover this chunk"
1258        );
1259        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1260        let pos_d = e.htod_i32(&pos)?;
1261
1262        let mut x_embed = self.embed(e, tokens)?; // [T, n_embd]
1263        if let Some(ov) = overlay {
1264            // Mixed-embedding splice: image rows overwrite the pad-token embeddings that
1265            // fall inside this chunk's prompt-relative window [chunk_off, chunk_off+t).
1266            // Images larger than one prime chunk straddle boundaries, hence the clipping.
1267            let n_embd = self.cfg.n_embd as usize;
1268            for &(pos, row_off, n_rows) in &ov.spans {
1269                let lo = pos.max(chunk_off);
1270                let hi = (pos + n_rows).min(chunk_off + t);
1271                if lo < hi {
1272                    let src_row = row_off + (lo - pos);
1273                    let view = ov
1274                        .rows
1275                        .slice(src_row * n_embd..(src_row + (hi - lo)) * n_embd);
1276                    e.copy_view_into(
1277                        &mut x_embed,
1278                        (lo - chunk_off) * n_embd,
1279                        &view,
1280                        (hi - lo) * n_embd,
1281                    )?;
1282                }
1283            }
1284        }
1285        let x = self.prime_layers(
1286            e,
1287            x_embed,
1288            0,
1289            self.layers.len(),
1290            &pos_d,
1291            t,
1292            base,
1293            cache,
1294            seq_end,
1295        )?;
1296        self.prime_chunk_epilogue(e, x, t, cache)
1297    }
1298
1299    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
1300    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
1301    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
1302    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
1303    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
1304    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
1305    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
1306    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
1307    ///     the plain add (materialize) and the next stage hoists its own first norm — the
1308    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
1309    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
1310    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
1311    ///     each stage walks through its own resident transients;
1312    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
1313    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
1314    #[allow(clippy::too_many_arguments)]
1315    fn prime_layers(
1316        &self,
1317        e: &Engine,
1318        x_in: CudaSlice<f32>,
1319        lo: usize,
1320        hi: usize,
1321        pos_d: &CudaSlice<i32>,
1322        t: usize,
1323        base: usize,
1324        cache: &mut Cache,
1325        seq_end: usize,
1326    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1327        let cfg = &self.cfg;
1328        let n_embd = cfg.n_embd as usize;
1329        let eps = cfg.rms_eps;
1330        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
1331        // standalone convert launches). Only when the f16 lane serves and T reaches the
1332        // GEMM tier; bit-identical either way.
1333        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1334        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
1335        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
1336        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
1337        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
1338        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
1339        let n_ff_max = self
1340            .layers
1341            .iter()
1342            .map(|l| match &l.ffn {
1343                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1344                _ => n_embd,
1345            })
1346            .max()
1347            .unwrap_or(n_embd)
1348            .max(n_embd);
1349        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1350        let slab = if use_slabs {
1351            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1352        } else {
1353            None
1354        };
1355        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1356        let mut x_own; // fallback storage when slabs are off
1357        type SlabRefs<'a> = (
1358            &'a mut CudaSlice<f32>,
1359            &'a mut CudaSlice<f32>,
1360            &'a mut CudaSlice<f32>,
1361            &'a mut CudaSlice<f32>,
1362            &'a mut CudaSlice<u8>,
1363            &'a mut CudaSlice<u8>,
1364            &'a mut CudaSlice<f32>,
1365            &'a mut CudaSlice<f32>,
1366            &'a mut CudaSlice<f32>,
1367        );
1368        let (mut x_cur, mut x_nxt, sl): (
1369            &mut CudaSlice<f32>,
1370            &mut CudaSlice<f32>,
1371            Option<SlabRefs>,
1372        );
1373        let mut seg: Option<(
1374            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1375            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1376            &mut CudaSlice<f32>,
1377            &mut usize,
1378        )> = None;
1379        let mut x_own2;
1380        match slab_guard.as_mut() {
1381            Some(g) => {
1382                let slabs = &mut **g;
1383                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1384                let PrimeSlabs {
1385                    xa,
1386                    xb,
1387                    h,
1388                    x1,
1389                    z,
1390                    act,
1391                    h16,
1392                    z16,
1393                    gate,
1394                    up,
1395                    ffn_out,
1396                    seg_glue,
1397                    mixed,
1398                    seg_mid,
1399                    seg_t,
1400                    ..
1401                } = slabs;
1402                x_cur = xa;
1403                x_nxt = xb;
1404                seg = Some((seg_glue, seg_mid, mixed, seg_t));
1405                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1406            }
1407            None => {
1408                x_own = x_in;
1409                x_own2 = e.uninit(t * n_embd)?;
1410                x_cur = &mut x_own;
1411                x_nxt = &mut x_own2;
1412                sl = None;
1413            }
1414        }
1415        let mut alloc_h;
1416        let mut alloc_x1;
1417        let mut alloc_z;
1418        let mut alloc_act;
1419        let mut alloc_h16;
1420        let mut alloc_z16;
1421        let mut alloc_gate;
1422        let mut alloc_up;
1423        let mut alloc_fo;
1424        let (h, x1, z, act): (
1425            &mut CudaSlice<f32>,
1426            &mut CudaSlice<f32>,
1427            &mut CudaSlice<f32>,
1428            &mut CudaSlice<f32>,
1429        );
1430        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1431        let (sl_gate, sl_up, sl_fo): (
1432            &mut CudaSlice<f32>,
1433            &mut CudaSlice<f32>,
1434            &mut CudaSlice<f32>,
1435        );
1436        match sl {
1437            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1438                h = a;
1439                x1 = b;
1440                z = c;
1441                act = d;
1442                h16 = e16;
1443                z16 = f16b;
1444                sl_gate = g;
1445                sl_up = u;
1446                sl_fo = fo;
1447            }
1448            None => {
1449                alloc_h = e.uninit(t * n_embd)?;
1450                alloc_x1 = e.uninit(t * n_embd)?;
1451                alloc_z = e.uninit(t * n_embd)?;
1452                alloc_act = e.uninit(t * n_ff_max)?;
1453                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1454                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1455                alloc_gate = e.uninit(t * n_ff_max)?;
1456                alloc_up = e.uninit(t * n_ff_max)?;
1457                alloc_fo = e.uninit(t * n_embd)?;
1458                h = &mut alloc_h;
1459                x1 = &mut alloc_x1;
1460                z = &mut alloc_z;
1461                act = &mut alloc_act;
1462                h16 = &mut alloc_h16;
1463                z16 = &mut alloc_z16;
1464                sl_gate = &mut alloc_gate;
1465                sl_up = &mut alloc_up;
1466                sl_fo = &mut alloc_fo;
1467            }
1468        }
1469        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
1470        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
1471        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
1472        // first prime at this t (capture does not execute -> launch right after).
1473        let n_layers = self.layers.len();
1474        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
1475        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
1476        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
1477        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
1478        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
1479        // machinery stays (byte-identical) as their foundation.
1480        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
1481        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
1482        // step35 rides its own mixer through the normal per-layer arm below.
1483        let use_seg = f16fuse
1484            && seg.is_some()
1485            && self.cfg.step35.is_none()
1486            && lo == 0
1487            && hi == n_layers
1488            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1489        if let Some((sg, sm, _, st)) = seg.as_mut() {
1490            if **st != t {
1491                sg.clear();
1492                sg.extend((0..n_layers).map(|_| None));
1493                sm.clear();
1494                sm.extend((0..n_layers).map(|_| None));
1495                **st = t;
1496            }
1497        }
1498        {
1499            let layer_lo = &self.layers[lo];
1500            if f16fuse {
1501                e.rms_norm_f16out(
1502                    x_cur,
1503                    layer_lo.attn_norm.float_data(),
1504                    h,
1505                    h16,
1506                    n_embd,
1507                    t,
1508                    eps,
1509                )?;
1510            } else {
1511                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1512            }
1513        }
1514        let anat = Self::prime_anatomy_on();
1515        let mut anat_last = if anat {
1516            e.stream().synchronize()?;
1517            Some(std::time::Instant::now())
1518        } else {
1519            None
1520        };
1521        // Closes the region that just ENDED into `slot`, restarting the clock.
1522        macro_rules! anat_mark {
1523            ($slot:expr) => {
1524                if let Some(ts) = anat_last.as_mut() {
1525                    e.stream().synchronize()?;
1526                    Self::prime_anatomy_slots()[$slot].fetch_add(
1527                        ts.elapsed().as_nanos() as u64,
1528                        std::sync::atomic::Ordering::Relaxed,
1529                    );
1530                    *ts = std::time::Instant::now();
1531                }
1532            };
1533        }
1534        for il in lo..hi {
1535            let layer = &self.layers[il];
1536            let hx16 = if f16fuse { Some(&*h16) } else { None };
1537            if use_seg {
1538                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
1539                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
1540                let (pre, pre16, w_out) = match &layer.mixer {
1541                    Mixer::Full(fa) => {
1542                        let g3 = match hx16 {
1543                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1544                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1545                        };
1546                        let (pre, pre16) =
1547                            self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1548                        (pre, pre16, &fa.wo)
1549                    }
1550                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1551                    Mixer::Linear(la) => {
1552                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1553                        let g4 = match hx16 {
1554                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1555                            None => e.matmul_group(&ws, h, t)?,
1556                        };
1557                        let (pre, pre16) =
1558                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1559                        (pre, pre16, &la.ssm_out)
1560                    }
1561                };
1562                {
1563                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
1564                    let pre_n = pre.len() / t;
1565                    let xh_pre = match pre16 {
1566                        Some(x) => x,
1567                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
1568                    };
1569                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1570                        let y = e.matmul(w_out, &pre, t)?;
1571                        e.copy_into(mslab, 0, &y, t * n_embd)?;
1572                    }
1573                    if sm[il].is_none() {
1574                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1575                        let w_post = layer.post_attn_norm.float_data();
1576                        e.stream().synchronize()?;
1577                        e.stream()
1578                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1579                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1580                            e.add(x_cur, mslab, x1, t * n_embd)?;
1581                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1582                            Ok(())
1583                        })();
1584                        let g = e.stream().end_capture(
1585                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1586                        r?;
1587                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1588                    }
1589                    sm[il].as_ref().unwrap().launch()?;
1590                }
1591            } else {
1592                let mixed = match &layer.mixer {
1593                    Mixer::Full(fa) => {
1594                        let y =
1595                            self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?;
1596                        anat_mark!(0);
1597                        y
1598                    }
1599                    Mixer::Linear(la) => {
1600                        let y = self.linear_attn_prime(e, la, h, hx16, t, cache, il)?;
1601                        anat_mark!(1);
1602                        y
1603                    }
1604                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1605                };
1606                if f16fuse {
1607                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
1608                    // bit-identical) — the standalone add pass disappears.
1609                    e.add_rms_norm_f16out(
1610                        x_cur,
1611                        &mixed,
1612                        layer.post_attn_norm.float_data(),
1613                        x1,
1614                        z,
1615                        z16,
1616                        n_embd,
1617                        t,
1618                        eps,
1619                    )?;
1620                } else {
1621                    e.add(x_cur, &mixed, x1, t * n_embd)?;
1622                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1623                }
1624                anat_mark!(4);
1625            }
1626            let zx16 = if f16fuse { Some(&*z16) } else { None };
1627            match &layer.ffn {
1628                crate::hybrid::Ffn::Dense {
1629                    ffn_gate,
1630                    ffn_up,
1631                    ffn_down,
1632                } => {
1633                    let n_ff = ffn_gate.out_features();
1634                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
1635                    // the allocating group + copy when a mirror is missing.
1636                    let mut into_ok = false;
1637                    if let Some(xh) = zx16 {
1638                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1639                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1640                    }
1641                    if !into_ok {
1642                        let mut g2 = match zx16 {
1643                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1644                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1645                        };
1646                        let up_y = g2.pop().unwrap();
1647                        let gate_y = g2.pop().unwrap();
1648                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1649                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1650                    }
1651                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
1652                    // operand in-epilogue; non-silu activations keep the standalone convert.
1653                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
1654                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
1655                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1656                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1657                    {
1658                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1659                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1660                        Some(a16)
1661                    } else {
1662                        Self::ffn_act_lim(
1663                            e,
1664                            &self.cfg,
1665                            sl_gate,
1666                            sl_up,
1667                            1.0,
1668                            1.0,
1669                            d_lim,
1670                            act,
1671                            t * n_ff,
1672                        )?;
1673                        None
1674                    };
1675                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
1676                    let xh_act = match act16 {
1677                        Some(x) => x,
1678                        None => e.f16_act(act, t * n_ff, n_ff)?,
1679                    };
1680                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1681                        let y = e.matmul(ffn_down, &*act, t)?;
1682                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1683                    }
1684                }
1685                crate::hybrid::Ffn::Moe(m) => {
1686                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1687                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1688                    anat_mark!(2);
1689                }
1690            }
1691            if let (crate::hybrid::Ffn::Dense { .. }, true) = (&layer.ffn, anat) {
1692                anat_mark!(3);
1693            }
1694            if use_seg && il + 1 < hi {
1695                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
1696                let w_next = self.layers[il + 1].attn_norm.float_data();
1697                let (sg, _, _, _) = seg.as_mut().unwrap();
1698                if sg[il].is_none() {
1699                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1700                    e.stream().synchronize()?;
1701                    e.stream()
1702                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1703                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1704                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1705                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1706                        Ok(())
1707                    })();
1708                    let g = e.stream().end_capture(
1709                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1710                    );
1711                    r?;
1712                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1713                }
1714                sg[il].as_ref().unwrap().launch()?;
1715            } else {
1716                if il + 1 < hi {
1717                    let w_next = self.layers[il + 1].attn_norm.float_data();
1718                    if f16fuse {
1719                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1720                    } else {
1721                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1722                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1723                    }
1724                } else {
1725                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1726                }
1727            }
1728            anat_mark!(4);
1729            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
1730            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
1731            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
1732            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
1733            // unset (the default) costs one OnceLock read per layer.
1734            if let Some(path) = Self::prime_trace_path() {
1735                let row = (base + t - 1) as usize;
1736                let host = e.dtoh(x_nxt)?;
1737                let last = &host[(t - 1) * n_embd..t * n_embd];
1738                use std::io::Write as _;
1739                let mut f = std::fs::OpenOptions::new()
1740                    .create(true)
1741                    .append(true)
1742                    .open(path)?;
1743                let mut h64: u64 = 0xcbf29ce484222325;
1744                for v in last {
1745                    h64 ^= v.to_bits() as u64;
1746                    h64 = h64.wrapping_mul(0x100000001b3);
1747                }
1748                writeln!(
1749                    f,
1750                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1751                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1752                    last[0], last[1], last[2]
1753                )?;
1754            }
1755            // dflash/dspark tap (no-op when no sink armed): post-layer residual rows for
1756            // drafter conditioning — the qwen twin of the gemma4 tap sites.
1757            self.dflash_tap(e, cache, il, x_nxt, t)?;
1758            std::mem::swap(&mut x_cur, &mut x_nxt);
1759        }
1760        if anat {
1761            let s = Self::prime_anatomy_slots();
1762            let ms = |i: usize| s[i].load(std::sync::atomic::Ordering::Relaxed) as f64 / 1.0e6;
1763            eprintln!(
1764                "[prime-anatomy] cumulative ms: attn_full={:.1} gdn_linear={:.1} moe={:.1} \
1765                 dense={:.1} norms_adds={:.1} (t={t}, layers {lo}..{hi})",
1766                ms(0),
1767                ms(1),
1768                ms(2),
1769                ms(3),
1770                ms(4)
1771            );
1772        }
1773        // hidden-stack return: clone the final x out of the slab
1774        let mut x = e.uninit(t * n_embd)?;
1775        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1776        drop(slab_guard);
1777        Ok(x)
1778    }
1779
1780    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
1781    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
1782    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
1783    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
1784    fn prime_chunk_epilogue(
1785        &self,
1786        e: &Engine,
1787        x: CudaSlice<f32>,
1788        t: usize,
1789        cache: &mut Cache,
1790    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1791        let n_embd = self.cfg.n_embd as usize;
1792        let eps = self.cfg.rms_eps;
1793        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
1794        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
1795        // the post-norm copy happens after hn exists).
1796        let mut h_seed = e.uninit(n_embd)?;
1797        if !crate::spec::spec_hpost() {
1798            e.copy_view_into(
1799                &mut h_seed,
1800                0,
1801                &x.slice((t - 1) * n_embd..t * n_embd),
1802                n_embd,
1803            )?;
1804        }
1805        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
1806        let mut hn = e.uninit(t * n_embd)?;
1807        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1808        if crate::spec::spec_hpost() {
1809            e.copy_view_into(
1810                &mut h_seed,
1811                0,
1812                &hn.slice((t - 1) * n_embd..t * n_embd),
1813                n_embd,
1814            )?;
1815        }
1816        let last = e.view(&hn, t * n_embd);
1817        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1818        let mut hlast = e.uninit(n_embd)?;
1819        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1820        let logits = e.matmul(&self.output, &hlast, 1)?;
1821        cache.pos += t;
1822        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
1823        // post-norm stack hn (MEMRA_SPEC_HPOST).
1824        Ok((
1825            e.dtoh(&logits)?,
1826            h_seed,
1827            if crate::spec::spec_hpost() { hn } else { x },
1828        ))
1829    }
1830
1831    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
1832    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
1833    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
1834    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
1835    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
1836    /// prefill kernels. Structure mirrors the verify split exactly:
1837    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
1838    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
1839    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
1840    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
1841    ///                  there via the sharded loader) → `publish_to`
1842    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
1843    /// round's stage-freed buffers must not be reused under the caller's queued reads);
1844    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
1845    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
1846    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
1847    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
1848    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
1849    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
1850    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
1851    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
1852    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
1853    /// and its liveness counter is bumped here — the gate goes green with this function.
1854    fn prime_chunk_ppn(
1855        &self,
1856        e: &Engine,
1857        tokens: &[u32],
1858        cache: &mut Cache,
1859        seq_end: usize,
1860        fence: &[usize],
1861    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1862        let rt = crate::pp::PpNRt::get(e)?;
1863        let n_st = fence.len() - 1;
1864        assert_eq!(
1865            rt.n_stages(),
1866            n_st,
1867            "PpNRt stage count {} != fence stages {n_st}",
1868            rt.n_stages()
1869        );
1870        let n_embd = self.cfg.n_embd as usize;
1871        let t = tokens.len();
1872        let base = cache.pos;
1873        debug_assert!(
1874            seq_end >= base + t,
1875            "prime_chunk_ppn: seq_end must cover this chunk"
1876        );
1877        let payload = t * n_embd;
1878        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1879        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1880        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1881        let caller_stream = e.stream();
1882        rt.fence_stages_behind(&caller_stream)?;
1883
1884        if n_st == 2 {
1885            let slot =
1886                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1887            let x =
1888                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1889            let out = {
1890                rt.bind_stage(1)?;
1891                let _st1 = rt.enter(1);
1892                let e1 = rt.engine(1, e);
1893                self.prime_chunk_epilogue(e1, x, t, cache)?
1894            };
1895            rt.publish_to(1, &caller_stream)?;
1896            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1897            return Ok(out);
1898        }
1899
1900        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1901
1902        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1903        let mut slot = {
1904            let _st0 = rt.enter(0);
1905            let e0 = rt.engine(0, e);
1906            let pos_d = e0.htod_i32(&pos)?;
1907            let x = self.embed(e0, tokens)?;
1908            let x =
1909                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1910            rt.tx(0, &x, payload)?
1911            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1912        };
1913
1914        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1915        for s in 1..n_st - 1 {
1916            let _st = rt.enter(s);
1917            let es = rt.engine(s, e);
1918            let pos_d = es.htod_i32(&pos)?;
1919            let x = rt.rx(s - 1, slot, payload)?;
1920            let x = self.prime_layers(
1921                es,
1922                x,
1923                fence[s],
1924                fence[s + 1],
1925                &pos_d,
1926                t,
1927                base,
1928                cache,
1929                seq_end,
1930            )?;
1931            slot = rt.tx(s, &x, payload)?;
1932        }
1933
1934        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1935        let _stl = rt.enter(n_st - 1);
1936        let el = rt.engine(n_st - 1, e);
1937        let pos_d = el.htod_i32(&pos)?;
1938        let x = rt.rx(n_st - 2, slot, payload)?;
1939        let x = self.prime_layers(
1940            el,
1941            x,
1942            fence[n_st - 1],
1943            fence[n_st],
1944            &pos_d,
1945            t,
1946            base,
1947            cache,
1948            seq_end,
1949        )?;
1950        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1951        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1952        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1953        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1954        // stage stream host-side, but the law is stated in events, not in a dtoh side
1955        // effect a later deferred form would remove.
1956        rt.publish_to(n_st - 1, &caller_stream)?;
1957        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1958        Ok(out)
1959    }
1960
1961    fn prime_pp2_stage0_enqueue(
1962        &self,
1963        e: &Engine,
1964        rt: &crate::pp::PpNRt,
1965        tokens: &[u32],
1966        cache: &mut Cache,
1967        seq_end: usize,
1968        fence: &[usize],
1969        base: usize,
1970        pipelined: bool,
1971    ) -> Result<usize, Box<dyn std::error::Error>> {
1972        let t = tokens.len();
1973        let n_embd = self.cfg.n_embd as usize;
1974        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1975        rt.bind_stage(0)?;
1976        let _st0 = rt.enter(0);
1977        let e0 = rt.engine(0, e);
1978        let pos_d = e0.htod_i32(&pos)?;
1979        let x = self.embed(e0, tokens)?;
1980        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1981        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1982        if pipelined {
1983            rt.tx_pipelined(0, &x, t * n_embd)
1984        } else {
1985            rt.tx(0, &x, t * n_embd)
1986        }
1987    }
1988
1989    fn prime_pp2_stage1_enqueue(
1990        &self,
1991        e: &Engine,
1992        rt: &crate::pp::PpNRt,
1993        slot: usize,
1994        t: usize,
1995        cache: &mut Cache,
1996        seq_end: usize,
1997        fence: &[usize],
1998        base: usize,
1999        pipelined: bool,
2000    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2001        let n_embd = self.cfg.n_embd as usize;
2002        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
2003        rt.bind_stage(1)?;
2004        let _st1 = rt.enter(1);
2005        let e1 = rt.engine(1, e);
2006        let pos_d = e1.htod_i32(&pos)?;
2007        let x = rt.rx(0, slot, t * n_embd)?;
2008        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
2009        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
2010    }
2011
2012    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
2013    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
2014    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
2015    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
2016    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
2017    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
2018    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
2019    /// bookkeeping still runs on the host per call — the real replay path moves the write
2020    /// slot to the len_d device counter (increment 3).
2021    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
2022    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
2023    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
2024    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
2025    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
2026    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
2027    pub fn prime_chunk_captured(
2028        &self,
2029        e: &Engine,
2030        x_in: &CudaSlice<f32>,
2031        pos_d: &CudaSlice<i32>,
2032        t: usize,
2033        cache: &mut Cache,
2034        len_d: &CudaSlice<i32>,
2035        logits_out: &mut CudaSlice<f32>,
2036        h_seed_out: &mut CudaSlice<f32>,
2037    ) -> Result<(), Box<dyn std::error::Error>> {
2038        let cfg = &self.cfg;
2039        let n_embd = cfg.n_embd as usize;
2040        let eps = cfg.rms_eps;
2041        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
2042        let mut x = e.uninit(t * n_embd)?;
2043        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
2044        for (il, layer) in self.layers.iter().enumerate() {
2045            let mut h = e.uninit(t * n_embd)?;
2046            let mut hx16: Option<CudaSlice<u8>> = None;
2047            if f16fuse {
2048                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
2049                e.rms_norm_f16out(
2050                    &x,
2051                    layer.attn_norm.float_data(),
2052                    &mut h,
2053                    &mut b16,
2054                    n_embd,
2055                    t,
2056                    eps,
2057                )?;
2058                hx16 = Some(b16);
2059            } else {
2060                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
2061            }
2062            let mixed = match &layer.mixer {
2063                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
2064                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
2065                // come from the caller (see step35_attn_pre_wo's doc note).
2066                Mixer::Full(fa) => {
2067                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
2068                }
2069                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2070                Mixer::Linear(la) => {
2071                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2072                    let g4 = match hx16.as_ref() {
2073                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
2074                        None => e.matmul_group(&ws, &h, t)?,
2075                    };
2076                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
2077                }
2078            };
2079            let mut x1 = e.uninit(t * n_embd)?;
2080            e.add(&x, &mixed, &mut x1, t * n_embd)?;
2081            let mut z = e.uninit(t * n_embd)?;
2082            let mut zx16: Option<CudaSlice<u8>> = None;
2083            if f16fuse {
2084                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
2085                e.rms_norm_f16out(
2086                    &x1,
2087                    layer.post_attn_norm.float_data(),
2088                    &mut z,
2089                    &mut b16,
2090                    n_embd,
2091                    t,
2092                    eps,
2093                )?;
2094                zx16 = Some(b16);
2095            } else {
2096                e.rms_norm(
2097                    &x1,
2098                    layer.post_attn_norm.float_data(),
2099                    &mut z,
2100                    n_embd,
2101                    t,
2102                    eps,
2103                )?;
2104            }
2105            let ffn_out = match &layer.ffn {
2106                crate::hybrid::Ffn::Dense {
2107                    ffn_gate,
2108                    ffn_up,
2109                    ffn_down,
2110                } => {
2111                    let n_ff = ffn_gate.out_features();
2112                    let mut g2 = match &zx16 {
2113                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
2114                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
2115                    };
2116                    let up = g2.pop().unwrap();
2117                    let gate = g2.pop().unwrap();
2118                    let mut act = e.uninit(t * n_ff)?;
2119                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
2120                    Self::ffn_act_lim(
2121                        e,
2122                        &self.cfg,
2123                        &gate,
2124                        &up,
2125                        1.0,
2126                        1.0,
2127                        self.cfg.clamp_shexp_at(il as u32),
2128                        &mut act,
2129                        t * n_ff,
2130                    )?;
2131                    e.matmul(ffn_down, &act, t)?
2132                }
2133                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
2134            };
2135            let mut x2 = e.uninit(t * n_embd)?;
2136            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
2137            x = x2;
2138        }
2139        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
2140        if !crate::spec::spec_hpost() {
2141            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
2142        }
2143        let mut hn = e.uninit(t * n_embd)?;
2144        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
2145        if crate::spec::spec_hpost() {
2146            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
2147        }
2148        let mut hlast = e.uninit(n_embd)?;
2149        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
2150        let logits = e.matmul(&self.output, &hlast, 1)?;
2151        let nv = logits.len();
2152        e.copy_into(logits_out, 0, &logits, nv)?;
2153        Ok(())
2154    }
2155
2156    fn step35_prime_batch_on() -> bool {
2157        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
2158    }
2159
2160    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
2161    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
2162    #[allow(clippy::too_many_arguments)]
2163    fn step35_prime_batch_layers(
2164        &self,
2165        e: &Engine,
2166        mut x: CudaSlice<f32>,
2167        lo: usize,
2168        hi: usize,
2169        ts: &[usize],
2170        offs: &[usize],
2171        pos_ds: &[CudaSlice<i32>],
2172        caches: &mut [&mut Cache],
2173    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2174        let cfg = &self.cfg;
2175        let n_embd = cfg.n_embd as usize;
2176        let eps = cfg.rms_eps;
2177        let b = ts.len();
2178        let total: usize = ts.iter().sum();
2179        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
2180
2181        let split = |e: &Engine,
2182                     y: &CudaSlice<f32>,
2183                     dim: usize|
2184         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2185            let mut out = Vec::with_capacity(b);
2186            for s in 0..b {
2187                let mut ys = e.uninit(ts[s] * dim)?;
2188                e.copy_view_into(
2189                    &mut ys,
2190                    0,
2191                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2192                    ts[s] * dim,
2193                )?;
2194                out.push(ys);
2195            }
2196            Ok(out)
2197        };
2198
2199        for il in lo..hi {
2200            let layer = &self.layers[il];
2201            let Mixer::Full(fa) = &layer.mixer else {
2202                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2203            };
2204
2205            let mut h = e.uninit(total * n_embd)?;
2206            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2207            if f16fuse {
2208                e.rms_norm_f16out(
2209                    &x,
2210                    layer.attn_norm.float_data(),
2211                    &mut h,
2212                    &mut hx16,
2213                    n_embd,
2214                    total,
2215                    eps,
2216                )?;
2217            } else {
2218                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2219            }
2220
2221            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
2222            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
2223            // application stay verbatim.
2224            let gate_w = fa
2225                .attn_gate
2226                .as_ref()
2227                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2228            let mut g4 = if f16fuse {
2229                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2230            } else {
2231                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2232            };
2233            let gate = g4.pop().unwrap();
2234            let mut parts: Vec<Vec<CudaSlice<f32>>> =
2235                (0..b).map(|_| Vec::with_capacity(3)).collect();
2236            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2237                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2238                    parts[s].push(ys);
2239                }
2240            }
2241            let gates = split(e, &gate, gate_w.out_features())?;
2242            let geometry = self.step35_geom(il);
2243            let hd = geometry.head_dim_k as usize;
2244            let nh = geometry.n_head as usize;
2245            let mut ag_cat = e.uninit(total * nh * hd)?;
2246            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2247                let ag = self.step35_attn_pre_wo(
2248                    e,
2249                    fa,
2250                    g3s,
2251                    None,
2252                    Some(&gate),
2253                    &pos_ds[s],
2254                    ts[s],
2255                    Some(&mut *caches[s]),
2256                    il,
2257                    ts[s],
2258                )?;
2259                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2260            }
2261            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2262
2263            let mut x1 = e.uninit(total * n_embd)?;
2264            let mut z = e.uninit(total * n_embd)?;
2265            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2266            if f16fuse {
2267                e.add_rms_norm_f16out(
2268                    &x,
2269                    &mixed,
2270                    layer.post_attn_norm.float_data(),
2271                    &mut x1,
2272                    &mut z,
2273                    &mut zx16,
2274                    n_embd,
2275                    total,
2276                    eps,
2277                )?;
2278            } else {
2279                e.add(&x, &mixed, &mut x1, total * n_embd)?;
2280                e.rms_norm(
2281                    &x1,
2282                    layer.post_attn_norm.float_data(),
2283                    &mut z,
2284                    n_embd,
2285                    total,
2286                    eps,
2287                )?;
2288            }
2289
2290            let ffn_out = match &layer.ffn {
2291                crate::hybrid::Ffn::Dense {
2292                    ffn_gate,
2293                    ffn_up,
2294                    ffn_down,
2295                } => {
2296                    let n_ff = ffn_gate.out_features();
2297                    let mut g2 = if f16fuse {
2298                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2299                    } else {
2300                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2301                    };
2302                    let up = g2.pop().unwrap();
2303                    let gate = g2.pop().unwrap();
2304                    let mut act = e.uninit(total * n_ff)?;
2305                    let d_lim = cfg.clamp_shexp_at(il as u32);
2306                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2307                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2308                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2309                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2310                            Some(y) => y,
2311                            None => e.matmul(ffn_down, &act, total)?,
2312                        }
2313                    } else {
2314                        Self::ffn_act_lim(
2315                            e,
2316                            cfg,
2317                            &gate,
2318                            &up,
2319                            1.0,
2320                            1.0,
2321                            d_lim,
2322                            &mut act,
2323                            total * n_ff,
2324                        )?;
2325                        e.matmul(ffn_down, &act, total)?
2326                    }
2327                }
2328                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2329            };
2330            let mut x2 = e.uninit(total * n_embd)?;
2331            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2332            x = x2;
2333        }
2334        Ok(x)
2335    }
2336
2337    fn step35_prime_batch_epilogue(
2338        &self,
2339        e: &Engine,
2340        x: CudaSlice<f32>,
2341        ts: &[usize],
2342        offs: &[usize],
2343        caches: &mut [&mut Cache],
2344    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2345        let n_embd = self.cfg.n_embd as usize;
2346        let total: usize = ts.iter().sum();
2347        let mut hn = e.uninit(total * n_embd)?;
2348        e.rms_norm(
2349            &x,
2350            self.output_norm.float_data(),
2351            &mut hn,
2352            n_embd,
2353            total,
2354            self.cfg.rms_eps,
2355        )?;
2356
2357        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2358        let mut out = Vec::with_capacity(ts.len());
2359        for s in 0..ts.len() {
2360            let mut hidden = e.uninit(ts[s] * n_embd)?;
2361            e.copy_view_into(
2362                &mut hidden,
2363                0,
2364                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2365                ts[s] * n_embd,
2366            )?;
2367            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2368            let mut h_seed = e.uninit(n_embd)?;
2369            e.copy_view_into(
2370                &mut h_seed,
2371                0,
2372                &hidden_src.slice(last0..last0 + n_embd),
2373                n_embd,
2374            )?;
2375            // Exactness-first: the serial reference runs the output head at m=1.
2376            let mut hlast = e.uninit(n_embd)?;
2377            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2378            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2379            caches[s].pos += ts[s];
2380            out.push((logits, h_seed, hidden));
2381        }
2382        Ok(out)
2383    }
2384
2385    fn step35_prime_cache_batch(
2386        &self,
2387        e: &Engine,
2388        prompts: &[&[u32]],
2389        caches: &mut [&mut Cache],
2390    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2391        if crate::pp::pp_host_bounce_active()
2392            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2393        {
2394            return Err(
2395                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2396                 stage split; refusing an unsplit remote-weight walk"
2397                    .into(),
2398            );
2399        }
2400        if !Self::step35_prime_batch_on() {
2401            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2402        }
2403        if caches.iter().any(|c| c.pos != 0) {
2404            return Err(
2405                "step35 batched prime currently supports complete fresh prompts only; \
2406                 continuation/tick chunks require per-request queued_after"
2407                    .into(),
2408            );
2409        }
2410
2411        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2412        for &t in &ts {
2413            assert!(
2414                t >= PRIME_MIN_T,
2415                "step35 batched prime needs T >= {PRIME_MIN_T}"
2416            );
2417        }
2418        for (s, c) in caches.iter().enumerate() {
2419            assert!(
2420                ts[s] <= c.max_ctx,
2421                "step35 batched prime exceeds cache max_ctx"
2422            );
2423        }
2424        let offs: Vec<usize> = ts
2425            .iter()
2426            .scan(0usize, |a, &t| {
2427                let o = *a;
2428                *a += t;
2429                Some(o)
2430            })
2431            .collect();
2432        let total: usize = ts.iter().sum();
2433        let payload = total * self.cfg.n_embd as usize;
2434        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2435        let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2436        let upload_positions =
2437            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2438                positions
2439                    .iter()
2440                    .map(|p| e.htod_i32(p))
2441                    .collect::<Result<_, _>>()
2442            };
2443
2444        static ONCE: std::sync::Once = std::sync::Once::new();
2445        ONCE.call_once(|| {
2446            eprintln!(
2447                "[step35-prime-batch] first concat prime: B={} tokens={total}",
2448                prompts.len()
2449            );
2450        });
2451
2452        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2453            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2454                let rt = crate::pp::PpNRt::get(e)?;
2455                let n_st = fence.len() - 1;
2456                assert_eq!(
2457                    rt.n_stages(),
2458                    n_st,
2459                    "step35 prime batch stage count mismatch"
2460                );
2461                let caller_stream = e.stream();
2462                rt.fence_stages_behind(&caller_stream)?;
2463
2464                let mut slot = {
2465                    let _st0 = rt.enter(0);
2466                    let e0 = rt.engine(0, e);
2467                    let pos_ds = upload_positions(e0)?;
2468                    let x = self.embed(e0, &cat_tokens)?;
2469                    let x = self.step35_prime_batch_layers(
2470                        e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2471                    )?;
2472                    rt.tx(0, &x, payload)?
2473                };
2474                for s in 1..n_st - 1 {
2475                    let _st = rt.enter(s);
2476                    let es = rt.engine(s, e);
2477                    let pos_ds = upload_positions(es)?;
2478                    let x = rt.rx(s - 1, slot, payload)?;
2479                    let x = self.step35_prime_batch_layers(
2480                        es,
2481                        x,
2482                        fence[s],
2483                        fence[s + 1],
2484                        &ts,
2485                        &offs,
2486                        &pos_ds,
2487                        caches,
2488                    )?;
2489                    slot = rt.tx(s, &x, payload)?;
2490                }
2491
2492                let _stl = rt.enter(n_st - 1);
2493                let el = rt.engine(n_st - 1, e);
2494                let pos_ds = upload_positions(el)?;
2495                let x = rt.rx(n_st - 2, slot, payload)?;
2496                let x = self.step35_prime_batch_layers(
2497                    el,
2498                    x,
2499                    fence[n_st - 1],
2500                    fence[n_st],
2501                    &ts,
2502                    &offs,
2503                    &pos_ds,
2504                    caches,
2505                )?;
2506                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2507                rt.publish_to(n_st - 1, &caller_stream)?;
2508                crate::pp::STEP35_PRIME_BATCH_SPLITS
2509                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2510                out
2511            } else {
2512                let pos_ds = upload_positions(e)?;
2513                let x = self.embed(e, &cat_tokens)?;
2514                let x = self.step35_prime_batch_layers(
2515                    e,
2516                    x,
2517                    0,
2518                    self.layers.len(),
2519                    &ts,
2520                    &offs,
2521                    &pos_ds,
2522                    caches,
2523                )?;
2524                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2525            }
2526        } else {
2527            let pos_ds = upload_positions(e)?;
2528            let x = self.embed(e, &cat_tokens)?;
2529            let x = self.step35_prime_batch_layers(
2530                e,
2531                x,
2532                0,
2533                self.layers.len(),
2534                &ts,
2535                &offs,
2536                &pos_ds,
2537                caches,
2538            )?;
2539            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2540        };
2541        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2542        Ok(out)
2543    }
2544
2545    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
2546    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
2547    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
2548    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
2549    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
2550    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
2551    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
2552    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
2553    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
2554    /// over the quantized past; Linear: the stateful pad_view twin — the same state
2555    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
2556    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
2557    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
2558    /// back to single-chunk serving).
2559    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
2560    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
2561    pub fn prime_cache_batch(
2562        &self,
2563        e: &Engine,
2564        prompts: &[&[u32]],
2565        caches: &mut [&mut Cache],
2566    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2567        let cfg = &self.cfg;
2568        let n_embd = cfg.n_embd as usize;
2569        let eps = cfg.rms_eps;
2570        let b = prompts.len();
2571        assert!(b >= 1 && b == caches.len());
2572        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2573        let carried = pos0s.iter().any(|&p| p > 0);
2574        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
2575        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
2576        // generic concat attn core below (uniform geometry, no per-layer swa window, no
2577        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
2578        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
2579        if cfg.gemma4.is_some() {
2580            return Err(
2581                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2582                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2583                    .into(),
2584            );
2585        }
2586        // Step35 has a dedicated concat walk: the generic core below cannot express its
2587        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
2588        if cfg.step35.is_some() {
2589            return self.step35_prime_cache_batch(e, prompts, caches);
2590        }
2591        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2592        for &t in &ts {
2593            assert!(
2594                t >= PRIME_MIN_T,
2595                "prime_cache_batch needs T >= {PRIME_MIN_T}"
2596            );
2597        }
2598        for (s, c) in caches.iter().enumerate() {
2599            assert!(
2600                c.pos + ts[s] <= c.max_ctx,
2601                "prime_cache_batch: prompt exceeds cache max_ctx"
2602            );
2603        }
2604        let total: usize = ts.iter().sum();
2605        let offs: Vec<usize> = ts
2606            .iter()
2607            .scan(0usize, |a, &t| {
2608                let o = *a;
2609                *a += t;
2610                Some(o)
2611            })
2612            .collect();
2613        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
2614        let pos_ds: Vec<CudaSlice<i32>> = ts
2615            .iter()
2616            .zip(&pos0s)
2617            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2618            .collect::<Result<_, _>>()?;
2619        // split a concat [total, dim] buffer into per-seq copies
2620        let split = |e: &Engine,
2621                     y: &CudaSlice<f32>,
2622                     dim: usize|
2623         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2624            let mut out = Vec::with_capacity(b);
2625            for s in 0..b {
2626                let mut ys = e.uninit(ts[s] * dim)?;
2627                e.copy_view_into(
2628                    &mut ys,
2629                    0,
2630                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2631                    ts[s] * dim,
2632                )?;
2633                out.push(ys);
2634            }
2635            Ok(out)
2636        };
2637
2638        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2639        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
2640        for (il, layer) in self.layers.iter().enumerate() {
2641            let mut h = e.uninit(total * n_embd)?;
2642            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2643            e.rms_norm_f16out(
2644                &x,
2645                layer.attn_norm.float_data(),
2646                &mut h,
2647                &mut hx16,
2648                n_embd,
2649                total,
2650                eps,
2651            )?;
2652            // mixer: projection GROUP on the concat (m = total), stateful core per seq
2653            let mut mixed = e.uninit(total * n_embd)?;
2654            match &layer.mixer {
2655                Mixer::Full(fa) => {
2656                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2657                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
2658                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
2659                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
2660                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
2661                    // back to the per-seq dispatch.
2662                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
2663                    let (n_head, n_head_kv, head_dim) = (
2664                        geometry.n_head as usize,
2665                        geometry.n_head_kv as usize,
2666                        geometry.head_dim_k as usize,
2667                    );
2668                    let fa_scale = geometry.attention_scale();
2669                    let use_favl = !carried
2670                        && (2..=8).contains(&b)
2671                        && (head_dim == 256 || head_dim == 128)
2672                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2673                        && std::env::var("MEMRA_NOFA").is_err()
2674                        && std::env::var("MEMRA_FA_FLOOR").is_err()
2675                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2676                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2677                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2678                    if use_favl {
2679                        let (qf_w, kf_w, vf_w) = (
2680                            fa.wq.out_features(),
2681                            fa.wk.out_features(),
2682                            fa.wv.out_features(),
2683                        );
2684                        // Same bounds contract as `Engine::q_gate_split`, applied to the varlen
2685                        // twin's PER-TOKEN stride. `attn_pre_vl8` takes raw device pointers so it
2686                        // cannot check its own extents; `qf_w` is the wq out-features that set
2687                        // them, and `q_gate_split_vl` reads 2*head_dim per head out of it.
2688                        memra_gguf::config::check_fused_q_gate_extent(qf_w, head_dim, n_head, 1)?;
2689                        struct APre {
2690                            q: CudaSlice<f32>,
2691                            gate: Option<CudaSlice<f32>>,
2692                            qn: CudaSlice<f32>,
2693                            kn: CudaSlice<f32>,
2694                        }
2695                        let mut aps = Vec::with_capacity(b);
2696                        for &t in ts.iter().take(b) {
2697                            aps.push(APre {
2698                                q: e.uninit(t * n_head * head_dim)?,
2699                                gate: Some(e.uninit(t * n_head * head_dim)?),
2700                                qn: e.uninit(t * n_head * head_dim)?,
2701                                kn: e.uninit(t * n_head_kv * head_dim)?,
2702                            });
2703                        }
2704                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2705                            let kvl = caches[0].kv[il].as_ref().unwrap();
2706                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2707                        };
2708                        let pargs: Vec<crate::AttnPreVl> = (0..b)
2709                            .map(|s| {
2710                                let (o, t) = (offs[s], ts[s]);
2711                                let kvl = caches[s].kv[il].as_ref().unwrap();
2712                                assert!(
2713                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2714                                    "prime_cache_batch attn vl: fresh + capacity"
2715                                );
2716                                crate::AttnPreVl {
2717                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2718                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2719                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2720                                    q: e.addr_f32(&aps[s].q),
2721                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2722                                    qn: e.addr_f32(&aps[s].qn),
2723                                    kn: e.addr_f32(&aps[s].kn),
2724                                    kc: e.addr_u8(&kvl.k),
2725                                    vc: e.addr_u8(&kvl.v),
2726                                    t: t as i32,
2727                                    pad: 0,
2728                                }
2729                            })
2730                            .collect();
2731                        e.attn_pre_vl8(
2732                            &pargs,
2733                            fa.q_norm.float_data(),
2734                            fa.k_norm.float_data(),
2735                            head_dim,
2736                            geometry.n_rot as usize,
2737                            n_head,
2738                            n_head_kv,
2739                            self.cfg.rms_eps,
2740                            geometry.rope_base,
2741                            1.0,
2742                            kv_dim_k,
2743                            kv_dim_v,
2744                            ktb,
2745                            vtb,
2746                        )?;
2747                        for s in 0..b {
2748                            let kvl = caches[s].kv[il].as_mut().unwrap();
2749                            kvl.len += ts[s];
2750                            let new_len = kvl.len as i32;
2751                            e.set_i32_one(&mut kvl.len_d, new_len)?;
2752                        }
2753                        let mut attns = Vec::with_capacity(b);
2754                        let mut mirrors = Vec::with_capacity(b);
2755                        for &t in ts.iter().take(b) {
2756                            attns.push(e.uninit(t * n_head * head_dim)?);
2757                            let n = t * n_head_kv * head_dim;
2758                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2759                        }
2760                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
2761                        // promoted single-seq config is on; else the mma favl.
2762                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2763                            Ok("0") => false,
2764                            Ok("1") => true,
2765                            _ => cfg!(memra_hopper_mma),
2766                        };
2767                        if fa3_on {
2768                            let mut q16s = Vec::with_capacity(b);
2769                            let mut v16s = Vec::with_capacity(b);
2770                            for s in 0..b {
2771                                let t = ts[s];
2772                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2773                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2774                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2775                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2776                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2777                                e.f32_to_bf16_v(
2778                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2779                                    &mut v16,
2780                                    t * n_head_kv * head_dim,
2781                                )?;
2782                                q16s.push(q16);
2783                                v16s.push((k16, v16));
2784                            }
2785                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2786                            let mut kp = qp;
2787                            let mut vp = qp;
2788                            let mut op = [core::ptr::null_mut::<f32>(); 8];
2789                            let mut tsv = [0i32; 8];
2790                            for s in 0..b {
2791                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2792                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2793                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2794                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
2795                                tsv[s] = ts[s] as i32;
2796                            }
2797                            let rc = unsafe {
2798                                crate::fa3_vl_raw(
2799                                    qp.as_ptr(),
2800                                    kp.as_ptr(),
2801                                    vp.as_ptr(),
2802                                    op.as_ptr(),
2803                                    tsv.as_ptr(),
2804                                    b as i32,
2805                                    n_head as i32,
2806                                    n_head_kv as i32,
2807                                    head_dim as i32,
2808                                    fa_scale,
2809                                    e.stream().cu_stream() as *mut core::ffi::c_void,
2810                                )
2811                            };
2812                            if rc != 0 {
2813                                return Err(format!("memra_fa3_vl rc={rc}").into());
2814                            }
2815                        } else {
2816                            let fargs: Vec<crate::FaSeqVl> = (0..b)
2817                                .map(|s| crate::FaSeqVl {
2818                                    q: e.addr_f32(&aps[s].qn),
2819                                    k16: e.addr_u8(&mirrors[s].0),
2820                                    v16: e.addr_u8(&mirrors[s].1),
2821                                    o: e.addr_f32(&attns[s]),
2822                                    kf: e.addr_f32(&aps[s].kn),
2823                                    vf: e.addr_f32v(
2824                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2825                                    ),
2826                                    t: ts[s] as i32,
2827                                    pad: 0,
2828                                })
2829                                .collect();
2830                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2831                        }
2832                        for (s, attn) in attns.into_iter().enumerate() {
2833                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
2834                                e,
2835                                attn,
2836                                &aps[s].gate,
2837                                ts[s],
2838                                n_head,
2839                                head_dim,
2840                            )?;
2841                            let mut done = false;
2842                            if let Some(xh) = &ag16 {
2843                                done = e.try_f16_gemm_pre_into_off(
2844                                    &fa.wo,
2845                                    xh,
2846                                    ts[s],
2847                                    &mut mixed,
2848                                    offs[s] * n_embd,
2849                                )?;
2850                            }
2851                            if !done {
2852                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2853                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2854                            }
2855                        }
2856                    } else {
2857                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
2858                            (0..b).map(|_| Vec::new()).collect();
2859                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2860                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2861                                parts[s].push(ys);
2862                            }
2863                        }
2864                        for (s, g3s) in parts.into_iter().enumerate() {
2865                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
2866                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
2867                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2868                            )?;
2869                            let mut done = false;
2870                            if let Some(xh) = &ag16 {
2871                                done = e.try_f16_gemm_pre_into_off(
2872                                    &fa.wo,
2873                                    xh,
2874                                    ts[s],
2875                                    &mut mixed,
2876                                    offs[s] * n_embd,
2877                                )?;
2878                            }
2879                            if !done {
2880                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2881                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2882                            }
2883                        }
2884                    }
2885                }
2886                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2887                Mixer::Linear(la) => {
2888                    // task #16: NO split copies (cores read row-offset views of the concat
2889                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
2890                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
2891                    // varlen K5 launch for all sequences.
2892                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2893                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2894                    let outs =
2895                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2896                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2897                        let (o, t) = (offs[s], ts[s]);
2898                        let mut done = false;
2899                        if let Some(xh) = &gn16 {
2900                            done = e.try_f16_gemm_pre_into_off(
2901                                &la.ssm_out,
2902                                xh,
2903                                t,
2904                                &mut mixed,
2905                                o * n_embd,
2906                            )?;
2907                        }
2908                        if !done {
2909                            let m = e.matmul(&la.ssm_out, &gn, t)?;
2910                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2911                        }
2912                    }
2913                }
2914            }
2915            let mut x1 = e.uninit(total * n_embd)?;
2916            let mut z = e.uninit(total * n_embd)?;
2917            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2918            e.add_rms_norm_f16out(
2919                &x,
2920                &mixed,
2921                layer.post_attn_norm.float_data(),
2922                &mut x1,
2923                &mut z,
2924                &mut zx16,
2925                n_embd,
2926                total,
2927                eps,
2928            )?;
2929            let ffn_out = match &layer.ffn {
2930                crate::hybrid::Ffn::Dense {
2931                    ffn_gate,
2932                    ffn_up,
2933                    ffn_down,
2934                } => {
2935                    let n_ff = ffn_gate.out_features();
2936                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2937                    let up = g2.pop().unwrap();
2938                    let gate = g2.pop().unwrap();
2939                    let mut act = e.uninit(total * n_ff)?;
2940                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
2941                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
2942                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
2943                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
2944                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2945                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2946                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2947                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2948                            Some(y) => y,
2949                            None => e.matmul(ffn_down, &act, total)?,
2950                        }
2951                    } else {
2952                        Self::ffn_act_lim(
2953                            e,
2954                            &self.cfg,
2955                            &gate,
2956                            &up,
2957                            1.0,
2958                            1.0,
2959                            d_lim,
2960                            &mut act,
2961                            total * n_ff,
2962                        )?;
2963                        e.matmul(ffn_down, &act, total)?
2964                    }
2965                }
2966                crate::hybrid::Ffn::Moe(m) => {
2967                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2968                }
2969            };
2970            let mut x2 = e.uninit(total * n_embd)?;
2971            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2972            x = x2;
2973        }
2974        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
2975        let mut hn = e.uninit(total * n_embd)?;
2976        e.rms_norm(
2977            &x,
2978            self.output_norm.float_data(),
2979            &mut hn,
2980            n_embd,
2981            total,
2982            eps,
2983        )?;
2984        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
2985        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
2986        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
2987        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
2988        // argmax battery arbitrates, same as every other prefill GEMM change.
2989        let mut hcat = e.uninit(b * n_embd)?;
2990        for s in 0..b {
2991            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2992            e.copy_view_into(
2993                &mut hcat,
2994                s * n_embd,
2995                &hn.slice(last0..last0 + n_embd),
2996                n_embd,
2997            )?;
2998        }
2999        let logits_cat = if b >= 2 {
3000            e.try_f16_gemm(&self.output, &hcat, b)?
3001        } else {
3002            None
3003        };
3004        let logits_host: Option<Vec<f32>> = match &logits_cat {
3005            Some(lc) => Some(e.dtoh(lc)?),
3006            None => None,
3007        };
3008        let n_vocab = self.output.out_features();
3009        let mut hidden_all = if crate::spec::spec_hpost() {
3010            split(e, &hn, n_embd)?
3011        } else {
3012            split(e, &x, n_embd)?
3013        };
3014        let mut out = Vec::with_capacity(b);
3015        for s in 0..b {
3016            let last0 = (offs[s] + ts[s] - 1) * n_embd;
3017            let mut h_seed = e.uninit(n_embd)?;
3018            if !crate::spec::spec_hpost() {
3019                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
3020            } else {
3021                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
3022            }
3023            let logits = match &logits_host {
3024                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
3025                None => {
3026                    let mut hlast = e.uninit(n_embd)?;
3027                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
3028                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
3029                }
3030            };
3031            caches[s].pos += ts[s];
3032            out.push((logits, h_seed, hidden_all.remove(0)));
3033        }
3034        Ok(out)
3035    }
3036
3037    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
3038    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
3039    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
3040    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
3041    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
3042    ///
3043    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
3044    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
3045    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
3046    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
3047    #[allow(clippy::too_many_arguments)]
3048    fn full_attn_prime(
3049        &self,
3050        e: &Engine,
3051        fa: &FullAttnLayer,
3052        h: &CudaSlice<f32>,
3053        hx: Option<&CudaSlice<u8>>,
3054        pos_d: &CudaSlice<i32>,
3055        t: usize,
3056        cache: &mut Cache,
3057        il: usize,
3058        seq_end: usize,
3059    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3060        if self.cfg.step35.is_some() {
3061            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
3062        }
3063        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
3064        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
3065        // this single-seq path composes proj+core identically (byte-for-byte the old body).
3066        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
3067        let g3 = match hx {
3068            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
3069            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
3070        };
3071        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
3072    }
3073
3074    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
3075    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
3076    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
3077    fn full_attn_prime_core(
3078        &self,
3079        e: &Engine,
3080        fa: &FullAttnLayer,
3081        g3: Vec<CudaSlice<f32>>,
3082        pos_d: &CudaSlice<i32>,
3083        t: usize,
3084        cache: &mut Cache,
3085        il: usize,
3086    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3087        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
3088        if let Some(xh) = &ag16 {
3089            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
3090                return Ok(y);
3091            }
3092        }
3093        Ok(e.matmul(&fa.wo, &attn_g, t)?)
3094    }
3095
3096    fn full_attn_prime_core_inner(
3097        &self,
3098        e: &Engine,
3099        fa: &FullAttnLayer,
3100        g3: Vec<CudaSlice<f32>>,
3101        pos_d: &CudaSlice<i32>,
3102        t: usize,
3103        cache: &mut Cache,
3104        il: usize,
3105    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3106        let cfg = &self.cfg;
3107        let geometry = cfg.full_attention_geometry_at(il as u32);
3108        let n_head = geometry.n_head as usize;
3109        let n_head_kv = geometry.n_head_kv as usize;
3110        let head_dim = geometry.head_dim_k as usize;
3111        let scale = geometry.attention_scale();
3112        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
3113        let AttnPre { q, k, v, gate } = pre;
3114        let mut attn = e.uninit(t * n_head * head_dim)?;
3115        self.full_attn_prime_fa_dispatch(
3116            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
3117        )?;
3118        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
3119    }
3120
3121    /// task #18 (attn side): projections tail through KV append — everything before the
3122    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
3123    /// present BEFORE this chunk's append (base_len; 0 == fresh).
3124    #[allow(clippy::type_complexity)]
3125    fn full_attn_prime_pre_fa(
3126        &self,
3127        e: &Engine,
3128        fa: &FullAttnLayer,
3129        mut g3: Vec<CudaSlice<f32>>,
3130        pos_d: &CudaSlice<i32>,
3131        t: usize,
3132        cache: &mut Cache,
3133        il: usize,
3134    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
3135        let cfg = &self.cfg;
3136        let geometry = cfg.full_attention_geometry_at(il as u32);
3137        let n_head = geometry.n_head as usize;
3138        let n_head_kv = geometry.n_head_kv as usize;
3139        let head_dim = geometry.head_dim_k as usize;
3140        let eps = cfg.rms_eps;
3141
3142        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
3143        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
3144        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
3145        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3146        let v = g3.pop().unwrap();
3147        let mut k = g3.pop().unwrap();
3148        let qf = g3.pop().unwrap();
3149        let (mut q, gate) = if gated {
3150            let mut q = e.uninit(t * n_head * head_dim)?;
3151            let mut gate = e.uninit(t * n_head * head_dim)?;
3152            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3153            (q, Some(gate))
3154        } else {
3155            (qf, None)
3156        };
3157
3158        let mut qn = e.uninit(t * n_head * head_dim)?;
3159        e.rms_norm(
3160            &q,
3161            fa.q_norm.float_data(),
3162            &mut qn,
3163            head_dim,
3164            n_head * t,
3165            eps,
3166        )?;
3167        q = qn;
3168        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3169        e.rms_norm(
3170            &k,
3171            fa.k_norm.float_data(),
3172            &mut kn,
3173            head_dim,
3174            n_head_kv * t,
3175            eps,
3176        )?;
3177        k = kn;
3178        let rope_dims = geometry.n_rot as usize;
3179        e.rope_neox(
3180            &mut q,
3181            pos_d,
3182            head_dim,
3183            rope_dims,
3184            n_head,
3185            t,
3186            geometry.rope_base,
3187            1.0,
3188        )?;
3189        e.rope_neox(
3190            &mut k,
3191            pos_d,
3192            head_dim,
3193            rope_dims,
3194            n_head_kv,
3195            t,
3196            geometry.rope_base,
3197            1.0,
3198        )?;
3199
3200        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
3201        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
3202        {
3203            let kvl = cache.kv[il].as_mut().unwrap();
3204            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3205            e.append_kv_quantized_rows(
3206                &k,
3207                &v,
3208                &mut kvl.k,
3209                &mut kvl.v,
3210                kvl.len,
3211                t,
3212                kvl.kv_dim_k,
3213                kvl.kv_dim_v,
3214                kvl.k_tok_bytes,
3215                kvl.v_tok_bytes,
3216                crate::Engine::kv_fp8_on(),
3217            )?;
3218            kvl.len += t;
3219            let new_len = kvl.len as i32;
3220            e.set_i32_one(&mut kvl.len_d, new_len)?;
3221        }
3222
3223        let base_len = {
3224            let kvl = cache.kv[il].as_ref().unwrap();
3225            kvl.len - t // KV rows present BEFORE this chunk's append above
3226        };
3227        Ok((AttnPre { q, k, v, gate }, base_len))
3228    }
3229
3230    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
3231    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
3232    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
3233    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
3234    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
3235    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
3236    #[allow(clippy::too_many_arguments)]
3237    fn full_attn_prime_fa_dispatch(
3238        &self,
3239        e: &Engine,
3240        q: &CudaSlice<f32>,
3241        k: &CudaSlice<f32>,
3242        v: &CudaSlice<f32>,
3243        attn: &mut CudaSlice<f32>,
3244        base_len: usize,
3245        t: usize,
3246        cache: &mut Cache,
3247        il: usize,
3248        head_dim: usize,
3249        n_head: usize,
3250        n_head_kv: usize,
3251        scale: f32,
3252    ) -> Result<(), Box<dyn std::error::Error>> {
3253        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
3254        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
3255        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
3256        // attend through the quantized cache exactly like every later chunk (quantize-then-
3257        // attend). One numeric class for every row => the chunk size cannot decide where a
3258        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
3259        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
3260        // pin-the-boundary approach).
3261        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
3262        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
3263        // with the fix unconditional, only re-introducing the class edge can prove the gate
3264        // still detects the mechanism. Never on in a measured default run.
3265        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3266            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3267                e.sdpa_naive(
3268                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3269                )?;
3270            } else {
3271                e.fa_prefill(
3272                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3273                )?;
3274            }
3275            return Ok(());
3276        }
3277        let kvl = cache.kv[il].as_ref().unwrap();
3278        let t_kv = base_len + t;
3279        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3280        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3281        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
3282        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
3283        // same numeric class, so the uniform contract holds on the fallback too.
3284        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3285            e.sdpa_naive_quantized_view(
3286                q,
3287                &k_view,
3288                &v_view,
3289                attn,
3290                head_dim,
3291                n_head,
3292                n_head_kv,
3293                t,
3294                t_kv,
3295                scale,
3296                true,
3297                kvl.k_tok_bytes,
3298                kvl.v_tok_bytes,
3299            )?;
3300            return Ok(());
3301        }
3302        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
3303        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
3304        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
3305        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
3306        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
3307        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
3308        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
3309        let deqw = std::env::var("MEMRA_PRIME_DEQW")
3310            .map(|v| v != "0")
3311            .unwrap_or(true);
3312        if deqw {
3313            e.fa_prefill_view_ws(
3314                q,
3315                &k_view,
3316                &v_view,
3317                attn,
3318                head_dim,
3319                n_head,
3320                n_head_kv,
3321                t,
3322                t_kv,
3323                scale,
3324                true,
3325                kvl.k_tok_bytes,
3326                kvl.v_tok_bytes,
3327                crate::Engine::kv_fp8_on(),
3328            )?;
3329        } else {
3330            e.fa_prefill_view(
3331                q,
3332                &k_view,
3333                &v_view,
3334                attn,
3335                head_dim,
3336                n_head,
3337                n_head_kv,
3338                t,
3339                t_kv,
3340                scale,
3341                true,
3342                kvl.k_tok_bytes,
3343                kvl.v_tok_bytes,
3344                crate::Engine::kv_fp8_on(),
3345            )?;
3346        }
3347        Ok(())
3348    }
3349
3350    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
3351    /// (bit-identical composition) and hands wo its fp16 operand directly.
3352    fn full_attn_prime_post_fa(
3353        &self,
3354        e: &Engine,
3355        attn: CudaSlice<f32>,
3356        gate: &Option<CudaSlice<f32>>,
3357        t: usize,
3358        n_head: usize,
3359        head_dim: usize,
3360    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3361        let (attn_g, ag16) = match gate {
3362            Some(gate) => {
3363                let n = t * n_head * head_dim;
3364                let mut ag = e.uninit(n)?;
3365                if Self::f16out_on(e, t) {
3366                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
3367                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3368                    (ag, Some(a16))
3369                } else {
3370                    let mut gsig = e.uninit(n)?;
3371                    e.sigmoid(gate, &mut gsig, n)?;
3372                    e.mul(&attn, &gsig, &mut ag, n)?;
3373                    (ag, None)
3374                }
3375            }
3376            None => (attn, None),
3377        };
3378        Ok((attn_g, ag16))
3379    }
3380
3381    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
3382    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
3383    /// carried THROUGH the cache like the spec verify does: carried-ring conv
3384    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
3385    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
3386    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
3387    fn linear_attn_prime(
3388        &self,
3389        e: &Engine,
3390        la: &LinearAttnLayer,
3391        h: &CudaSlice<f32>,
3392        hx: Option<&CudaSlice<u8>>,
3393        t: usize,
3394        cache: &mut Cache,
3395        il: usize,
3396    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3397        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
3398        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3399        let g4 = match hx {
3400            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3401            None => e.matmul_group(&ws, h, t)?,
3402        };
3403        self.linear_attn_prime_core(e, la, g4, t, cache, il)
3404    }
3405
3406    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
3407    fn linear_attn_prime_core(
3408        &self,
3409        e: &Engine,
3410        la: &LinearAttnLayer,
3411        mut g4: Vec<CudaSlice<f32>>,
3412        t: usize,
3413        cache: &mut Cache,
3414        il: usize,
3415    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3416        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3417    }
3418
3419    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
3420    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
3421    /// conv ring writes back from the true tail. None = classic path, byte-identical.
3422    #[allow(clippy::too_many_arguments)]
3423    fn linear_attn_prime_core_pad_inner(
3424        &self,
3425        e: &Engine,
3426        la: &LinearAttnLayer,
3427        mut g4: Vec<CudaSlice<f32>>,
3428        t: usize,
3429        cache: &mut Cache,
3430        il: usize,
3431        pad_len: Option<&CudaSlice<i32>>,
3432    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3433        // shim over the view twin (task #16): full-range views of the owned buffers.
3434        let ssm = self.cfg.ssm.as_ref().unwrap();
3435        let d_state = ssm.state_size as usize;
3436        let num_k = ssm.group_count as usize;
3437        let num_v = ssm.time_step_rank as usize;
3438        let key_dim = d_state * num_k;
3439        let value_dim = d_state * num_v;
3440        let conv_dim = key_dim * 2 + value_dim;
3441        let alpha = g4.pop().unwrap(); // [T, num_v]
3442        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3443        let z = g4.pop().unwrap(); // [T, value_dim]
3444        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3445        self.linear_attn_prime_core_pad_view(
3446            e,
3447            la,
3448            &qkv_mixed.slice(0..t * conv_dim),
3449            &z.slice(0..t * value_dim),
3450            &beta_raw.slice(0..t * num_v),
3451            &alpha.slice(0..t * num_v),
3452            t,
3453            cache,
3454            il,
3455            pad_len,
3456        )
3457    }
3458
3459    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
3460    /// shared verbatim by the per-seq scan path and the varlen batched path.
3461    #[allow(clippy::too_many_arguments)]
3462    fn linear_attn_gdn_prep(
3463        &self,
3464        e: &Engine,
3465        la: &LinearAttnLayer,
3466        qkv_mixed: &cudarc::driver::CudaView<f32>,
3467        beta_raw: &cudarc::driver::CudaView<f32>,
3468        alpha: &cudarc::driver::CudaView<f32>,
3469        t: usize,
3470        cache: &mut Cache,
3471        il: usize,
3472        pad_len: Option<&CudaSlice<i32>>,
3473    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3474        let cfg = &self.cfg;
3475        let ssm = cfg.ssm.as_ref().unwrap();
3476        let d_state = ssm.state_size as usize; // 128
3477        let num_k = ssm.group_count as usize; // 16
3478        let num_v = ssm.time_step_rank as usize; // 32
3479        let d_conv = ssm.conv_kernel as usize; // 4
3480        let key_dim = d_state * num_k; // 2048
3481        let value_dim = d_state * num_v; // 4096
3482        let conv_dim = key_dim * 2 + value_dim; // 8192
3483        let eps = cfg.rms_eps;
3484        debug_assert!(
3485            t >= d_conv - 1,
3486            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3487        );
3488
3489        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
3490        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
3491        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
3492        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
3493        let rl = cache.recur[il].as_mut().unwrap();
3494        let hk = Self::gdn_hk(e, t, num_v, num_k);
3495        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3496        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
3497        let mut q_g = e.uninit(d_state * hk * t)?;
3498        let mut k_g = e.uninit(d_state * hk * t)?;
3499        let mut v_g = e.uninit(d_state * num_v * t)?;
3500        if conv_fuse {
3501            e.ssm_conv1d_gdn_state_pad(
3502                qkv_mixed,
3503                &mut rl.conv_state,
3504                la.ssm_conv1d.float_data(),
3505                &mut q_g,
3506                &mut k_g,
3507                &mut v_g,
3508                conv_dim,
3509                t,
3510                d_conv,
3511                d_state,
3512                num_v,
3513                num_k,
3514                key_dim,
3515                hk,
3516                pad_len,
3517            )?;
3518        } else {
3519            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
3520            e.ssm_conv1d_tm_state_pad_v(
3521                qkv_mixed,
3522                &mut rl.conv_state,
3523                la.ssm_conv1d.float_data(),
3524                &mut conv_out,
3525                conv_dim,
3526                t,
3527                d_conv,
3528                pad_len,
3529            )?;
3530            e.qkv_to_gdn_repack(
3531                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3532            )?;
3533        }
3534        let mut q_l2 = e.uninit(d_state * hk * t)?;
3535        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
3536        // Emitted only where a consumer exists (the wgmma config) — on other arches the
3537        // alloc + epilogue stores would be pure waste.
3538        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3539            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3540            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3541            Some(qb)
3542        } else {
3543            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3544            None
3545        };
3546        let mut k_l2 = e.uninit(d_state * hk * t)?;
3547        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
3548        let kb16 = if Engine::l2_v2_on(d_state) {
3549            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3550            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3551            Some(kb)
3552        } else {
3553            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3554            None
3555        };
3556        let mut beta = e.uninit(t * num_v)?;
3557        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3558        let mut g_log = e.uninit(t * num_v)?;
3559        e.gdn_glog_v(
3560            alpha,
3561            la.ssm_dt.float_data(),
3562            la.ssm_a.float_data(),
3563            &mut g_log,
3564            num_v,
3565            t,
3566        )?;
3567        if let Some(len_d) = pad_len {
3568            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3569        }
3570        Ok(GdnPrep {
3571            hk,
3572            q_l2,
3573            k_l2,
3574            v_g,
3575            beta,
3576            g_log,
3577            kb16,
3578            qb16,
3579        })
3580    }
3581
3582    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
3583    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
3584    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
3585    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
3586    #[allow(clippy::too_many_arguments)]
3587    fn linear_attn_prime_core_batch(
3588        &self,
3589        e: &Engine,
3590        la: &LinearAttnLayer,
3591        g4: &[CudaSlice<f32>],
3592        offs: &[usize],
3593        ts: &[usize],
3594        caches: &mut [&mut Cache],
3595        il: usize,
3596    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3597        let ssm = self.cfg.ssm.as_ref().unwrap();
3598        let d_state = ssm.state_size as usize;
3599        let num_k = ssm.group_count as usize;
3600        let num_v = ssm.time_step_rank as usize;
3601        let key_dim = d_state * num_k;
3602        let value_dim = d_state * num_v;
3603        let conv_dim = key_dim * 2 + value_dim;
3604        let eps = self.cfg.rms_eps;
3605        let scale = 1.0 / (d_state as f32).sqrt();
3606        let b = ts.len();
3607        let c = Engine::gdn_chunk_size();
3608        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
3609        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
3610        let carried = caches.iter().any(|c| c.pos > 0);
3611        let use_vl = !carried
3612            && (2..=8).contains(&b)
3613            && Engine::gdn_chunked_enabled()
3614            && ts.iter().all(|&t| t >= 16)
3615            && e.gdn_mma_enabled(c)
3616            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3617        if !use_vl {
3618            return (0..b)
3619                .map(|s| {
3620                    let (o, t) = (offs[s], ts[s]);
3621                    self.linear_attn_prime_core_pad_view(
3622                        e,
3623                        la,
3624                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3625                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
3626                        &g4[2].slice(o * num_v..(o + t) * num_v),
3627                        &g4[3].slice(o * num_v..(o + t) * num_v),
3628                        t,
3629                        caches[s],
3630                        il,
3631                        None,
3632                    )
3633                })
3634                .collect();
3635        }
3636        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
3637        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
3638        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
3639        struct SeqBufs {
3640            conv_out: CudaSlice<f32>,
3641            q_g: CudaSlice<f32>,
3642            k_g: CudaSlice<f32>,
3643            v_g: CudaSlice<f32>,
3644            q_l2: CudaSlice<f32>,
3645            k_l2: CudaSlice<f32>,
3646            beta: CudaSlice<f32>,
3647            g_log: CudaSlice<f32>,
3648            gn: CudaSlice<f32>,
3649            gn16: CudaSlice<u8>,
3650        }
3651        let d_conv = ssm.conv_kernel as usize;
3652        let f16o = Self::f16out_on(e, 16);
3653        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
3654        let mut sb = Vec::with_capacity(b);
3655        let mut pres = Vec::with_capacity(b);
3656        for &t in ts.iter().take(b) {
3657            sb.push(SeqBufs {
3658                conv_out: e.uninit(conv_dim * t)?,
3659                q_g: e.uninit(d_state * hk * t)?,
3660                k_g: e.uninit(d_state * hk * t)?,
3661                v_g: e.uninit(d_state * num_v * t)?,
3662                q_l2: e.uninit(d_state * hk * t)?,
3663                k_l2: e.uninit(d_state * hk * t)?,
3664                beta: e.uninit(t * num_v)?,
3665                g_log: e.uninit(t * num_v)?,
3666                gn: e.uninit(d_state * num_v * t)?,
3667                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3668            });
3669            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3670        }
3671        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3672            .map(|s| {
3673                let (o, t) = (offs[s], ts[s]);
3674                let rl = caches[s].recur[il].as_ref().unwrap();
3675                crate::GdnPrepVl {
3676                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3677                    conv_state: e.addr_f32(&rl.conv_state),
3678                    conv_out: e.addr_f32(&sb[s].conv_out),
3679                    q_g: e.addr_f32(&sb[s].q_g),
3680                    k_g: e.addr_f32(&sb[s].k_g),
3681                    v_g: e.addr_f32(&sb[s].v_g),
3682                    q_l2: e.addr_f32(&sb[s].q_l2),
3683                    k_l2: e.addr_f32(&sb[s].k_l2),
3684                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3685                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3686                    beta: e.addr_f32(&sb[s].beta),
3687                    g_log: e.addr_f32(&sb[s].g_log),
3688                    o: e.addr_f32(&pres[s].o),
3689                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3690                    gn: e.addr_f32(&sb[s].gn),
3691                    gn16: e.addr_u8(&sb[s].gn16),
3692                    kb16: if Engine::l2_v2_on(d_state) {
3693                        e.addr_u8(&pres[s].kb16)
3694                    } else {
3695                        0
3696                    },
3697                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3698                        e.addr_u8(&pres[s].qb16)
3699                    } else {
3700                        0
3701                    },
3702                    t: t as i32,
3703                    pad: 0,
3704                }
3705            })
3706            .collect();
3707        let args: Vec<crate::GdnSeqVl> = (0..b)
3708            .map(|s| {
3709                let rl = caches[s].recur[il].as_ref().unwrap();
3710                crate::GdnSeqVl {
3711                    kb16: e.addr_u8(&pres[s].kb16),
3712                    gcum: e.addr_f32(&pres[s].gcum),
3713                    beta: e.addr_f32(&sb[s].beta),
3714                    u: e.addr_f32(&pres[s].u),
3715                    wb16: e.addr_u8(&pres[s].wb16),
3716                    y: e.addr_u8(&pres[s].y16),
3717                    ssnap: e.addr_u8(&pres[s].ssnap16),
3718                    state_in: e.addr_f32(&rl.ssm_state),
3719                    state_out: e.addr_f32(&rl.ssm_state_alt),
3720                    q: e.addr_f32(&sb[s].q_l2),
3721                    p: e.addr_f32(&pres[s].p),
3722                    o: e.addr_f32(&pres[s].o),
3723                    k: e.addr_f32(&sb[s].k_l2),
3724                    v: e.addr_f32(&sb[s].v_g),
3725                    g: e.addr_f32(&sb[s].g_log),
3726                    a: e.addr_f32(&pres[s].a),
3727                    w: e.addr_f32(&pres[s].w),
3728                    t: ts[s] as i32,
3729                    nc: pres[s].nc as i32,
3730                }
3731            })
3732            .collect();
3733        e.gdn_prep_vl8(
3734            &prep_args,
3735            la.ssm_conv1d.float_data(),
3736            la.ssm_dt.float_data(),
3737            la.ssm_a.float_data(),
3738            conv_dim,
3739            d_conv,
3740            d_state,
3741            num_v,
3742            num_k,
3743            key_dim,
3744            hk,
3745            eps,
3746        )?;
3747        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
3748        // both standalone mirror launches vanish on the default config.
3749        if !Engine::l2_v2_on(d_state) {
3750            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3751        }
3752        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
3753        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3754            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
3755            if !Engine::l2_v2_on(d_state) {
3756                for s in 0..b {
3757                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3758                }
3759            }
3760            let mut wa = [crate::GdnWVl::default(); 8];
3761            for s in 0..b {
3762                wa[s] = crate::GdnWVl {
3763                    qb16: e.addr_u8(&pres[s].qb16),
3764                    pb16: e.addr_u8(&pres[s].pb16),
3765                };
3766            }
3767            Some(crate::GdnWVl8(wa))
3768        } else {
3769            None
3770        };
3771        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3772        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3773        if f16o {
3774            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3775        }
3776        // per-seq state swap (+ non-f16out tail fallback)
3777        let mut out = Vec::with_capacity(b);
3778        for (s, bufs) in sb.into_iter().enumerate() {
3779            let rl = caches[s].recur[il].as_mut().unwrap();
3780            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3781            let (o, t) = (offs[s], ts[s]);
3782            let SeqBufs { mut gn, gn16, .. } = bufs;
3783            if f16o {
3784                out.push((gn, Some(gn16)));
3785            } else {
3786                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3787                e.gated_rmsnorm_zv(
3788                    &pres[s].o,
3789                    la.ssm_norm.float_data(),
3790                    &z_v,
3791                    &mut gn,
3792                    d_state,
3793                    num_v * t,
3794                    eps,
3795                )?;
3796                out.push((gn, None));
3797            }
3798        }
3799        Ok(out)
3800    }
3801
3802    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
3803    /// views of the CONCAT projection outputs directly (no per-seq split copies).
3804    /// Same kernels, same values, byte-identical to the Vec shim above.
3805    #[allow(clippy::too_many_arguments)]
3806    fn linear_attn_prime_core_pad_view(
3807        &self,
3808        e: &Engine,
3809        la: &LinearAttnLayer,
3810        qkv_mixed: &cudarc::driver::CudaView<f32>,
3811        z: &cudarc::driver::CudaView<f32>,
3812        beta_raw: &cudarc::driver::CudaView<f32>,
3813        alpha: &cudarc::driver::CudaView<f32>,
3814        t: usize,
3815        cache: &mut Cache,
3816        il: usize,
3817        pad_len: Option<&CudaSlice<i32>>,
3818    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3819        let cfg = &self.cfg;
3820        let ssm = cfg.ssm.as_ref().unwrap();
3821        let d_state = ssm.state_size as usize; // 128
3822        let num_v = ssm.time_step_rank as usize; // 32
3823        let eps = cfg.rms_eps;
3824        let scale = 1.0 / (d_state as f32).sqrt();
3825
3826        let prep =
3827            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3828
3829        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
3830        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
3831        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
3832        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
3833        // verify keep the sequential kernel).
3834        let mut o = e.uninit(d_state * num_v * t)?;
3835        let rl = cache.recur[il].as_mut().unwrap();
3836        {
3837            let crate::cache::RecurLayer {
3838                ssm_state,
3839                ssm_state_alt,
3840                ..
3841            } = rl;
3842            e.gdn_scan_prefill(
3843                &prep.q_l2,
3844                &prep.k_l2,
3845                &prep.v_g,
3846                &prep.g_log,
3847                &prep.beta,
3848                prep.kb16.as_ref(),
3849                prep.qb16.as_ref(),
3850                ssm_state,
3851                ssm_state_alt,
3852                &mut o,
3853                num_v,
3854                t,
3855                scale,
3856                prep.hk,
3857            )?;
3858        }
3859        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3860
3861        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
3862        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
3863        let mut gn = e.uninit(d_state * num_v * t)?;
3864        let gn16 = if Self::f16out_on(e, t) {
3865            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3866            e.gated_rmsnorm_f16out_zv(
3867                &o,
3868                la.ssm_norm.float_data(),
3869                z,
3870                &mut gn,
3871                &mut g16,
3872                d_state,
3873                num_v * t,
3874                eps,
3875            )?;
3876            Some(g16)
3877        } else {
3878            e.gated_rmsnorm_zv(
3879                &o,
3880                la.ssm_norm.float_data(),
3881                z,
3882                &mut gn,
3883                d_state,
3884                num_v * t,
3885                eps,
3886            )?;
3887            None
3888        };
3889        Ok((gn, gn16))
3890    }
3891
3892    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
3893    #[allow(clippy::too_many_arguments)]
3894    fn linear_attn_prime_core_pad(
3895        &self,
3896        e: &Engine,
3897        la: &LinearAttnLayer,
3898        g4: Vec<CudaSlice<f32>>,
3899        t: usize,
3900        cache: &mut Cache,
3901        il: usize,
3902        pad_len: Option<&CudaSlice<i32>>,
3903    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3904        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3905        if let Some(xh) = &gn16 {
3906            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3907                return Ok(y);
3908            }
3909        }
3910        Ok(e.matmul(&la.ssm_out, &gn, t)?)
3911    }
3912
3913    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
3914    ///
3915    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
3916    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
3917    pub fn full_attn(
3918        &self,
3919        e: &Engine,
3920        fa: &FullAttnLayer,
3921        h: &CudaSlice<f32>,
3922        pos_d: &CudaSlice<i32>,
3923        t: usize,
3924        il: usize,
3925    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3926        if self.cfg.step35.is_some() {
3927            return self.step35_attn(e, fa, h, pos_d, t, il);
3928        }
3929        let cfg = &self.cfg;
3930        let _n_embd = cfg.n_embd as usize;
3931        let geometry = cfg.full_attention_geometry_at(il as u32);
3932        let n_head = geometry.n_head as usize;
3933        let n_head_kv = geometry.n_head_kv as usize;
3934        let head_dim = geometry.head_dim_k as usize;
3935        let eps = cfg.rms_eps;
3936        let scale = geometry.attention_scale();
3937
3938        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
3939        // gate — wq out = n_head*head_dim, no split (see prime-path note).
3940        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3941        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
3942        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3943        let v = g3.pop().unwrap();
3944        let mut k = g3.pop().unwrap();
3945        let qf = g3.pop().unwrap();
3946        let (mut q, gate) = if gated {
3947            let mut q = e.uninit(t * n_head * head_dim)?;
3948            let mut gate = e.uninit(t * n_head * head_dim)?;
3949            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3950            (q, Some(gate))
3951        } else {
3952            (qf, None)
3953        };
3954
3955        // QK-norm (per head_dim row), then partial RoPE.
3956        let mut qn = e.uninit(t * n_head * head_dim)?;
3957        e.rms_norm(
3958            &q,
3959            fa.q_norm.float_data(),
3960            &mut qn,
3961            head_dim,
3962            n_head * t,
3963            eps,
3964        )?;
3965        q = qn;
3966        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3967        e.rms_norm(
3968            &k,
3969            fa.k_norm.float_data(),
3970            &mut kn,
3971            head_dim,
3972            n_head_kv * t,
3973            eps,
3974        )?;
3975        k = kn;
3976        let rope_dims = geometry.n_rot as usize;
3977        e.rope_neox(
3978            &mut q,
3979            pos_d,
3980            head_dim,
3981            rope_dims,
3982            n_head,
3983            t,
3984            geometry.rope_base,
3985            1.0,
3986        )?;
3987        e.rope_neox(
3988            &mut k,
3989            pos_d,
3990            head_dim,
3991            rope_dims,
3992            n_head_kv,
3993            t,
3994            geometry.rope_base,
3995            1.0,
3996        )?;
3997
3998        // SDPA
3999        let mut attn = e.uninit(t * n_head * head_dim)?;
4000        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
4001        // falls back to naive sdpa.
4002        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
4003            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
4004            e.sdpa_naive(
4005                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
4006            )?;
4007        } else {
4008            e.fa_prefill(
4009                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
4010            )?;
4011        }
4012
4013        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
4014        let attn_g = match &gate {
4015            Some(gate) => {
4016                let mut gsig = e.uninit(t * n_head * head_dim)?;
4017                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
4018                let mut ag = e.uninit(t * n_head * head_dim)?;
4019                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
4020                ag
4021            }
4022            None => attn,
4023        };
4024
4025        // o projection
4026        let o = e.matmul(&fa.wo, &attn_g, t)?;
4027        Ok(o)
4028    }
4029
4030    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
4031    pub fn linear_attn(
4032        &self,
4033        e: &Engine,
4034        la: &LinearAttnLayer,
4035        h: &CudaSlice<f32>,
4036        t: usize,
4037    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4038        let cfg = &self.cfg;
4039        let _n_embd = cfg.n_embd as usize;
4040        let ssm = cfg.ssm.as_ref().unwrap();
4041        let d_state = ssm.state_size as usize; // 128
4042        let num_k = ssm.group_count as usize; // 16
4043        let num_v = ssm.time_step_rank as usize; // 32
4044        let d_conv = ssm.conv_kernel as usize; // 4
4045        let head_k = d_state;
4046        let head_v = d_state;
4047        let key_dim = head_k * num_k; // 2048
4048        let value_dim = head_v * num_v; // 4096
4049        let conv_dim = key_dim * 2 + value_dim; // 8192
4050        let eps = cfg.rms_eps;
4051        let scale = 1.0 / (d_state as f32).sqrt();
4052
4053        // projections
4054        // grouped: one f16 activation convert feeds all four projections (matmul_group)
4055        let mut g4 = e.matmul_group(
4056            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
4057            h,
4058            t,
4059        )?;
4060        let alpha = g4.pop().unwrap(); // [T, num_v]
4061        let beta_raw = g4.pop().unwrap(); // [T, num_v]
4062        let z = g4.pop().unwrap(); // [T, value_dim]
4063        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
4064
4065        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
4066        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
4067        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
4068        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
4069        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
4070        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
4071        let _ = (head_k, head_v);
4072        let mut q_g = e.uninit(d_state * num_v * t)?;
4073        let mut k_g = e.uninit(d_state * num_v * t)?;
4074        let mut v_g = e.uninit(d_state * num_v * t)?;
4075        e.ssm_conv1d_gdn(
4076            &qkv_mixed,
4077            la.ssm_conv1d.float_data(),
4078            &mut q_g,
4079            &mut k_g,
4080            &mut v_g,
4081            conv_dim,
4082            t,
4083            d_conv,
4084            d_state,
4085            num_v,
4086            num_k,
4087            key_dim,
4088        )?;
4089        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
4090        let mut q_l2 = e.uninit(d_state * num_v * t)?;
4091        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
4092        let mut k_l2 = e.uninit(d_state * num_v * t)?;
4093        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
4094        let v_gd = v_g;
4095
4096        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
4097        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
4098        let mut beta = e.uninit(t * num_v)?;
4099        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
4100        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
4101        let mut g_log = e.uninit(t * num_v)?;
4102        e.gdn_glog(
4103            &alpha,
4104            la.ssm_dt.float_data(),
4105            la.ssm_a.float_data(),
4106            &mut g_log,
4107            num_v,
4108            t,
4109        )?;
4110
4111        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
4112        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
4113        let mut state_out = e.zeros(d_state * d_state * num_v)?;
4114        let mut o = e.uninit(d_state * num_v * t)?;
4115        e.gdn_scan_prefill(
4116            &q_l2,
4117            &k_l2,
4118            &v_gd,
4119            &g_log,
4120            &beta,
4121            None,
4122            None,
4123            &state_in,
4124            &mut state_out,
4125            &mut o,
4126            num_v,
4127            t,
4128            scale,
4129            num_v,
4130        )?;
4131
4132        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
4133        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
4134        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
4135        // o rows are (t*num_v+vh) too. Good.
4136        let mut gn = e.uninit(d_state * num_v * t)?;
4137        e.gated_rmsnorm(
4138            &o,
4139            la.ssm_norm.float_data(),
4140            &z,
4141            &mut gn,
4142            d_state,
4143            num_v * t,
4144            eps,
4145        )?;
4146
4147        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
4148        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
4149        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
4150        let out = e.matmul(&la.ssm_out, &gn, t)?;
4151        Ok(out)
4152    }
4153}
4154
4155impl HybridModel {
4156    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
4157    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
4158    ///
4159    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
4160    /// different 860160-byte block than the same expert of layer 7).
4161    ///
4162    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
4163    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
4164    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
4165    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
4166    pub fn moe_ffn_il(
4167        &self,
4168        e: &Engine,
4169        m: &MoeWeights,
4170        z: &CudaSlice<f32>,
4171        t: usize,
4172        il: u16,
4173    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4174        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
4175    }
4176
4177    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
4178    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
4179    pub fn moe_ffn_il_prefill(
4180        &self,
4181        e: &Engine,
4182        m: &MoeWeights,
4183        z: &CudaSlice<f32>,
4184        t: usize,
4185        il: u16,
4186    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4187        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
4188    }
4189
4190    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
4191    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
4192    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
4193    pub fn moe_ffn_il_zq8(
4194        &self,
4195        e: &Engine,
4196        m: &MoeWeights,
4197        z: &CudaSlice<f32>,
4198        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4199        t: usize,
4200        il: u16,
4201    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4202        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4203    }
4204
4205    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
4206    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
4207    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
4208    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
4209    ///
4210    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
4211    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
4212    pub(crate) fn moe_ffn(
4213        e: &Engine,
4214        m: &MoeWeights,
4215        z: &CudaSlice<f32>,
4216        t: usize,
4217        cfg: &ModelConfig,
4218        il: u16,
4219        max_block: usize,
4220    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4221        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4222    }
4223
4224    #[allow(clippy::too_many_arguments)]
4225    pub(crate) fn moe_ffn_inner(
4226        e: &Engine,
4227        m: &MoeWeights,
4228        z: &CudaSlice<f32>,
4229        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4230        t: usize,
4231        cfg: &ModelConfig,
4232        il: u16,
4233        max_block: usize,
4234        prefill: bool,
4235    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4236        let worker_io = crate::spill_pread::worker_enabled();
4237        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4238        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4239            e.with_moe_cache(max_block, |cache, _| {
4240                cache.begin_forward_epoch(il, t);
4241                if worker_io {
4242                    cache.begin_worker_scope();
4243                }
4244                Ok(())
4245            })?;
4246        }
4247        if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4248            let moe = cfg.moe.as_ref().unwrap();
4249            let n_expert = moe.expert_count as usize;
4250            let n_used = moe.expert_used_count as usize;
4251            let sigmoid = cfg.sigmoid_router().unwrap();
4252            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4253            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4254            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4255        }
4256        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
4257        // current caller into this research arm; the naked default stays on the established path.
4258        if t > 1 && moe_grouped_enabled(cfg, prefill) {
4259            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4260            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
4261            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
4262            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
4263            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
4264            if std::env::var("MEMRA_MOE_GATE").is_ok() {
4265                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4266                let g_host = e.dtoh(&grouped_out)?;
4267                let s_host = e.dtoh(&seq_out)?;
4268                let g_bytes: &[u8] = unsafe {
4269                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4270                };
4271                let s_bytes: &[u8] = unsafe {
4272                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4273                };
4274                if g_bytes == s_bytes {
4275                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4276                } else {
4277                    let diffs = g_host
4278                        .iter()
4279                        .zip(s_host.iter())
4280                        .enumerate()
4281                        .filter(|(_, (a, b))| a != b)
4282                        .count();
4283                    let maxdiff = g_host
4284                        .iter()
4285                        .zip(s_host.iter())
4286                        .map(|(a, b)| (a - b).abs())
4287                        .fold(0.0f32, f32::max);
4288                    panic!(
4289                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4290                        g_host.len()
4291                    );
4292                }
4293            }
4294            return Ok(grouped_out);
4295        }
4296        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4297    }
4298
4299    fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4300        let Some(moe) = cfg.moe.as_ref() else {
4301            return false;
4302        };
4303        // Cached once per process: this predicate runs per MoE layer per decode step, and five
4304        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
4305        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4306        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4307            std::env::var("MEMRA_MOE_STATS").is_ok()
4308                || std::env::var("MEMRA_MOE_TRACE").is_ok()
4309                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4310                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4311                || std::env::var("MEMRA_MOE_GATE").is_ok()
4312        });
4313        cfg.step35.is_some()
4314            && sigmoid_router_enabled()
4315            && moe_dev_enabled()
4316            && moe_slab_enabled()
4317            && !observation_mode
4318            && moe.expert_used_count <= 8
4319            && m.has_uniform_expert_layout()
4320            && m.gate_exps.macros.is_none()
4321            && m.up_exps.macros.is_none()
4322            && m.down_exps.macros.is_none()
4323            && !m.has_macros
4324            && moe_q8_enabled()
4325            && q8_expert_supported(m.gate_exps.qtype)
4326            && q8_expert_supported(m.up_exps.qtype)
4327            && q8_expert_supported(m.down_exps.qtype)
4328            && m.dev_exps
4329                .as_ref()
4330                .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4331    }
4332
4333    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
4334    pub(crate) fn moe_ffn_sequential(
4335        e: &Engine,
4336        m: &MoeWeights,
4337        z: &CudaSlice<f32>,
4338        t: usize,
4339        cfg: &ModelConfig,
4340        il: u16,
4341        max_block: usize,
4342    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4343        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4344    }
4345
4346    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
4347    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
4348    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
4349    fn moe_router_logits(
4350        e: &Engine,
4351        m: &MoeWeights,
4352        z: &CudaSlice<f32>,
4353        t: usize,
4354        cfg: &ModelConfig,
4355    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4356        if t < PRIME_MIN_T {
4357            // Decode and speculative verify use one fixed per-row reduction program.
4358            if crate::router_kernel_on() {
4359                e.router_gemv(
4360                    m.gate_inp.float_data(),
4361                    z,
4362                    cfg.n_embd as usize,
4363                    m.gate_exps.n_expert,
4364                    t,
4365                )
4366            } else {
4367                e.matmul_decode_exact(&m.gate_inp, z, t)
4368            }
4369        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4370            e.router_gemv(
4371                m.gate_inp.float_data(),
4372                z,
4373                cfg.n_embd as usize,
4374                m.gate_exps.n_expert,
4375                t,
4376            )
4377        } else {
4378            e.matmul(&m.gate_inp, z, t)
4379        }
4380    }
4381
4382    /// Append the host-visible router selection for one layer/forward when calibration tracing is
4383    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
4384    /// trace is independent of the dispatch optimization selected for the forward.
4385    fn trace_moe_routes(
4386        il: u16,
4387        t: usize,
4388        sel_all: &[u32],
4389        weights: &[f32],
4390    ) -> Result<(), Box<dyn std::error::Error>> {
4391        use std::io::Write as _;
4392        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4393            let mut f = std::fs::OpenOptions::new()
4394                .create(true)
4395                .append(true)
4396                .open(path)?;
4397            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4398            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4399        }
4400        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4401            let mut f = std::fs::OpenOptions::new()
4402                .create(true)
4403                .append(true)
4404                .open(path)?;
4405            let pairs: Vec<String> = sel_all
4406                .iter()
4407                .zip(weights)
4408                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4409                .collect();
4410            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4411        }
4412        Ok(())
4413    }
4414
4415    #[allow(clippy::too_many_arguments)]
4416    fn trace_sigmoid_router_logits(
4417        e: &Engine,
4418        il: u16,
4419        t: usize,
4420        n_expert: usize,
4421        n_used: usize,
4422        logits: &CudaSlice<f32>,
4423        m: &MoeWeights,
4424        (scaling_factor, route_norm): (f32, bool),
4425    ) -> Result<(), Box<dyn std::error::Error>> {
4426        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4427            return Ok(());
4428        }
4429        let logits = e.dtoh(logits)?;
4430        let active: Vec<u8> = m
4431            .active_experts
4432            .as_ref()
4433            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4434            .unwrap_or_else(|| vec![1; n_expert]);
4435        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4436        crate::sigrouter_contract::capture_served_logits(
4437            il as u32,
4438            t,
4439            n_expert,
4440            n_used,
4441            scaling_factor,
4442            route_norm,
4443            &active,
4444            &bias,
4445            &logits,
4446        )?;
4447        Ok(())
4448    }
4449
4450    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
4451    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
4452    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
4453    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
4454    fn trace_moe_input(
4455        e: &Engine,
4456        il: u16,
4457        t: usize,
4458        n_embd: usize,
4459        z: &CudaSlice<f32>,
4460    ) -> Result<(), Box<dyn std::error::Error>> {
4461        use std::io::Write as _;
4462        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4463            return Ok(());
4464        };
4465        let host = e.dtoh(z)?;
4466        if host.len() != t * n_embd {
4467            return Err(format!(
4468                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4469                host.len(),
4470                t,
4471                n_embd
4472            )
4473            .into());
4474        }
4475        let bytes = unsafe {
4476            std::slice::from_raw_parts(
4477                host.as_ptr().cast::<u8>(),
4478                host.len() * std::mem::size_of::<f32>(),
4479            )
4480        };
4481        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4482        let mut state = state
4483            .lock()
4484            .map_err(|_| "MoE input trace writer lock is poisoned")?;
4485        if state.is_none() {
4486            let dir = std::path::PathBuf::from(&dir);
4487            std::fs::create_dir_all(&dir)?;
4488            let index = std::fs::OpenOptions::new()
4489                .create(true)
4490                .append(true)
4491                .open(dir.join("index.jsonl"))?;
4492            *state = Some(MoeInputTraceWriter {
4493                dir,
4494                index,
4495                payloads: std::collections::HashMap::new(),
4496            });
4497        }
4498        let writer = state.as_mut().unwrap();
4499        if writer.dir != std::path::Path::new(&dir) {
4500            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4501        }
4502        let file_name = format!("layer-{il:03}.f32");
4503        if !writer.payloads.contains_key(&il) {
4504            let payload = std::fs::OpenOptions::new()
4505                .create(true)
4506                .append(true)
4507                .open(writer.dir.join(&file_name))?;
4508            let offset = payload.metadata()?.len();
4509            writer.payloads.insert(il, (payload, offset));
4510        }
4511        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4512        let row_offset = *offset;
4513        payload.write_all(bytes)?;
4514        *offset += bytes.len() as u64;
4515        writeln!(
4516            writer.index,
4517            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4518             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4519             \"payload_bytes\":{}}}",
4520            bytes.len()
4521        )?;
4522        Ok(())
4523    }
4524
4525    #[allow(clippy::too_many_arguments)]
4526    pub(crate) fn moe_ffn_sequential_zq8(
4527        e: &Engine,
4528        m: &MoeWeights,
4529        z: &CudaSlice<f32>,
4530        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4531        t: usize,
4532        cfg: &ModelConfig,
4533        il: u16,
4534        max_block: usize,
4535    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4536        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4537        let moe = cfg.moe.as_ref().unwrap();
4538        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
4539        let n_expert = moe.expert_count as usize; // 256
4540        let n_used = moe.expert_used_count as usize; // 8
4541        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
4542
4543        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
4544        debug_assert_eq!(m.gate_exps.in_f, n_embd);
4545        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4546        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
4547        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
4548        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
4549
4550        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
4551        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
4552        let lim_exp = cfg.clamp_exp_at(il as u32);
4553        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4554        let use_cache = Engine::moe_cache_enabled();
4555        let uniform_experts = m.has_uniform_expert_layout();
4556        let moe_q8 = uniform_experts
4557            && moe_q8_enabled()
4558            && q8_expert_supported(m.gate_exps.qtype)
4559            && q8_expert_supported(m.up_exps.qtype)
4560            && q8_expert_supported(m.down_exps.qtype);
4561        // Experimental secondary backend: complete experts already resident in the SLRU stay on
4562        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
4563        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
4564        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
4565        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
4566        // commands and CI have no llama.cpp or OpenMP dependency.
4567        let cpu_expert_requested = crate::cpu_experts::configured();
4568        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4569            return Err(std::io::Error::other(
4570                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4571            )
4572            .into());
4573        }
4574        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4575        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
4576        // Those backends are each deterministic but are different numeric configurations, so a
4577        // later prefill eviction can change greedy output. Freeze after the first real prefill;
4578        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
4579        // staging below and cannot change backend assignment.
4580        let freeze_cpu_residency = cpu_expert_requested
4581            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4582        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4583            .ok()
4584            .and_then(|value| value.parse::<usize>().ok())
4585            .is_some_and(|tokens| tokens > 0);
4586        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4587            e.freeze_moe_cache();
4588        }
4589        let cache_frozen = use_cache && e.moe_cache_frozen();
4590        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4591
4592        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
4593        // cannot change logits, selected expert ids, or routing weights.
4594        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4595        if let Some(sig) = cfg.sigmoid_router() {
4596            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4597        }
4598
4599        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
4600        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
4601        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
4602        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
4603        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
4604        // per-token host stall that dominated the 35B decode wall after stages 1+2.
4605        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
4606        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
4607        // only difference is where sel/w/pointers are READ from (device instead of params).
4608        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
4609        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
4610        // Any non-resident layer falls through to host routing + the gdec/sequential path.
4611        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
4612        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
4613        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
4614        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
4615        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
4616        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
4617        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
4618        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
4619        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
4620        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
4621        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
4622        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
4623        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
4624        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
4625        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
4626        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
4627        // now rides the dev loop below (same kernels per token as decode); pairs serves real
4628        // prefill (t >= 16, where spec never verifies).
4629        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
4630        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
4631        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
4632        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
4633        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
4634        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
4635        // ride the macro-aware sequential/staged paths below or every expert output is off by
4636        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
4637        let no_exp_macros = m.gate_exps.macros.is_none()
4638            && m.up_exps.macros.is_none()
4639            && m.down_exps.macros.is_none();
4640        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
4641        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
4642        // so it cannot even see the per-layer limit.
4643        if cfg.sigmoid_router().is_none()
4644            && cfg.m3.is_none()
4645            && cfg.hy3.is_none()
4646            && !cfg.swiglu_clamped_at(il as u32)
4647            && no_exp_macros
4648            && t >= PRIME_MIN_T
4649            && m.dev_exps.is_some()
4650            && moe_q8_enabled()
4651            && q8_expert_supported(m.gate_exps.qtype)
4652            && q8_expert_supported(m.up_exps.qtype)
4653            && q8_expert_supported(m.down_exps.qtype)
4654            && std::env::var("MEMRA_MOE_PAIRS")
4655                .map(|v| v != "0")
4656                .unwrap_or(true)
4657            && std::env::var("MEMRA_MOE_STATS").is_err()
4658        {
4659            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4660        }
4661
4662        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
4663        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
4664        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
4665        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
4666        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
4667        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
4668        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
4669        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
4670        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
4671        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
4672        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
4673        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
4674        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
4675        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
4676        // Keyed off sigmoid_router() so arch #4 is denied by construction.
4677        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
4678        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
4679        let dev_ok = uniform_experts
4680            && cfg.sigmoid_router().is_none()
4681            && cfg.m3.is_none()
4682            && cfg.hy3.is_none()
4683            && !cfg.swiglu_clamped_at(il as u32);
4684        // Observation modes must route through the host-visible selection below. Otherwise a fully
4685        // resident layer returns through device dispatch before its trace/stats row is recorded,
4686        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
4687        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4688            || std::env::var("MEMRA_MOE_TRACE").is_ok()
4689            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4690            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4691        if dev_ok
4692            && t < PRIME_MIN_T
4693            && m.dev_exps.is_some()
4694            && n_used <= 8
4695            && moe_dev_enabled()
4696            && !observe_routes
4697        {
4698            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4699        }
4700        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4701            let row_ok = e.with_moe_cache(max_block, |c, eng| {
4702                if moe_prewarm_enabled() {
4703                    c.prewarm_layer(il, m, eng)?;
4704                }
4705                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4706            })?;
4707            if row_ok {
4708                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4709            }
4710        }
4711
4712        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
4713        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4714            if cpu_hybrid {
4715                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4716                    e,
4717                    &logits,
4718                    z,
4719                    t,
4720                    n_expert,
4721                    n_used,
4722                    m.exp_probs_b.as_deref(),
4723                    sig,
4724                    m.active_experts.as_deref(),
4725                )?;
4726                (sel, w, Some(input))
4727            } else {
4728                let (sel, w) =
4729                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4730                (sel, w, None)
4731            }
4732        } else {
4733            let (sel, w) =
4734                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4735            (sel, w, None)
4736        };
4737        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4738
4739        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
4740        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
4741        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
4742        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4743        Self::trace_moe_input(e, il, t, n_embd, z)?;
4744
4745        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
4746        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
4747        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
4748        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
4749        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
4750        // wait for each pending block, so later copies can overlap the earlier expert kernels while
4751        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
4752        // T=1; batched forwards can have token-local consumers still in flight between selections.
4753        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
4754        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
4755        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
4756        let worker_disk_prefetch =
4757            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4758        let promote_worker_h2d =
4759            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4760        if promote_worker_h2d {
4761            let mut selected_blocks = Vec::with_capacity(n_used * 3);
4762            for &ex in sel_all.iter().take(n_used) {
4763                let ex = ex as u16;
4764                selected_blocks.extend([
4765                    BlockId::new(il, PROJ_GATE, ex),
4766                    BlockId::new(il, PROJ_UP, ex),
4767                    BlockId::new(il, PROJ_DOWN, ex),
4768                ]);
4769            }
4770            for &ex in sel_all.iter().take(n_used) {
4771                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4772            }
4773            e.with_moe_cache(max_block, |cache, eng| {
4774                cache.promote_worker_reads_at_safe_boundary(
4775                    &selected_blocks,
4776                    &selected_blocks,
4777                    eng,
4778                )?;
4779                Ok(())
4780            })?;
4781        }
4782
4783        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
4784        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
4785        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4786            let mut cnt = vec![0u32; n_expert];
4787            for &s in sel_all.iter() {
4788                cnt[s as usize] += 1;
4789            }
4790            let total = sel_all.len() as f64;
4791            let mut h = 0.0f64;
4792            let mut active = 0usize;
4793            for &c in &cnt {
4794                if c > 0 {
4795                    active += 1;
4796                    let p = c as f64 / total;
4797                    h -= p * p.log2();
4798                }
4799            }
4800            let maxc = cnt.iter().copied().max().unwrap_or(0);
4801            println!(
4802                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4803                il,
4804                t,
4805                sel_all.len(),
4806                active,
4807                n_expert,
4808                h,
4809                (n_expert as f64).log2(),
4810                total / active.max(1) as f64,
4811                maxc
4812            );
4813        }
4814
4815        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
4816        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
4817        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
4818        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
4819        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
4820        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
4821        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
4822        // zeroed-then-accumulated exactly as before (fallback).
4823        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
4824        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
4825        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
4826        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
4827        let gdec_may_fire = uniform_experts
4828            && use_cache
4829            && n_used <= 8
4830            && gdec_enabled()
4831            && !cfg.swiglu_clamped_at(il as u32);
4832        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
4833        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
4834        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
4835        // archs the slabs were uploaded but never read, and every expert went through the
4836        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
4837        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
4838        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
4839        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
4840        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
4841        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
4842        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
4843        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
4844        // strictly worse than staging); under PP-2 without the prime walker this admits
4845        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
4846        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
4847        let slab_local = m
4848            .dev_exps
4849            .as_ref()
4850            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4851        let slab_bases = slab_local.map(|d| {
4852            use cudarc::driver::DevicePtr;
4853            let s = e.stream();
4854            let (pg, _g0) = d.gate.device_ptr(&s);
4855            let (pu, _g1) = d.up.device_ptr(&s);
4856            let (pd, _g2) = d.down.device_ptr(&s);
4857            (pg as u64, pu as u64, pd as u64)
4858        });
4859        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
4860        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
4861        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
4862        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
4863        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
4864        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
4865        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
4866        // all-resident tokens, staged loop for misses), which is a dispatch-class
4867        // comparison, not a provenance one.
4868        let slab_fused_may_fire = slab_bases.is_some()
4869            && n_used <= 8
4870            && gdec_enabled()
4871            && !cfg.swiglu_clamped_at(il as u32)
4872            && cfg.m3.is_none()
4873            && no_exp_macros
4874            && moe_q8;
4875        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
4876        // uninit; a token that falls through to any accumulating loop zeroes its own row.
4877        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4878            e.uninit(t * n_embd)?
4879        } else {
4880            e.zeros(t * n_embd)?
4881        };
4882        // The router readback above already established a host boundary. Copy each small-t hidden
4883        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
4884        let cpu_input = if cpu_hybrid {
4885            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4886        } else {
4887            None
4888        };
4889
4890        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
4891        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
4892        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
4893        // measured ~123 memsets/token of the decode wall).
4894        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
4895        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
4896        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
4897        let mut scratch_g: Option<CudaSlice<u8>> = None;
4898        let mut scratch_u: Option<CudaSlice<u8>> = None;
4899        let mut scratch_d: Option<CudaSlice<u8>> = None;
4900        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
4901        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
4902
4903        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
4904        // the copy stream before launching the current expert's compute. Pending slots stay invisible
4905        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
4906        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
4907        let page_window = moe_page_prefetch_window();
4908
4909        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
4910        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
4911        for tok in 0..t {
4912            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4913            let w = &w_all[tok * n_used..(tok + 1) * n_used];
4914            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
4915            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4916
4917            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
4918            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
4919            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
4920            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
4921            // memcpy, zero admission, so no slot can move under the collected pointers) — any
4922            // miss falls through to the sequential loop below, which admits as before. In steady
4923            // state on a fully-resident rig every token-layer takes the grouped path.
4924            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
4925            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
4926            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
4927            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
4928            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
4929            // per-expert macro-scales the fused kernels don't fold — those fall through too.
4930            let no_macros = m.gate_exps.macros.is_none()
4931                && m.up_exps.macros.is_none()
4932                && m.down_exps.macros.is_none();
4933            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
4934            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
4935            // with pointers computed from the resident slab base + ex*stride instead of
4936            // collected SLRU slot addresses. No cache lock, no residency predicate — the
4937            // slab holds every expert by construction, so this arm never falls through
4938            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
4939            // staging both die). Bit-identity class: pointer provenance only, the same
4940            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
4941            // slab exists it is strictly better (no lock, no miss).
4942            if slab_fused_may_fire {
4943                let (pg, pu, pd) = slab_bases.unwrap();
4944                let mut gp = [0u64; 8];
4945                let mut up = [0u64; 8];
4946                let mut dp = [0u64; 8];
4947                for (j, &ex) in sel.iter().enumerate() {
4948                    let ex = ex as usize;
4949                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4950                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4951                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4952                }
4953                let mut wv = [0f32; 8];
4954                wv[..n_used].copy_from_slice(w);
4955                if tok_q8.is_none() {
4956                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4957                }
4958                let (zq, zd) = tok_q8.as_ref().unwrap();
4959                let act = e.moe_gate_up_silu8_q8(
4960                    crate::WPtr8(gp),
4961                    crate::WPtr8(up),
4962                    zq,
4963                    zd,
4964                    n_embd,
4965                    n_ff_exp,
4966                    n_used,
4967                    m.gate_exps.qtype,
4968                    m.up_exps.qtype,
4969                    m.gate_exps.row_bytes,
4970                    m.up_exps.row_bytes,
4971                )?;
4972                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4973                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4974                e.moe_down8_fma_q8(
4975                    crate::WPtr8(dp),
4976                    crate::F32x8(wv),
4977                    &aq2,
4978                    &ad2,
4979                    &mut dst,
4980                    n_ff_exp,
4981                    n_embd,
4982                    n_used,
4983                    m.down_exps.qtype,
4984                    m.down_exps.row_bytes,
4985                )?;
4986                continue;
4987            }
4988            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4989                if tok_q8.is_none() {
4990                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4991                }
4992                let (zq, zd) = tok_q8.as_ref().unwrap();
4993                if Self::moe_gdec_token_q8(
4994                    e,
4995                    m,
4996                    il,
4997                    max_block,
4998                    zq,
4999                    zd,
5000                    sel,
5001                    w,
5002                    &mut moe_out,
5003                    tok,
5004                    n_embd,
5005                    n_ff_exp,
5006                    n_used,
5007                )? {
5008                    continue;
5009                }
5010            } else if gdec_may_fire
5011                && cfg.m3.is_none()
5012                && no_macros
5013                && Self::moe_gdec_token(
5014                    e,
5015                    m,
5016                    il,
5017                    max_block,
5018                    &zt,
5019                    sel,
5020                    w,
5021                    &mut moe_out,
5022                    tok,
5023                    n_embd,
5024                    n_ff_exp,
5025                    n_used,
5026                )?
5027            {
5028                continue;
5029            }
5030
5031            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
5032            // slab pair could fire. This token fell through to a sequential axpy loop, which
5033            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
5034            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
5035            // has no fallible predicate), included for the allocation invariant's symmetry.
5036            if gdec_may_fire || slab_fused_may_fire {
5037                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5038                e.memset_zeros_view(&mut row)?;
5039            }
5040
5041            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
5042            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
5043            // stall this path exists to remove, while mixing projections would require another
5044            // activation round-trip. Weight addresses remain valid until this worker is joined at
5045            // the bottom of the token scope.
5046            let mut cpu_mask = vec![false; sel.len()];
5047            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
5048                let gpu_resident = if use_cache {
5049                    e.with_moe_cache(max_block, |cache, _| {
5050                        Ok(sel
5051                            .iter()
5052                            .map(|&expert| {
5053                                let expert = expert as u16;
5054                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
5055                                    .into_iter()
5056                                    .filter(|&projection| {
5057                                        cache
5058                                            .resident(BlockId::new(il, projection, expert))
5059                                            .is_some()
5060                                    })
5061                                    .count()
5062                            })
5063                            .collect::<Vec<_>>())
5064                    })?
5065                } else {
5066                    vec![0; sel.len()]
5067                };
5068                let mut cpu_selected = Vec::new();
5069                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
5070                    if gpu_resident[index] != 3 {
5071                        cpu_mask[index] = true;
5072                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
5073                        let expert = expert as usize;
5074                        cpu_selected.push((expert, route_weight));
5075                    }
5076                }
5077                if crate::cpu_experts::predictor_enabled() {
5078                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
5079                    // from this layer's MoE input and prefetches predicted-and-missing
5080                    // experts into the companion RAM cache. Never blocks this thread.
5081                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
5082                    crate::cpu_experts::predictor_submit(il, row);
5083                }
5084                if cpu_selected.is_empty() {
5085                    None
5086                } else {
5087                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
5088                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
5089                        .map_err(std::io::Error::other)?;
5090                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
5091                }
5092            } else {
5093                None
5094            };
5095
5096            let worker_window = worker_disk_prefetch
5097                .then(worker_prefetch_window)
5098                .unwrap_or(0);
5099            for (j, &ex) in sel.iter().enumerate() {
5100                if cpu_mask[j] {
5101                    continue;
5102                }
5103                let ex = ex as usize;
5104                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
5105                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
5106                // fused form) and macro-carrying artifacts — still have their bytes in the
5107                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
5108                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
5109                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
5110                if let Some(d) = slab_local {
5111                    let gl = m.gate_exps.expert_layout(ex);
5112                    let ul = m.up_exps.expert_layout(ex);
5113                    let dl = m.down_exps.expert_layout(ex);
5114                    let (g0, u0, d0) = (
5115                        ex * m.gate_exps.expert_stride,
5116                        ex * m.up_exps.expert_stride,
5117                        ex * m.down_exps.expert_stride,
5118                    );
5119                    let (gate, up) = if moe_q8 {
5120                        if tok_q8.is_none() {
5121                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5122                        }
5123                        let (zq, zd) = tok_q8.as_ref().unwrap();
5124                        (
5125                            e.qmatvec_expert_q8(
5126                                &d.gate,
5127                                g0..g0 + gl.len,
5128                                zq,
5129                                zd,
5130                                1,
5131                                m.gate_exps.in_f,
5132                                m.gate_exps.out_f,
5133                                gl.qtype,
5134                                gl.row_bytes,
5135                            )?,
5136                            e.qmatvec_expert_q8(
5137                                &d.up,
5138                                u0..u0 + ul.len,
5139                                zq,
5140                                zd,
5141                                1,
5142                                m.up_exps.in_f,
5143                                m.up_exps.out_f,
5144                                ul.qtype,
5145                                ul.row_bytes,
5146                            )?,
5147                        )
5148                    } else {
5149                        (
5150                            e.qmatvec_view(
5151                                &d.gate,
5152                                g0..g0 + gl.len,
5153                                &zt,
5154                                1,
5155                                m.gate_exps.in_f,
5156                                m.gate_exps.out_f,
5157                                gl.qtype,
5158                                gl.row_bytes,
5159                            )?,
5160                            e.qmatvec_view(
5161                                &d.up,
5162                                u0..u0 + ul.len,
5163                                &zt,
5164                                1,
5165                                m.up_exps.in_f,
5166                                m.up_exps.out_f,
5167                                ul.qtype,
5168                                ul.row_bytes,
5169                            )?,
5170                        )
5171                    };
5172                    let mut act = e.uninit(n_ff_exp)?;
5173                    Self::ffn_act_lim(
5174                        e,
5175                        cfg,
5176                        &gate,
5177                        &up,
5178                        m.gate_exps.macro_scale(ex),
5179                        m.up_exps.macro_scale(ex),
5180                        lim_exp,
5181                        &mut act,
5182                        n_ff_exp,
5183                    )?;
5184                    let y = if moe_q8 {
5185                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5186                        e.qmatvec_expert_q8(
5187                            &d.down,
5188                            d0..d0 + dl.len,
5189                            &aq2,
5190                            &ad2,
5191                            1,
5192                            m.down_exps.in_f,
5193                            m.down_exps.out_f,
5194                            dl.qtype,
5195                            dl.row_bytes,
5196                        )?
5197                    } else {
5198                        let actv = act.slice(0..n_ff_exp);
5199                        e.qmatvec_view(
5200                            &d.down,
5201                            d0..d0 + dl.len,
5202                            &actv,
5203                            1,
5204                            m.down_exps.in_f,
5205                            m.down_exps.out_f,
5206                            dl.qtype,
5207                            dl.row_bytes,
5208                        )?
5209                    };
5210                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5211                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5212                    continue;
5213                }
5214                for next in page_prefetch_positions(j, sel.len(), page_window) {
5215                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
5216                }
5217                let keep = [
5218                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5219                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5220                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5221                ];
5222                if worker_disk_prefetch && worker_window > 0 {
5223                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5224                        Self::moe_prefetch_disk_expert(
5225                            e,
5226                            il,
5227                            sel[next] as usize,
5228                            m,
5229                            max_block,
5230                            &keep,
5231                        )?;
5232                    }
5233                } else if cache_dispatch
5234                    && !cpu_hybrid
5235                    && moe_prefetch_enabled()
5236                    && j + 1 < sel.len()
5237                {
5238                    let next = sel[j + 1] as usize;
5239                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5240                }
5241                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5242                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5243                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
5244                    // layouts stay on the metadata-aware f32 path.
5245                    if (gate_q8 || up_q8) && tok_q8.is_none() {
5246                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5247                    }
5248                    let gate = if gate_q8 {
5249                        let (zq, zd) = tok_q8.as_ref().unwrap();
5250                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5251                    } else {
5252                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5253                    };
5254                    let up = if up_q8 {
5255                        let (zq, zd) = tok_q8.as_ref().unwrap();
5256                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5257                    } else {
5258                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5259                    };
5260                    let mut act = e.uninit(n_ff_exp)?;
5261                    Self::ffn_act_lim(
5262                        e,
5263                        cfg,
5264                        &gate,
5265                        &up,
5266                        m.gate_exps.macro_scale(ex),
5267                        m.up_exps.macro_scale(ex),
5268                        lim_exp,
5269                        &mut act,
5270                        n_ff_exp,
5271                    )?;
5272                    let y = if down_q8 {
5273                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5274                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5275                    } else {
5276                        let actv = act.slice(0..n_ff_exp);
5277                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5278                    };
5279                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5280                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
5281                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5282                } else if cache_dispatch {
5283                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
5284                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
5285                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
5286                    // only difference between HIT and MISS is whether the memcpy_htod ran.
5287                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5288                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5289                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5290                    Self::ffn_act_lim(
5291                        e,
5292                        cfg,
5293                        &gate,
5294                        &up,
5295                        m.gate_exps.macro_scale(ex),
5296                        m.up_exps.macro_scale(ex),
5297                        lim_exp,
5298                        &mut act,
5299                        n_ff_exp,
5300                    )?;
5301                    let actv = act.slice(0..n_ff_exp);
5302                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5303                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5304                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
5305                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5306                } else if cache_frozen {
5307                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
5308                    // first prime. Reuse every fixed resident projection directly and stage only a
5309                    // true miss through the ordinary scratch slot. This preserves the established
5310                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
5311                    let gate = Self::moe_frozen_gemm(
5312                        e,
5313                        il,
5314                        PROJ_GATE,
5315                        ex,
5316                        m,
5317                        max_block,
5318                        &zt,
5319                        &mut scratch_g,
5320                        g_len,
5321                    )?;
5322                    let up = Self::moe_frozen_gemm(
5323                        e,
5324                        il,
5325                        PROJ_UP,
5326                        ex,
5327                        m,
5328                        max_block,
5329                        &zt,
5330                        &mut scratch_u,
5331                        u_len,
5332                    )?;
5333                    let mut act = e.uninit(n_ff_exp)?;
5334                    Self::ffn_act_lim(
5335                        e,
5336                        cfg,
5337                        &gate,
5338                        &up,
5339                        m.gate_exps.macro_scale(ex),
5340                        m.up_exps.macro_scale(ex),
5341                        lim_exp,
5342                        &mut act,
5343                        n_ff_exp,
5344                    )?;
5345                    let actv = act.slice(0..n_ff_exp);
5346                    let y = Self::moe_frozen_gemm(
5347                        e,
5348                        il,
5349                        PROJ_DOWN,
5350                        ex,
5351                        m,
5352                        max_block,
5353                        &actv,
5354                        &mut scratch_d,
5355                        d_len,
5356                    )?;
5357                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5358                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5359                } else {
5360                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
5361                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
5362                    // fully overwrites the byte range the GEMM reads).
5363                    if scratch_g.is_none() {
5364                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5365                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5366                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5367                    }
5368                    let (sg, su, sd) = (
5369                        scratch_g.as_mut().unwrap(),
5370                        scratch_u.as_mut().unwrap(),
5371                        scratch_d.as_mut().unwrap(),
5372                    );
5373                    let gl = m.gate_exps.expert_layout(ex);
5374                    let ul = m.up_exps.expert_layout(ex);
5375                    let dl = m.down_exps.expert_layout(ex);
5376                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5377                    let gate = e.qmatvec_view(
5378                        sg,
5379                        0..gl.len,
5380                        &zt,
5381                        1,
5382                        m.gate_exps.in_f,
5383                        m.gate_exps.out_f,
5384                        gl.qtype,
5385                        gl.row_bytes,
5386                    )?;
5387
5388                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5389                    let up = e.qmatvec_view(
5390                        su,
5391                        0..ul.len,
5392                        &zt,
5393                        1,
5394                        m.up_exps.in_f,
5395                        m.up_exps.out_f,
5396                        ul.qtype,
5397                        ul.row_bytes,
5398                    )?;
5399
5400                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5401                    Self::ffn_act_lim(
5402                        e,
5403                        cfg,
5404                        &gate,
5405                        &up,
5406                        m.gate_exps.macro_scale(ex),
5407                        m.up_exps.macro_scale(ex),
5408                        lim_exp,
5409                        &mut act,
5410                        n_ff_exp,
5411                    )?;
5412
5413                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5414                    let actv = act.slice(0..n_ff_exp);
5415                    let y = e.qmatvec_view(
5416                        sd,
5417                        0..dl.len,
5418                        &actv,
5419                        1,
5420                        m.down_exps.in_f,
5421                        m.down_exps.out_f,
5422                        dl.qtype,
5423                        dl.row_bytes,
5424                    )?;
5425
5426                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5427                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5428                }
5429            }
5430            if let Some(worker) = cpu_worker {
5431                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5432                let cpu_output = e.htod(&cpu_output)?;
5433                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5434                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5435            }
5436            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5437                for (j, &ex) in sel.iter().enumerate() {
5438                    if cpu_mask[j] {
5439                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5440                    }
5441                }
5442            }
5443        }
5444
5445        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
5446        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
5447        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5448        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5449        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5450            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5451        {
5452            let n_ff_sh = gate_shexp.out_features(); // 512
5453            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
5454            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
5455            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
5456            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
5457            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
5458            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
5459            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
5460            let verify_t = t > 1 && t < PRIME_MIN_T;
5461            let (sg_gate, sg_up) = if t == 1 {
5462                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5463                    Some(pair) => pair,
5464                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5465                }
5466            } else if verify_t {
5467                (
5468                    e.matmul_decode_exact(gate_shexp, z, t)?,
5469                    e.matmul_decode_exact(up_shexp, z, t)?,
5470                )
5471            } else {
5472                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
5473            };
5474            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
5475            Self::ffn_act_lim(
5476                e,
5477                cfg,
5478                &sg_gate,
5479                &sg_up,
5480                1.0,
5481                1.0,
5482                lim_shexp,
5483                &mut sa,
5484                t * n_ff_sh,
5485            )?;
5486            let sh = if verify_t {
5487                e.matmul_decode_exact(down_shexp, &sa, t)?
5488            } else {
5489                e.matmul(down_shexp, &sa, t)?
5490            }; // [T, n_embd]
5491
5492            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
5493            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
5494            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
5495            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
5496            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
5497            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
5498            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
5499            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
5500            // expert's contribution into every token's residual, so under cross-request
5501            // concat prefill a session's hidden state depended on its co-arrivals' token
5502            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
5503            let g = match &m.gate_inp_shexp {
5504                Some(gate_inp_shexp) => {
5505                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5506                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5507                    } else {
5508                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5509                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
5510                        e.sigmoid(&gs, &mut g, t)?;
5511                        g
5512                    }
5513                }
5514                None => e.htod(&vec![1.0f32; t])?,
5515            };
5516            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
5517            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5518        }
5519
5520        Ok(moe_out)
5521    }
5522
5523    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
5524    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
5525    pub fn stage1_h2d_per_token(&self) -> u64 {
5526        use crate::hybrid::Ffn;
5527        let n_used = self
5528            .cfg
5529            .moe
5530            .as_ref()
5531            .map(|m| m.expert_used_count as u64)
5532            .unwrap_or(0);
5533        let mut bytes = 0u64;
5534        for l in self.layers.iter() {
5535            if let Ffn::Moe(m) = &l.ffn {
5536                bytes += n_used
5537                    * (m.gate_exps.max_expert_bytes()
5538                        + m.up_exps.max_expert_bytes()
5539                        + m.down_exps.max_expert_bytes()) as u64;
5540            }
5541        }
5542        bytes
5543    }
5544
5545    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
5546    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
5547    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
5548    pub(crate) fn max_moe_block(&self) -> usize {
5549        use crate::hybrid::Ffn;
5550        let mut mx = 0usize;
5551        let mut scan = |ffn: &Ffn| {
5552            if let Ffn::Moe(m) = ffn {
5553                mx = mx
5554                    .max(m.gate_exps.max_expert_bytes())
5555                    .max(m.up_exps.max_expert_bytes())
5556                    .max(m.down_exps.max_expert_bytes());
5557            }
5558        };
5559        for l in self.layers.iter() {
5560            scan(&l.ffn);
5561        }
5562        if let Some(mtp) = self.mtp.as_ref() {
5563            scan(&mtp.ffn);
5564        }
5565        mx
5566    }
5567
5568    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
5569    /// but have no bytes and therefore consume no residency slot.
5570    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5571        use crate::hybrid::Ffn;
5572        let mut sizes = Vec::new();
5573        let mut scan = |ffn: &Ffn| {
5574            let Ffn::Moe(m) = ffn else { return };
5575            for ex in 0..m.gate_exps.n_expert {
5576                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5577                    continue;
5578                }
5579                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5580                    let len = exps.expert_layout(ex).len;
5581                    if len > 0 {
5582                        sizes.push(len);
5583                    }
5584                }
5585            }
5586        };
5587        for layer in &self.layers {
5588            scan(&layer.ffn);
5589        }
5590        if let Some(mtp) = &self.mtp {
5591            scan(&mtp.ffn);
5592        }
5593        sizes
5594    }
5595
5596    /// Persist the frozen residency set so a later process can restage it directly and skip
5597    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
5598    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
5599    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
5600    /// post-freeze argmax gate still validates the serving assignment.
5601    pub fn save_cpu_expert_residency_profile(
5602        &self,
5603        e: &Engine,
5604        path: &std::path::Path,
5605    ) -> Result<(), Box<dyn std::error::Error>> {
5606        let Some(ids) = e.export_moe_residency() else {
5607            return Err("no MoE residency cache to persist".into());
5608        };
5609        let mut body = format!(
5610            "memra-freeze-profile v1 max_block={} blocks={}\n",
5611            self.max_moe_block(),
5612            ids.len()
5613        );
5614        for (layer, proj, ex) in &ids {
5615            body.push_str(&format!("{layer} {proj} {ex}\n"));
5616        }
5617        let tmp = path.with_extension("tmp");
5618        std::fs::write(&tmp, body)?;
5619        std::fs::rename(&tmp, path)?;
5620        println!(
5621            "[moe-cache] freeze profile saved: {} blocks -> {}",
5622            ids.len(),
5623            path.display()
5624        );
5625        Ok(())
5626    }
5627
5628    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
5629    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
5630    /// missing or its header does not match this model's slot geometry.
5631    pub fn restore_cpu_expert_residency_profile(
5632        &self,
5633        e: &Engine,
5634        path: &std::path::Path,
5635    ) -> Result<bool, Box<dyn std::error::Error>> {
5636        use crate::hybrid::Ffn;
5637        use crate::moe_cache::BlockId;
5638        let Ok(content) = std::fs::read_to_string(path) else {
5639            return Ok(false);
5640        };
5641        let mut lines = content.lines();
5642        let Some(header) = lines.next() else {
5643            return Ok(false);
5644        };
5645        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5646        if !header.starts_with(&expected) {
5647            println!(
5648                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5649                path.display()
5650            );
5651            return Ok(false);
5652        }
5653        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5654            std::collections::HashMap::new();
5655        for line in lines {
5656            let mut fields = line.split_whitespace();
5657            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5658            else {
5659                continue;
5660            };
5661            let (Ok(layer), Ok(proj), Ok(ex)) =
5662                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5663            else {
5664                continue;
5665            };
5666            by_layer
5667                .entry(layer)
5668                .or_default()
5669                .push(BlockId::new(layer, proj, ex));
5670        }
5671        let requested: usize = by_layer.values().map(Vec::len).sum();
5672        if requested == 0 {
5673            return Ok(false);
5674        }
5675        let max_block = self.max_moe_block();
5676        let mut restaged = 0usize;
5677        let mut stage_layer =
5678            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5679                let Ffn::Moe(m) = ffn else { return Ok(()) };
5680                let Some(ids) = by_layer.get(&layer_index) else {
5681                    return Ok(());
5682                };
5683                e.with_moe_cache(max_block, |cache, eng| {
5684                    for id in ids {
5685                        if cache.restage_block(*id, m, eng)? {
5686                            restaged += 1;
5687                        }
5688                    }
5689                    Ok(())
5690                })
5691            };
5692        for (index, layer) in self.layers.iter().enumerate() {
5693            stage_layer(index as u16, &layer.ffn)?;
5694        }
5695        if let Some(mtp) = self.mtp.as_ref() {
5696            stage_layer(u16::MAX, &mtp.ffn)?;
5697        }
5698        e.freeze_moe_cache();
5699        println!(
5700            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5701            path.display()
5702        );
5703        Ok(true)
5704    }
5705
5706    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
5707    pub fn freeze_cpu_expert_residency(
5708        &self,
5709        e: &Engine,
5710    ) -> Result<(), Box<dyn std::error::Error>> {
5711        e.freeze_moe_cache();
5712        Ok(())
5713    }
5714
5715    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
5716    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
5717    /// the model's activation exactly.
5718    ///
5719    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
5720    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
5721    /// form for anything that can land on a clamped layer.
5722    pub fn ffn_act(
5723        e: &Engine,
5724        cfg: &ModelConfig,
5725        gate: &CudaSlice<f32>,
5726        up: &CudaSlice<f32>,
5727        act: &mut CudaSlice<f32>,
5728        n: usize,
5729    ) -> Result<(), Box<dyn std::error::Error>> {
5730        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5731    }
5732
5733    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
5734    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
5735    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
5736    #[allow(clippy::too_many_arguments)]
5737    pub(crate) fn ffn_act_scaled(
5738        e: &Engine,
5739        cfg: &ModelConfig,
5740        gate: &CudaSlice<f32>,
5741        up: &CudaSlice<f32>,
5742        gs: f32,
5743        us: f32,
5744        act: &mut CudaSlice<f32>,
5745        n: usize,
5746    ) -> Result<(), Box<dyn std::error::Error>> {
5747        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5748    }
5749
5750    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
5751    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
5752    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
5753    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
5754    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
5755    ///                 arrays are SEPARATE and a layer can have one without the other.
5756    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
5757    /// already known live.
5758    #[allow(clippy::too_many_arguments)]
5759    pub(crate) fn ffn_act_lim(
5760        e: &Engine,
5761        cfg: &ModelConfig,
5762        gate: &CudaSlice<f32>,
5763        up: &CudaSlice<f32>,
5764        gs: f32,
5765        us: f32,
5766        limit: Option<f32>,
5767        act: &mut CudaSlice<f32>,
5768        n: usize,
5769    ) -> Result<(), Box<dyn std::error::Error>> {
5770        if let Some(m3) = cfg.m3.as_ref() {
5771            debug_assert!(
5772                limit.is_none(),
5773                "m3 swigluoai and step35 clamp are different archs"
5774            );
5775            return e.swigluoai_mul_scaled(
5776                gate,
5777                up,
5778                gs,
5779                us,
5780                m3.swiglu_alpha,
5781                m3.swiglu_limit,
5782                act,
5783                n,
5784            );
5785        }
5786        if let Some(l) = limit {
5787            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5788        }
5789        if gs == 1.0 && us == 1.0 {
5790            return e.silu_mul(gate, up, act, n);
5791        }
5792        e.silu_mul_scaled(gate, up, gs, us, act, n)
5793    }
5794
5795    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
5796    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
5797    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
5798    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
5799    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
5800    fn moe_route(
5801        e: &Engine,
5802        logits: &CudaSlice<f32>,
5803        t: usize,
5804        n_expert: usize,
5805        n_used: usize,
5806    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5807        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5808    }
5809
5810    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
5811    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
5812    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
5813    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
5814    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
5815    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
5816    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
5817    #[allow(clippy::too_many_arguments)]
5818    fn moe_route_sigmoid_cfg(
5819        e: &Engine,
5820        logits: &CudaSlice<f32>,
5821        t: usize,
5822        n_expert: usize,
5823        n_used: usize,
5824        m: &MoeWeights,
5825        (sf, route_norm): (f32, bool),
5826    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5827        if sigmoid_router_enabled() {
5828            return e.moe_router_sigmoid_topk_host(
5829                logits,
5830                t,
5831                n_expert,
5832                n_used,
5833                m.active_count(),
5834                &m.exp_probs_b_dev,
5835                &m.active_experts_dev,
5836                sf,
5837                route_norm,
5838            );
5839        }
5840        let lg = e.dtoh(logits)?;
5841        Self::moe_route_sigmoid_host(
5842            &lg,
5843            t,
5844            n_expert,
5845            n_used,
5846            m.exp_probs_b.as_deref(),
5847            sf,
5848            route_norm,
5849            m.active_experts.as_deref(),
5850        )
5851    }
5852
5853    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
5854    /// the existing softmax device kernel has no mask input.
5855    fn moe_route_cfg(
5856        e: &Engine,
5857        logits: &CudaSlice<f32>,
5858        t: usize,
5859        n_expert: usize,
5860        n_used: usize,
5861        active: Option<&[bool]>,
5862    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5863        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
5864        // rollback) via the single-sync pinned readback — softmax arch only.
5865        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5866            return e.moe_router_topk_host(logits, t, n_expert, n_used);
5867        }
5868        // Host oracle (the §D bit-identity reference).
5869        let lg = e.dtoh(logits)?; // [T*n_expert] host
5870        let mut sel = vec![0u32; t * n_used];
5871        let mut w_out = vec![0f32; t * n_used];
5872        for tok in 0..t {
5873            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5874            // softmax over ALL n_expert (stable: subtract max)
5875            let maxl = row
5876                .iter()
5877                .enumerate()
5878                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5879                .map(|(_, &x)| x)
5880                .fold(f32::NEG_INFINITY, f32::max);
5881            let mut probs = vec![0f32; n_expert];
5882            let mut den = 0f32;
5883            for i in 0..n_expert {
5884                if active.is_some_and(|mask| !mask[i]) {
5885                    continue;
5886                }
5887                let x = (row[i] - maxl).exp();
5888                probs[i] = x;
5889                den += x;
5890            }
5891            for p in probs.iter_mut() {
5892                *p /= den;
5893            }
5894            // stable DESC sort: prob DESC, ascending-index tiebreak.
5895            let mut idx: Vec<usize> = (0..n_expert)
5896                .filter(|&i| active.is_none_or(|mask| mask[i]))
5897                .collect();
5898            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5899            let sl = &idx[..n_used];
5900            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5901            let mut ws: f32 = wv.iter().sum();
5902            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
5903            for x in wv.iter_mut() {
5904                *x /= ws;
5905            }
5906            for j in 0..n_used {
5907                sel[tok * n_used + j] = sl[j] as u32;
5908                w_out[tok * n_used + j] = wv[j];
5909            }
5910        }
5911        Ok((sel, w_out))
5912    }
5913
5914    #[allow(clippy::too_many_arguments)]
5915    fn moe_route_sigmoid_with_input(
5916        e: &Engine,
5917        logits: &CudaSlice<f32>,
5918        input: &CudaSlice<f32>,
5919        t: usize,
5920        n_expert: usize,
5921        n_used: usize,
5922        bias: Option<&[f32]>,
5923        (sf, route_norm): (f32, bool),
5924        active: Option<&[bool]>,
5925    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5926        let (lg, input) = e.dtoh_pair(logits, input)?;
5927        let (sel, w) =
5928            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5929        Ok((sel, w, input))
5930    }
5931
5932    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
5933    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
5934    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
5935    /// active mask, prebuilt projection descriptors) so no model reference escapes.
5936    pub fn start_moe_prefetch_predictor(
5937        &self,
5938        e: &Engine,
5939        cfg: &ModelConfig,
5940    ) -> Result<(), Box<dyn std::error::Error>> {
5941        use crate::hybrid::Ffn;
5942        let Some(sig) = cfg.sigmoid_router() else {
5943            return Err("prefetch predictor requires a sigmoid-router arch".into());
5944        };
5945        let resident: std::collections::HashSet<(u16, u8, u16)> = e
5946            .export_moe_residency()
5947            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5948            .into_iter()
5949            .collect();
5950        let mut layers = Vec::new();
5951        for (index, layer) in self.layers.iter().enumerate() {
5952            let Ffn::Moe(m) = &layer.ffn else { continue };
5953            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5954                continue;
5955            };
5956            let router = e.dtoh(data)?;
5957            let n_expert = m.gate_exps.n_expert;
5958            let n_embd = m.gate_exps.in_f;
5959            if router.len() != n_embd * n_expert {
5960                continue;
5961            }
5962            let build = |exps: &crate::model::HostExps| {
5963                (0..n_expert)
5964                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5965                    .collect::<Vec<_>>()
5966            };
5967            layers.push((
5968                index as u16,
5969                crate::cpu_experts::PredictLayerInit {
5970                    router,
5971                    bias: m.exp_probs_b.clone(),
5972                    active: m.active_experts.clone(),
5973                    n_embd,
5974                    n_used: cfg
5975                        .moe
5976                        .as_ref()
5977                        .map(|moe| moe.expert_used_count as usize)
5978                        .ok_or("prefetch predictor requires MoE config")?,
5979                    sig,
5980                    weights_n_expert: n_expert,
5981                    gate: build(&m.gate_exps),
5982                    up: build(&m.up_exps),
5983                    down: build(&m.down_exps),
5984                },
5985            ));
5986        }
5987        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5988    }
5989
5990    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
5991    /// selection math to the rollback runtime, applied to host-computed logits.
5992    #[allow(clippy::too_many_arguments)]
5993    pub fn moe_route_sigmoid_host_public(
5994        logits: &[f32],
5995        t: usize,
5996        n_expert: usize,
5997        n_used: usize,
5998        bias: Option<&[f32]>,
5999        sf: f32,
6000        route_norm: bool,
6001        active: Option<&[bool]>,
6002    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6003        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
6004    }
6005
6006    #[allow(clippy::too_many_arguments)]
6007    fn moe_route_sigmoid_host(
6008        lg: &[f32],
6009        t: usize,
6010        n_expert: usize,
6011        n_used: usize,
6012        bias: Option<&[f32]>,
6013        sf: f32,
6014        route_norm: bool,
6015        active: Option<&[bool]>,
6016    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
6017        let active_count = active
6018            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
6019            .unwrap_or(n_expert);
6020        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
6021        if lg.len() != t * n_expert {
6022            return Err(format!(
6023                "sigmoid router logits length mismatch: got {}, expected {}",
6024                lg.len(),
6025                t * n_expert,
6026            )
6027            .into());
6028        }
6029        let mut sel = vec![0u32; t * n_used];
6030        let mut w_out = vec![0f32; t * n_used];
6031        for tok in 0..t {
6032            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
6033            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
6034            // selection score = sigmoid + bias; weight = plain sigmoid.
6035            let selsc: Vec<f32> = match bias {
6036                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
6037                None => scores.clone(),
6038            };
6039            let mut idx: Vec<usize> = (0..n_expert)
6040                .filter(|&i| active.is_none_or(|mask| mask[i]))
6041                .collect();
6042            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
6043            let sl = &idx[..n_used];
6044            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
6045            if route_norm {
6046                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
6047                for x in wv.iter_mut() {
6048                    *x = *x / ws * sf;
6049                }
6050            } else {
6051                for x in wv.iter_mut() {
6052                    *x *= sf;
6053                }
6054            }
6055            for j in 0..n_used {
6056                sel[tok * n_used + j] = sl[j] as u32;
6057                w_out[tok * n_used + j] = wv[j];
6058            }
6059        }
6060        Ok((sel, w_out))
6061    }
6062
6063    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
6064    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
6065    /// macro-scaled experts, and observation modes are denied by the caller.
6066    #[allow(clippy::too_many_arguments)]
6067    fn moe_ffn_sigmoid_dev(
6068        e: &Engine,
6069        m: &MoeWeights,
6070        z: &CudaSlice<f32>,
6071        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6072        logits: &CudaSlice<f32>,
6073        t: usize,
6074        cfg: &ModelConfig,
6075        il: u16,
6076        (scaling_factor, route_norm): (f32, bool),
6077    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6078        let moe = cfg.moe.as_ref().unwrap();
6079        let n_embd = cfg.n_embd as usize;
6080        let n_expert = moe.expert_count as usize;
6081        let n_used = moe.expert_used_count as usize;
6082        let n_ff_exp = moe.expert_ff_length as usize;
6083        let dev = m.dev_exps.as_ref().unwrap();
6084        debug_assert!(cfg.step35.is_some());
6085        debug_assert_eq!(dev.dev, e.ctx().ordinal());
6086        debug_assert!(m.has_uniform_expert_layout());
6087        debug_assert!(!m.has_macros);
6088
6089        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
6090            logits,
6091            t,
6092            n_expert,
6093            n_used,
6094            m.active_count(),
6095            &m.exp_probs_b_dev,
6096            &m.active_experts_dev,
6097            scaling_factor,
6098            route_norm,
6099        )?;
6100        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
6101        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
6102            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6103            (combined, combined)
6104        } else {
6105            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6106        };
6107        let (zq, zd) = match (t, zq8) {
6108            (1, Some((q, d))) => (q.clone(), d.clone()),
6109            _ => e.quantize_q8_1(z, t, n_embd)?,
6110        };
6111        let n_pairs = t * n_used;
6112        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
6113            // The final Step layers retain the established separate gate/up -> clamp -> down
6114            // arithmetic. Pair rows are derived from token position; selected expert ids and
6115            // routing weights remain the device router's buffers throughout.
6116            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
6117            let pair_tok_d = e.htod_i32(&pair_tok)?;
6118            let gate = e.moe_pairs_matvec_q8(
6119                &dev.ptr_row,
6120                0,
6121                &pair_tok_d,
6122                &sel_d,
6123                &zq,
6124                &zd,
6125                n_embd,
6126                n_ff_exp,
6127                n_expert,
6128                n_pairs,
6129                m.gate_exps.qtype,
6130                gate_row_bytes,
6131            )?;
6132            let up = e.moe_pairs_matvec_q8(
6133                &dev.ptr_row,
6134                1,
6135                &pair_tok_d,
6136                &sel_d,
6137                &zq,
6138                &zd,
6139                n_embd,
6140                n_ff_exp,
6141                n_expert,
6142                n_pairs,
6143                m.up_exps.qtype,
6144                up_row_bytes,
6145            )?;
6146            let mut act = e.uninit(n_pairs * n_ff_exp)?;
6147            Self::ffn_act_lim(
6148                e,
6149                cfg,
6150                &gate,
6151                &up,
6152                1.0,
6153                1.0,
6154                cfg.clamp_exp_at(il as u32),
6155                &mut act,
6156                n_pairs * n_ff_exp,
6157            )?;
6158            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6159            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6160            let pair_self_d = e.htod_i32(&pair_self)?;
6161            let down = e.moe_pairs_matvec_q8(
6162                &dev.ptr_row,
6163                2,
6164                &pair_self_d,
6165                &sel_d,
6166                &aq2,
6167                &ad2,
6168                n_ff_exp,
6169                n_embd,
6170                n_expert,
6171                n_pairs,
6172                m.down_exps.qtype,
6173                m.down_exps.row_bytes,
6174            )?;
6175            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6176            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6177            let tok_off_d = e.htod_i32(&tok_off)?;
6178            let tok_ids_d = e.htod_i32(&tok_ids)?;
6179            let mut output = e.uninit(t * n_embd)?;
6180            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
6181            output
6182        } else {
6183            let act = e.moe_gate_up_silu8_dev_q8_rows(
6184                &dev.ptr_row,
6185                &sel_d,
6186                &zq,
6187                &zd,
6188                t,
6189                n_embd,
6190                n_ff_exp,
6191                n_used,
6192                n_expert,
6193                m.gate_exps.qtype,
6194                m.up_exps.qtype,
6195                gate_row_bytes,
6196                up_row_bytes,
6197                &m.dev_macros,
6198            )?;
6199            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6200            let mut output = e.uninit(t * n_embd)?;
6201            e.moe_down8_fma_dev_q8_rows_g(
6202                &dev.ptr_row,
6203                &sel_d,
6204                &w_d,
6205                &aq2,
6206                &ad2,
6207                &mut output,
6208                t,
6209                n_ff_exp,
6210                n_embd,
6211                n_used,
6212                n_expert,
6213                m.down_exps.qtype,
6214                m.down_exps.row_bytes,
6215            )?;
6216            output
6217        };
6218
6219        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6220            eprintln!(
6221                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6222                cfg.clamp_exp_at(il as u32).is_some(),
6223                dev.gu_il,
6224            );
6225        }
6226        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6227        Ok(moe_out)
6228    }
6229
6230    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
6231    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
6232    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
6233    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
6234    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
6235    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
6236    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
6237    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
6238    fn moe_ffn_pairs(
6239        e: &Engine,
6240        m: &MoeWeights,
6241        z: &CudaSlice<f32>,
6242        logits: &CudaSlice<f32>,
6243        t: usize,
6244        cfg: &ModelConfig,
6245    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6246        let moe = cfg.moe.as_ref().unwrap();
6247        let n_embd = cfg.n_embd as usize;
6248        let n_expert = moe.expert_count as usize;
6249        let n_used = moe.expert_used_count as usize;
6250        let n_ff_exp = moe.expert_ff_length as usize;
6251        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
6252        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
6253        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
6254        // that forgets the gate fails loudly in debug instead of returning wrong logits.
6255        debug_assert!(
6256            !cfg.swiglu_clamped_anywhere(),
6257            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6258        );
6259        let dev = m.dev_exps.as_ref().unwrap();
6260        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
6261        let (rbg_d, rbu_d) = if dev.gu_il {
6262            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6263            (sxx, sxx)
6264        } else {
6265            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6266        };
6267
6268        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6269        let n_pairs = t * n_used;
6270        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
6271        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
6272        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6273        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6274        let pair_w: Vec<f32> = w_all.clone();
6275        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6276        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6277        let pt = e.htod_i32(&pair_tok)?;
6278        let px = e.htod_i32(&pair_ex)?;
6279        let pw = e.htod(&pair_w)?;
6280        let toff = e.htod_i32(&tok_off)?;
6281        let tids = e.htod_i32(&tok_ids)?;
6282
6283        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
6284        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
6285        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
6286        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6287        for p in 0..n_pairs {
6288            by_ex[pair_ex[p] as usize].push(p as i32);
6289        }
6290        let mut ex_ids: Vec<i32> = Vec::new();
6291        let mut ex_off: Vec<i32> = vec![0];
6292        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6293        for (ex, list) in by_ex.iter().enumerate() {
6294            if list.is_empty() {
6295                continue;
6296            }
6297            ex_ids.push(ex as i32);
6298            ex_pairs.extend_from_slice(list);
6299            ex_off.push(ex_pairs.len() as i32);
6300        }
6301        let n_active = ex_ids.len();
6302        let exi = e.htod_i32(&ex_ids)?;
6303        let exo = e.htod_i32(&ex_off)?;
6304        let exp_d = e.htod_i32(&ex_pairs)?;
6305        let _ = &px; // pair-major twin keeps it; em path uses CSR
6306
6307        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
6308        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
6309        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
6310        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
6311        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
6312        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
6313        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
6314        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
6315        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
6316        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
6317        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
6318        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
6319        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
6320        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
6321        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
6322        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
6323        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
6324        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
6325        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6326        let mma_t = *MMA_T.get_or_init(|| {
6327            std::env::var("MEMRA_MOE_MMA_T")
6328                .ok()
6329                .and_then(|v| v.parse().ok())
6330                .unwrap_or(16)
6331        });
6332        let use_mma = std::env::var("MEMRA_MOE_MMA")
6333            .map(|v| v != "0")
6334            .unwrap_or(true)
6335            && t >= mma_t
6336            && q8_expert_dec_supported(m.gate_exps.qtype)
6337            && q8_expert_dec_supported(m.up_exps.qtype)
6338            && q8_expert_dec_supported(m.down_exps.qtype)
6339            && n_embd % 256 == 0
6340            && n_ff_exp % 256 == 0;
6341        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
6342        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
6343        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
6344        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
6345        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
6346        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
6347        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
6348        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
6349        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
6350        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
6351        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
6352        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
6353        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
6354        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
6355        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
6356        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6357            && q8_expert_dec_supported(m.up_exps.qtype)
6358            && q8_expert_dec_supported(m.down_exps.qtype)
6359            && n_embd % 256 == 0
6360            && n_ff_exp % 256 == 0;
6361        let f16g_mode = crate::moe_f16g_mode();
6362        let f16g = f16g_mode != 0
6363            && t >= mma_t
6364            && (f16g_mode != 3 || !mma_capable)
6365            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6366            && f16g_proj_ok(m.up_exps.qtype, n_embd)
6367            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6368        if use_mma || f16g {
6369            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
6370            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
6371            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
6372            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
6373            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
6374            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
6375            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
6376            let y_down = if f16g {
6377                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
6378                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
6379                // permute at the very end back to pair-id order for the scatter.
6380                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6381                let csr_tok_d = e.htod_i32(&csr_tok)?;
6382                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6383                let g_csr = e.moe_f16_grouped(
6384                    &dev.ptr_row,
6385                    0,
6386                    n_expert,
6387                    &exi,
6388                    &ex_off,
6389                    &exo,
6390                    &z_f16,
6391                    &z_s,
6392                    n_embd,
6393                    n_ff_exp,
6394                    n_active,
6395                    n_pairs,
6396                    m.gate_exps.qtype,
6397                    rbg_d,
6398                )?;
6399                let u_csr = e.moe_f16_grouped(
6400                    &dev.ptr_row,
6401                    1,
6402                    n_expert,
6403                    &exi,
6404                    &ex_off,
6405                    &exo,
6406                    &z_f16,
6407                    &z_s,
6408                    n_embd,
6409                    n_ff_exp,
6410                    n_active,
6411                    n_pairs,
6412                    m.up_exps.qtype,
6413                    rbu_d,
6414                )?;
6415                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6416                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6417                let d_csr = e.moe_f16_grouped(
6418                    &dev.ptr_row,
6419                    2,
6420                    n_expert,
6421                    &exi,
6422                    &ex_off,
6423                    &exo,
6424                    &a_f16,
6425                    &a_s,
6426                    n_ff_exp,
6427                    n_embd,
6428                    n_active,
6429                    n_pairs,
6430                    m.down_exps.qtype,
6431                    m.down_exps.row_bytes,
6432                )?;
6433                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6434            } else {
6435                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
6436                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6437                let gate = e.mmq_iq_experts(
6438                    &dev.ptr_row,
6439                    0,
6440                    n_expert,
6441                    &exi,
6442                    &exo,
6443                    &exp_d,
6444                    &pt,
6445                    &z_scr,
6446                    n_embd,
6447                    n_ff_exp,
6448                    n_active,
6449                    n_pairs,
6450                    t,
6451                    m.gate_exps.qtype,
6452                    rbg_d,
6453                )?;
6454                let up = e.mmq_iq_experts(
6455                    &dev.ptr_row,
6456                    1,
6457                    n_expert,
6458                    &exi,
6459                    &exo,
6460                    &exp_d,
6461                    &pt,
6462                    &z_scr,
6463                    n_embd,
6464                    n_ff_exp,
6465                    n_active,
6466                    n_pairs,
6467                    t,
6468                    m.up_exps.qtype,
6469                    rbu_d,
6470                )?;
6471                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
6472                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
6473                // registers and writes ONLY the quantized scratch — the two-pass chain
6474                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
6475                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
6476                let a_scr = if crate::moe_fuse_actq_on() {
6477                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6478                } else {
6479                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6480                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6481                };
6482                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6483                let pself = e.htod_i32(&pair_self)?;
6484                e.mmq_iq_experts(
6485                    &dev.ptr_row,
6486                    2,
6487                    n_expert,
6488                    &exi,
6489                    &exo,
6490                    &exp_d,
6491                    &pself,
6492                    &a_scr,
6493                    n_ff_exp,
6494                    n_embd,
6495                    n_active,
6496                    n_pairs,
6497                    n_pairs,
6498                    m.down_exps.qtype,
6499                    m.down_exps.row_bytes,
6500                )?
6501            };
6502            let mut moe_out = e.uninit(t * n_embd)?;
6503            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6504            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6505                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6506            {
6507                let n_ff_sh = gate_shexp.out_features();
6508                let sg_gate = e.matmul(gate_shexp, z, t)?;
6509                let sg_up = e.matmul(up_shexp, z, t)?;
6510                let mut sa = e.uninit(t * n_ff_sh)?;
6511                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6512                let sh = e.matmul(down_shexp, &sa, t)?;
6513                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6514                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
6515                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
6516                // i.e. the one real prefill actually takes on a resident-expert MoE model,
6517                // so the concat-prime isolation fix has to land here as well.
6518                let g = match &m.gate_inp_shexp {
6519                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6520                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6521                    }
6522                    Some(gate_inp_shexp) => {
6523                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6524                        let mut g = e.uninit(t)?;
6525                        e.sigmoid(&gs, &mut g, t)?;
6526                        g
6527                    }
6528                    None => e.htod(&vec![1.0f32; t])?,
6529                };
6530                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6531            }
6532            return Ok(moe_out);
6533        }
6534
6535        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
6536        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
6537        let dec = std::env::var("MEMRA_MOE_DEC")
6538            .map(|v| v != "0")
6539            .unwrap_or(true);
6540        let matvec = |proj,
6541                      exi: &_,
6542                      exo: &_,
6543                      exp_d: &_,
6544                      pt: &_,
6545                      aq: &_,
6546                      ad: &_,
6547                      inf,
6548                      outf,
6549                      qtype,
6550                      rb|
6551         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6552            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
6553            let dec = dec && q8_expert_dec_supported(qtype);
6554            if dec {
6555                e.moe_pairs_matvec_q8_dec(
6556                    &dev.ptr_row,
6557                    proj,
6558                    exi,
6559                    exo,
6560                    exp_d,
6561                    pt,
6562                    aq,
6563                    ad,
6564                    inf,
6565                    outf,
6566                    n_expert,
6567                    n_active,
6568                    n_pairs,
6569                    qtype,
6570                    rb,
6571                )
6572            } else {
6573                e.moe_pairs_matvec_q8_em(
6574                    &dev.ptr_row,
6575                    proj,
6576                    exi,
6577                    exo,
6578                    exp_d,
6579                    pt,
6580                    aq,
6581                    ad,
6582                    inf,
6583                    outf,
6584                    n_expert,
6585                    n_active,
6586                    n_pairs,
6587                    qtype,
6588                    rb,
6589                )
6590            }
6591        };
6592        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6593        let gate = matvec(
6594            0,
6595            &exi,
6596            &exo,
6597            &exp_d,
6598            &pt,
6599            &zq,
6600            &zd,
6601            n_embd,
6602            n_ff_exp,
6603            m.gate_exps.qtype,
6604            rbg_d,
6605        )?;
6606        let up = matvec(
6607            1,
6608            &exi,
6609            &exo,
6610            &exp_d,
6611            &pt,
6612            &zq,
6613            &zd,
6614            n_embd,
6615            n_ff_exp,
6616            m.up_exps.qtype,
6617            rbu_d,
6618        )?;
6619        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6620        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6621        // down consumes PAIR-major activation rows: pair_tok = identity.
6622        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6623        let pself = e.htod_i32(&pair_self)?;
6624        let y_down = matvec(
6625            2,
6626            &exi,
6627            &exo,
6628            &exp_d,
6629            &pself,
6630            &aq2,
6631            &ad2,
6632            n_ff_exp,
6633            n_embd,
6634            m.down_exps.qtype,
6635            m.down_exps.row_bytes,
6636        )?;
6637        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
6638        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6639
6640        // SHARED EXPERT epilogue — same as the other paths.
6641        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
6642        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
6643        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6644            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6645        {
6646            let n_ff_sh = gate_shexp.out_features();
6647            // These decode-exact forms are required by the new Step resident arm. Keep the
6648            // established grouped shared-expert program for every other architecture: widening
6649            // this to Gemma changed its speculative acceptance despite green argmax gates.
6650            let step_exact = cfg.step35.is_some();
6651            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6652            let (sg_gate, sg_up) = if step_exact && t == 1 {
6653                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6654                    Some(pair) => pair,
6655                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6656                }
6657            } else if verify_t {
6658                let mut fused = None;
6659                if crate::spec::spec_fused_t()
6660                    && (2..=4).contains(&t)
6661                    && e.uses_q8_1_fast(gate_shexp)
6662                    && e.uses_q8_1_fast(up_shexp)
6663                {
6664                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6665                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6666                }
6667                match fused {
6668                    Some(pair) => pair,
6669                    None => (
6670                        e.matmul_decode_exact(gate_shexp, z, t)?,
6671                        e.matmul_decode_exact(up_shexp, z, t)?,
6672                    ),
6673                }
6674            } else {
6675                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6676            };
6677            let mut sa = e.uninit(t * n_ff_sh)?;
6678            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6679            let sh = if verify_t {
6680                e.matmul_decode_exact(down_shexp, &sa, t)?
6681            } else {
6682                e.matmul(down_shexp, &sa, t)?
6683            };
6684            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6685            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
6686            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
6687            // dispatch choice cannot change bits.
6688            let g = match &m.gate_inp_shexp {
6689                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6690                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6691                }
6692                Some(gate_inp_shexp) => {
6693                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6694                    let mut g = e.uninit(t)?;
6695                    e.sigmoid(&gs, &mut g, t)?;
6696                    g
6697                }
6698                None => e.htod(&vec![1.0f32; t])?,
6699            };
6700            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6701        }
6702        Ok(moe_out)
6703    }
6704
6705    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
6706    #[allow(clippy::too_many_arguments)]
6707    #[allow(clippy::too_many_arguments)]
6708    fn moe_ffn_dev(
6709        e: &Engine,
6710        m: &MoeWeights,
6711        z: &CudaSlice<f32>,
6712        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6713        logits: &CudaSlice<f32>,
6714        t: usize,
6715        cfg: &ModelConfig,
6716        il: u16,
6717        max_block: usize,
6718    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6719        let moe = cfg.moe.as_ref().unwrap();
6720        let n_embd = cfg.n_embd as usize;
6721        let n_expert = moe.expert_count as usize;
6722        let n_used = moe.expert_used_count as usize;
6723        let n_ff_exp = moe.expert_ff_length as usize;
6724        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
6725        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
6726        // clamped layers; assert both so a future caller that skips the gate fails loudly.
6727        debug_assert!(
6728            cfg.sigmoid_router().is_none(),
6729            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6730        );
6731        debug_assert!(
6732            !cfg.swiglu_clamped_at(il as u32),
6733            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6734        );
6735
6736        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
6737        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6738        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
6739        // skipped entirely for macro-free experts (every k-quant GGUF).
6740        if m.has_macros {
6741            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6742        }
6743
6744        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
6745        let mut moe_out = e.uninit(t * n_embd)?;
6746
6747        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
6748        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
6749        if let Some(dev) = m.dev_exps.as_ref() {
6750            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
6751            // the combined stride; up's base is offset in the ptr table. Down unchanged.
6752            let (rbg_d, rbu_d) = if dev.gu_il {
6753                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6754                (sxx, sxx)
6755            } else {
6756                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6757            };
6758            let q8 = moe_q8_enabled()
6759                && q8_expert_supported(m.gate_exps.qtype)
6760                && q8_expert_supported(m.up_exps.qtype)
6761                && q8_expert_supported(m.down_exps.qtype);
6762            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
6763            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
6764            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
6765            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
6766            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
6767            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
6768            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
6769            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
6770            let rows_arm = q8
6771                && t > 1
6772                && crate::spec::spec_m2()
6773                && n_ff_exp == 512
6774                && n_used <= 8
6775                && std::env::var("MEMRA_MOE_DEVQ8_GU")
6776                    .map(|v| v.is_empty() || v == "v")
6777                    .unwrap_or(true)
6778                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6779                    .map(|v| v.is_empty() || v == "w8h2v")
6780                    .unwrap_or(true);
6781            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
6782            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
6783            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
6784            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
6785            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
6786            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
6787            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
6788            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
6789            let csr_mode = std::env::var("MEMRA_MOE_CSR")
6790                .ok()
6791                .and_then(|v| v.parse::<i32>().ok())
6792                .unwrap_or(1);
6793            // NVFP4 admission REVERTED 2026-08-21 (lane/samplat, decode-batch-gate2 find):
6794            // the csr_nvfp4 kernel drifts last-ULP vs the rows program (11041/32768 ACT
6795            // elements at t=8) and the drift is BATCH-COMPOSITION-DEPENDENT — gate2 (B=8 vs
6796            // isolated) FAILED on the ornith15 artifact, the one-numeric-program law's batch
6797            // axis. Three chain-pinning attempts did not close it (receipts,
6798            // research/samplat-20260821/); a source-verbatim per-pair helper form IS
6799            // bit-identical but loses the dedup win (-3% vs rows). NVFP4 stays on the rows
6800            // twins until a cached form passes gate2 + the =2 byte-compare at t=8. The
6801            // increment-1 qualification hole: =2 ran across run-spec (solo verify shapes),
6802            // never decode-batch-gate at B=8 on the MoE model itself.
6803            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6804            let csr_uniform = m.gate_exps.qtype == m.up_exps.qtype;
6805            let csr_arm = rows_arm
6806                && csr_mode > 0
6807                && t <= 10
6808                && csr_uniform
6809                && csr_qt(m.gate_exps.qtype)
6810                && csr_qt(m.up_exps.qtype)
6811                && csr_qt(m.down_exps.qtype);
6812            if csr_arm {
6813                if csr_mode == 2 {
6814                    static ENGAGED: std::sync::Once = std::sync::Once::new();
6815                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6816                }
6817                let n_pairs = t * n_used;
6818                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6819                let act = e.moe_gate_up_silu8_dev_q8_csr(
6820                    &dev.ptr_row,
6821                    &sel_d,
6822                    &zq,
6823                    &zd,
6824                    n_pairs,
6825                    n_embd,
6826                    n_ff_exp,
6827                    n_used,
6828                    n_expert,
6829                    m.gate_exps.qtype,
6830                    m.up_exps.qtype,
6831                    rbg_d,
6832                    rbu_d,
6833                )?;
6834                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6835                // down stays on the _rows twin — BOTH CSR down variants measured negative
6836                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
6837                // 16-group rows have too little decode to amortize any dedup structure.
6838                e.moe_down8_fma_dev_q8_rows(
6839                    &dev.ptr_row,
6840                    &sel_d,
6841                    &w_d,
6842                    &aq2,
6843                    &ad2,
6844                    &mut moe_out,
6845                    t,
6846                    n_ff_exp,
6847                    n_embd,
6848                    n_used,
6849                    n_expert,
6850                    m.down_exps.qtype,
6851                    m.down_exps.row_bytes,
6852                )?;
6853                if csr_mode == 2 {
6854                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
6855                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6856                        &dev.ptr_row,
6857                        &sel_d,
6858                        &zq,
6859                        &zd,
6860                        t,
6861                        n_embd,
6862                        n_ff_exp,
6863                        n_used,
6864                        n_expert,
6865                        m.gate_exps.qtype,
6866                        m.up_exps.qtype,
6867                        rbg_d,
6868                        rbu_d,
6869                        &m.dev_macros,
6870                    )?;
6871                    let mut out_r = e.uninit(t * n_embd)?;
6872                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6873                    e.moe_down8_fma_dev_q8_rows(
6874                        &dev.ptr_row,
6875                        &sel_d,
6876                        &w_d,
6877                        &aq2r,
6878                        &ad2r,
6879                        &mut out_r,
6880                        t,
6881                        n_ff_exp,
6882                        n_embd,
6883                        n_used,
6884                        n_expert,
6885                        m.down_exps.qtype,
6886                        m.down_exps.row_bytes,
6887                    )?;
6888                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6889                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6890                    let ba = a1
6891                        .iter()
6892                        .zip(&a2)
6893                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6894                        .count();
6895                    let bo = o1
6896                        .iter()
6897                        .zip(&o2)
6898                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6899                        .count();
6900                    if ba + bo > 0 {
6901                        eprintln!(
6902                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6903                            a1.len(),
6904                            o1.len()
6905                        );
6906                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
6907                        let sel_h = e.dtoh_i32(&sel_d)?;
6908                        let mut shown = 0;
6909                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6910                            if x.to_bits() != y.to_bits() && shown < 4 {
6911                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6912                                let ex = sel_h[p];
6913                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
6914                                eprintln!(
6915                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6916                                );
6917                                shown += 1;
6918                            }
6919                        }
6920                        std::process::exit(3);
6921                    }
6922                }
6923            } else if rows_arm {
6924                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
6925                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
6926                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6927                    use std::sync::atomic::{AtomicU64, Ordering};
6928                    static PAIRS: AtomicU64 = AtomicU64::new(0);
6929                    static UNIQ: AtomicU64 = AtomicU64::new(0);
6930                    static CALLS: AtomicU64 = AtomicU64::new(0);
6931                    let sel_h = e.dtoh_i32(&sel_d)?;
6932                    let mut u: Vec<i32> = sel_h.clone();
6933                    u.sort_unstable();
6934                    u.dedup();
6935                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6936                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6937                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6938                    if c % 480 == 0 {
6939                        let p = PAIRS.load(Ordering::Relaxed);
6940                        let q = UNIQ.load(Ordering::Relaxed);
6941                        eprintln!(
6942                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6943                            q as f64 / p as f64
6944                        );
6945                    }
6946                }
6947                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6948                let act = e.moe_gate_up_silu8_dev_q8_rows(
6949                    &dev.ptr_row,
6950                    &sel_d,
6951                    &zq,
6952                    &zd,
6953                    t,
6954                    n_embd,
6955                    n_ff_exp,
6956                    n_used,
6957                    n_expert,
6958                    m.gate_exps.qtype,
6959                    m.up_exps.qtype,
6960                    rbg_d,
6961                    rbu_d,
6962                    &m.dev_macros,
6963                )?;
6964                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6965                e.moe_down8_fma_dev_q8_rows(
6966                    &dev.ptr_row,
6967                    &sel_d,
6968                    &w_d,
6969                    &aq2,
6970                    &ad2,
6971                    &mut moe_out,
6972                    t,
6973                    n_ff_exp,
6974                    n_embd,
6975                    n_used,
6976                    n_expert,
6977                    m.down_exps.qtype,
6978                    m.down_exps.row_bytes,
6979                )?;
6980            } else {
6981                for tok in 0..t {
6982                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6983                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6984                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6985                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6986                    if q8 {
6987                        let (zq, zd) = match (t, zq8) {
6988                            (1, Some((q, d))) => (q.clone(), d.clone()),
6989                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6990                        };
6991                        let act = e.moe_gate_up_silu8_dev_q8(
6992                            &dev.ptr_row,
6993                            &selt,
6994                            &zq,
6995                            &zd,
6996                            n_embd,
6997                            n_ff_exp,
6998                            n_used,
6999                            n_expert,
7000                            m.gate_exps.qtype,
7001                            m.up_exps.qtype,
7002                            rbg_d,
7003                            rbu_d,
7004                            &m.dev_macros,
7005                        )?;
7006                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7007                        e.moe_down8_fma_dev_q8(
7008                            &dev.ptr_row,
7009                            &selt,
7010                            &wt,
7011                            &aq2,
7012                            &ad2,
7013                            &mut dst,
7014                            n_ff_exp,
7015                            n_embd,
7016                            n_used,
7017                            n_expert,
7018                            m.down_exps.qtype,
7019                            m.down_exps.row_bytes,
7020                        )?;
7021                    } else {
7022                        let act = e.moe_gate_up_silu8_dev(
7023                            &dev.ptr_row,
7024                            &selt,
7025                            &zt,
7026                            n_embd,
7027                            n_ff_exp,
7028                            n_used,
7029                            n_expert,
7030                            m.gate_exps.qtype,
7031                            m.up_exps.qtype,
7032                            rbg_d,
7033                            rbu_d,
7034                            &m.dev_macros,
7035                        )?;
7036                        e.moe_down8_fma_dev(
7037                            &dev.ptr_row,
7038                            &selt,
7039                            &wt,
7040                            &act,
7041                            &mut dst,
7042                            n_ff_exp,
7043                            n_embd,
7044                            n_used,
7045                            n_expert,
7046                            m.down_exps.qtype,
7047                            m.down_exps.row_bytes,
7048                        )?;
7049                    }
7050                }
7051            }
7052        } else {
7053            // Launch under the cache lock: the row borrow lives as long as the closure, and the
7054            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
7055            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
7056            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
7057            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
7058            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
7059            let q8 = moe_q8_enabled()
7060                && q8_expert_supported(m.gate_exps.qtype)
7061                && q8_expert_supported(m.up_exps.qtype)
7062                && q8_expert_supported(m.down_exps.qtype);
7063            e.with_moe_cache(max_block, |c, eng| {
7064                let row = c
7065                    .layer_dev_row(il, n_expert, eng)?
7066                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
7067                for tok in 0..t {
7068                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
7069                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
7070                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
7071                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7072                    if q8 {
7073                        let (zq, zd) = match (t, zq8) {
7074                            (1, Some((q, d))) => (q.clone(), d.clone()),
7075                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
7076                        };
7077                        let act = eng.moe_gate_up_silu8_dev_q8(
7078                            row,
7079                            &selt,
7080                            &zq,
7081                            &zd,
7082                            n_embd,
7083                            n_ff_exp,
7084                            n_used,
7085                            n_expert,
7086                            m.gate_exps.qtype,
7087                            m.up_exps.qtype,
7088                            m.gate_exps.row_bytes,
7089                            m.up_exps.row_bytes,
7090                            &m.dev_macros,
7091                        )?;
7092                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
7093                        eng.moe_down8_fma_dev_q8(
7094                            row,
7095                            &selt,
7096                            &wt,
7097                            &aq2,
7098                            &ad2,
7099                            &mut dst,
7100                            n_ff_exp,
7101                            n_embd,
7102                            n_used,
7103                            n_expert,
7104                            m.down_exps.qtype,
7105                            m.down_exps.row_bytes,
7106                        )?;
7107                    } else {
7108                        let act = eng.moe_gate_up_silu8_dev(
7109                            row,
7110                            &selt,
7111                            &zt,
7112                            n_embd,
7113                            n_ff_exp,
7114                            n_used,
7115                            n_expert,
7116                            m.gate_exps.qtype,
7117                            m.up_exps.qtype,
7118                            m.gate_exps.row_bytes,
7119                            m.up_exps.row_bytes,
7120                            &m.dev_macros,
7121                        )?;
7122                        eng.moe_down8_fma_dev(
7123                            row,
7124                            &selt,
7125                            &wt,
7126                            &act,
7127                            &mut dst,
7128                            n_ff_exp,
7129                            n_embd,
7130                            n_used,
7131                            n_expert,
7132                            m.down_exps.qtype,
7133                            m.down_exps.row_bytes,
7134                        )?;
7135                    }
7136                }
7137                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
7138                c.hits += (t * 3 * n_used) as u64;
7139                Ok(())
7140            })?;
7141        }
7142
7143        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
7144        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
7145        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
7146        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
7147        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7148            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7149        {
7150            let n_ff_sh = gate_shexp.out_features();
7151            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
7152            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
7153            let verify_t = t > 1 && t < PRIME_MIN_T;
7154            let (sg_gate, sg_up) = if t == 1 {
7155                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
7156                    Some(pair) => pair,
7157                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
7158                }
7159            } else if verify_t {
7160                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
7161                // rides one shared quantize + one fused2 batched launch instead of two
7162                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
7163                let mut fused = None;
7164                if crate::spec::spec_fused_t()
7165                    && (2..=4).contains(&t)
7166                    && e.uses_q8_1_fast(gate_shexp)
7167                    && e.uses_q8_1_fast(up_shexp)
7168                {
7169                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7170                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
7171                }
7172                match fused {
7173                    Some(pair) => pair,
7174                    None => (
7175                        e.matmul_decode_exact(gate_shexp, z, t)?,
7176                        e.matmul_decode_exact(up_shexp, z, t)?,
7177                    ),
7178                }
7179            } else {
7180                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
7181            };
7182            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
7183            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
7184            let sh = if verify_t {
7185                e.matmul_decode_exact(down_shexp, &sa, t)?
7186            } else {
7187                e.matmul(down_shexp, &sa, t)?
7188            };
7189            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
7190            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
7191            // between the two arms; prefill keeps the batched cuBLASLt linear).
7192            let g = match &m.gate_inp_shexp {
7193                Some(gate_inp_shexp) => {
7194                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
7195                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
7196                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7197                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7198                    } else {
7199                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7200                        let mut g = e.uninit(t)?;
7201                        e.sigmoid(&gs, &mut g, t)?;
7202                        g
7203                    }
7204                }
7205                None => e.htod(&vec![1.0f32; t])?,
7206            };
7207            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7208        }
7209
7210        Ok(moe_out)
7211    }
7212
7213    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
7214    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
7215    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
7216    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
7217    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
7218    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
7219    /// the collected raw pointers cannot move between collection and launch (single-threaded
7220    /// decode; the lock is held only for collection, launches are stream-ordered after any
7221    /// prior same-stream staging writes).
7222    #[allow(clippy::too_many_arguments)]
7223    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
7224    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
7225    #[allow(clippy::too_many_arguments)]
7226    fn moe_gdec_token_q8(
7227        e: &Engine,
7228        m: &MoeWeights,
7229        il: u16,
7230        max_block: usize,
7231        zq: &CudaSlice<i8>,
7232        zd: &CudaSlice<f32>,
7233        sel: &[u32],
7234        w: &[f32],
7235        moe_out: &mut CudaSlice<f32>,
7236        tok: usize,
7237        n_embd: usize,
7238        n_ff_exp: usize,
7239        n_used: usize,
7240    ) -> Result<bool, Box<dyn std::error::Error>> {
7241        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7242        use cudarc::driver::DevicePtr;
7243        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7244            let mut g = [0u64; 8];
7245            let mut u = [0u64; 8];
7246            let mut d = [0u64; 8];
7247            for (j, &ex) in sel.iter().enumerate() {
7248                let ex = ex as u16;
7249                let (Some(sg), Some(su), Some(sd)) = (
7250                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7251                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7252                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7253                ) else {
7254                    return Ok(None);
7255                };
7256                let __s = eng.stream();
7257                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7258                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7259                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7260                g[j] = pg as u64;
7261                u[j] = pu as u64;
7262                d[j] = pd as u64;
7263            }
7264            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7265                for &ex in sel {
7266                    let ex = ex as u16;
7267                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7268                        c.note_profile_hit(BlockId::new(il, proj, ex));
7269                    }
7270                }
7271            }
7272            c.hits += (3 * n_used) as u64;
7273            Ok(Some((g, u, d)))
7274        })?;
7275        let Some((g, u, d)) = ptrs else {
7276            return Ok(false);
7277        };
7278        let mut wv = [0f32; 8];
7279        wv[..n_used].copy_from_slice(w);
7280        let act = e.moe_gate_up_silu8_q8(
7281            crate::WPtr8(g),
7282            crate::WPtr8(u),
7283            zq,
7284            zd,
7285            n_embd,
7286            n_ff_exp,
7287            n_used,
7288            m.gate_exps.qtype,
7289            m.up_exps.qtype,
7290            m.gate_exps.row_bytes,
7291            m.up_exps.row_bytes,
7292        )?;
7293        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
7294        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7295        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7296        e.moe_down8_fma_q8(
7297            crate::WPtr8(d),
7298            crate::F32x8(wv),
7299            &aq2,
7300            &ad2,
7301            &mut dst,
7302            n_ff_exp,
7303            n_embd,
7304            n_used,
7305            m.down_exps.qtype,
7306            m.down_exps.row_bytes,
7307        )?;
7308        Ok(true)
7309    }
7310
7311    fn moe_gdec_token(
7312        e: &Engine,
7313        m: &MoeWeights,
7314        il: u16,
7315        max_block: usize,
7316        zt: &cudarc::driver::CudaView<f32>,
7317        sel: &[u32],
7318        w: &[f32],
7319        moe_out: &mut CudaSlice<f32>,
7320        tok: usize,
7321        n_embd: usize,
7322        n_ff_exp: usize,
7323        n_used: usize,
7324    ) -> Result<bool, Box<dyn std::error::Error>> {
7325        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7326        use cudarc::driver::DevicePtr;
7327        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
7328        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7329            let mut g = [0u64; 8];
7330            let mut u = [0u64; 8];
7331            let mut d = [0u64; 8];
7332            for (j, &ex) in sel.iter().enumerate() {
7333                let ex = ex as u16;
7334                let (Some(sg), Some(su), Some(sd)) = (
7335                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7336                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7337                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7338                ) else {
7339                    return Ok(None);
7340                };
7341                let __s = eng.stream();
7342                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7343                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7344                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7345                g[j] = pg as u64;
7346                u[j] = pu as u64;
7347                d[j] = pd as u64;
7348            }
7349            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7350                for &ex in sel {
7351                    let ex = ex as u16;
7352                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7353                        c.note_profile_hit(BlockId::new(il, proj, ex));
7354                    }
7355                }
7356            }
7357            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
7358            Ok(Some((g, u, d)))
7359        })?;
7360        let Some((g, u, d)) = ptrs else {
7361            return Ok(false);
7362        };
7363        let mut wv = [0f32; 8];
7364        wv[..n_used].copy_from_slice(w);
7365        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
7366        let act = e.moe_gate_up_silu8(
7367            crate::WPtr8(g),
7368            crate::WPtr8(u),
7369            zt,
7370            n_embd,
7371            n_ff_exp,
7372            n_used,
7373            m.gate_exps.qtype,
7374            m.up_exps.qtype,
7375            m.gate_exps.row_bytes,
7376            m.up_exps.row_bytes,
7377        )?;
7378        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7379        e.moe_down8_fma_into(
7380            crate::WPtr8(d),
7381            crate::F32x8(wv),
7382            &act,
7383            &mut dst,
7384            n_ff_exp,
7385            n_embd,
7386            n_used,
7387            m.down_exps.qtype,
7388            m.down_exps.row_bytes,
7389        )?;
7390        Ok(true)
7391    }
7392
7393    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
7394    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
7395    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
7396    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
7397    fn moe_cached_gemm_q8(
7398        e: &Engine,
7399        il: u16,
7400        proj: u8,
7401        ex: usize,
7402        m: &MoeWeights,
7403        max_block: usize,
7404        aq: &CudaSlice<i8>,
7405        ad: &CudaSlice<f32>,
7406    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7407        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7408        let exps = match proj {
7409            PROJ_GATE => &m.gate_exps,
7410            PROJ_UP => &m.up_exps,
7411            _ => &m.down_exps,
7412        };
7413        let layout = exps.expert_layout(ex);
7414        let id = BlockId::new(il, proj, ex as u16);
7415        let source = exps.expert_source(ex);
7416        e.with_moe_cache(max_block, |c, eng| {
7417            let slot = c.dispatch_source(id, source, eng)?;
7418            let DispatchSlot::Resident(sl) = slot;
7419            let buf = c.slot(sl);
7420            eng.qmatvec_expert_q8(
7421                buf,
7422                0..layout.len,
7423                aq,
7424                ad,
7425                1,
7426                exps.in_f,
7427                exps.out_f,
7428                layout.qtype,
7429                layout.row_bytes,
7430            )
7431        })
7432    }
7433
7434    fn moe_cached_gemm(
7435        e: &Engine,
7436        il: u16,
7437        proj: u8,
7438        ex: usize,
7439        m: &MoeWeights,
7440        max_block: usize,
7441        x: &cudarc::driver::CudaView<f32>,
7442    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7443        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7444        let exps = match proj {
7445            PROJ_GATE => &m.gate_exps,
7446            PROJ_UP => &m.up_exps,
7447            _ => &m.down_exps,
7448        };
7449        let layout = exps.expert_layout(ex);
7450        let id = BlockId::new(il, proj, ex as u16);
7451        let source = exps.expert_source(ex);
7452        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
7453        e.with_moe_cache(max_block, |c, eng| {
7454            let slot = c.dispatch_source(id, source, eng)?;
7455            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
7456            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
7457            let DispatchSlot::Resident(sl) = slot;
7458            let buf = c.slot(sl);
7459            eng.qmatvec_view(
7460                buf,
7461                0..layout.len,
7462                x,
7463                1,
7464                exps.in_f,
7465                exps.out_f,
7466                layout.qtype,
7467                layout.row_bytes,
7468            )
7469        })
7470    }
7471
7472    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
7473    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
7474    /// so the current forward's backend assignment and output remain unchanged.
7475    fn moe_profile_admit_expert(
7476        e: &Engine,
7477        il: u16,
7478        ex: usize,
7479        m: &MoeWeights,
7480        max_block: usize,
7481    ) -> Result<(), Box<dyn std::error::Error>> {
7482        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7483        e.with_moe_cache(max_block, |cache, eng| {
7484            for (proj, exps) in [
7485                (PROJ_GATE, &m.gate_exps),
7486                (PROJ_UP, &m.up_exps),
7487                (PROJ_DOWN, &m.down_exps),
7488            ] {
7489                let id = BlockId::new(il, proj, ex as u16);
7490                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7491            }
7492            Ok(())
7493        })
7494    }
7495
7496    /// Read a projection from the immutable residency set when present; otherwise use one
7497    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
7498    #[allow(clippy::too_many_arguments)]
7499    fn moe_frozen_gemm(
7500        e: &Engine,
7501        il: u16,
7502        proj: u8,
7503        ex: usize,
7504        m: &MoeWeights,
7505        max_block: usize,
7506        x: &cudarc::driver::CudaView<f32>,
7507        scratch: &mut Option<CudaSlice<u8>>,
7508        scratch_len: usize,
7509    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7510        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7511        let exps = match proj {
7512            PROJ_GATE => &m.gate_exps,
7513            PROJ_UP => &m.up_exps,
7514            _ => &m.down_exps,
7515        };
7516        let layout = exps.expert_layout(ex);
7517        let id = BlockId::new(il, proj, ex as u16);
7518        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7519            let Some(slot) = cache.resident(id) else {
7520                return Ok(None);
7521            };
7522            let buf = cache.slot(slot);
7523            Ok(Some(eng.qmatvec_view(
7524                buf,
7525                0..layout.len,
7526                x,
7527                1,
7528                exps.in_f,
7529                exps.out_f,
7530                layout.qtype,
7531                layout.row_bytes,
7532            )?))
7533        })? {
7534            return Ok(output);
7535        }
7536        if scratch.is_none() {
7537            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7538        }
7539        let scratch = scratch.as_mut().unwrap();
7540        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7541        e.qmatvec_view(
7542            scratch,
7543            0..layout.len,
7544            x,
7545            1,
7546            exps.in_f,
7547            exps.out_f,
7548            layout.qtype,
7549            layout.row_bytes,
7550        )
7551    }
7552
7553    fn moe_prefetch_expert(
7554        e: &Engine,
7555        il: u16,
7556        ex: usize,
7557        m: &MoeWeights,
7558        max_block: usize,
7559        keep: &[crate::moe_cache::BlockId],
7560    ) -> Result<(), Box<dyn std::error::Error>> {
7561        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7562        e.with_moe_cache(max_block, |c, eng| {
7563            for (proj, exps) in [
7564                (PROJ_GATE, &m.gate_exps),
7565                (PROJ_UP, &m.up_exps),
7566                (PROJ_DOWN, &m.down_exps),
7567            ] {
7568                let id = BlockId::new(il, proj, ex as u16);
7569                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7570            }
7571            Ok(())
7572        })
7573    }
7574
7575    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
7576    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
7577    fn moe_prefetch_disk_expert(
7578        e: &Engine,
7579        il: u16,
7580        ex: usize,
7581        m: &MoeWeights,
7582        max_block: usize,
7583        keep: &[crate::moe_cache::BlockId],
7584    ) -> Result<(), Box<dyn std::error::Error>> {
7585        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7586        e.with_moe_cache(max_block, |c, eng| {
7587            for (proj, exps) in [
7588                (PROJ_GATE, &m.gate_exps),
7589                (PROJ_UP, &m.up_exps),
7590                (PROJ_DOWN, &m.down_exps),
7591            ] {
7592                let source = exps.expert_source(ex);
7593                if let crate::model::ExpertSource::Disk { .. } = &source {
7594                    let id = BlockId::new(il, proj, ex as u16);
7595                    let _ = c.prefetch_source(id, source, keep, eng)?;
7596                }
7597            }
7598            Ok(())
7599        })
7600    }
7601
7602    #[inline]
7603    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7604        let _ = m.gate_exps.prefetch_expert_pages(ex);
7605        let _ = m.up_exps.prefetch_expert_pages(ex);
7606        let _ = m.down_exps.prefetch_expert_pages(ex);
7607    }
7608}
7609
7610// ================================================================================================
7611// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
7612//
7613// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
7614// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
7615// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
7616//
7617// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
7618// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
7619// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
7620// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
7621// identical to the per-token loop regardless of expert processing order.
7622//
7623// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
7624// ================================================================================================
7625
7626impl HybridModel {
7627    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
7628    /// sequential fused q8 program over the token axis; clamped layers use the separate
7629    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
7630    #[allow(clippy::too_many_arguments)]
7631    fn moe_ffn_grouped_resident_q8(
7632        e: &Engine,
7633        m: &MoeWeights,
7634        z: &CudaSlice<f32>,
7635        t: usize,
7636        cfg: &ModelConfig,
7637        il: u16,
7638        sel_all: &[u32],
7639        w_all: &[f32],
7640        table: &CudaSlice<u64>,
7641        gu_il: bool,
7642    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7643        let moe = cfg.moe.as_ref().unwrap();
7644        let n_embd = cfg.n_embd as usize;
7645        let n_expert = moe.expert_count as usize;
7646        let n_used = moe.expert_used_count as usize;
7647        let n_ff_exp = moe.expert_ff_length as usize;
7648        let n_pairs = t * n_used;
7649        debug_assert_eq!(sel_all.len(), n_pairs);
7650        debug_assert_eq!(w_all.len(), n_pairs);
7651        debug_assert!(
7652            m.gate_exps.macros.is_none()
7653                && m.up_exps.macros.is_none()
7654                && m.down_exps.macros.is_none(),
7655            "resident grouped q8 does not fold per-expert macro scales",
7656        );
7657
7658        // The rows twins run the resident sequential program verbatim on grid.z = token:
7659        // fused gate/up/SiLU per slot, batched activation quantization, then the original
7660        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
7661        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
7662        // never enter the softmax router.
7663        if !cfg.swiglu_clamped_at(il as u32) {
7664            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7665            let sel_d = e.htod_i32(&sel)?;
7666            let w_d = e.htod(w_all)?;
7667            let (gate_row_bytes, up_row_bytes) = if gu_il {
7668                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7669                (combined, combined)
7670            } else {
7671                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7672            };
7673            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7674            let act = e.moe_gate_up_silu8_dev_q8_rows(
7675                table,
7676                &sel_d,
7677                &zq,
7678                &zd,
7679                t,
7680                n_embd,
7681                n_ff_exp,
7682                n_used,
7683                n_expert,
7684                m.gate_exps.qtype,
7685                m.up_exps.qtype,
7686                gate_row_bytes,
7687                up_row_bytes,
7688                &m.dev_macros,
7689            )?;
7690            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7691            let mut moe_out = e.uninit(t * n_embd)?;
7692            e.moe_down8_fma_dev_q8_rows_g(
7693                table,
7694                &sel_d,
7695                &w_d,
7696                &aq2,
7697                &ad2,
7698                &mut moe_out,
7699                t,
7700                n_ff_exp,
7701                n_embd,
7702                n_used,
7703                n_expert,
7704                m.down_exps.qtype,
7705                m.down_exps.row_bytes,
7706            )?;
7707
7708            if std::env::var("MEMRA_MOE_STATS").is_ok() {
7709                let mut counts = vec![0usize; n_expert];
7710                for &expert in sel_all {
7711                    counts[expert as usize] += 1;
7712                }
7713                let mut sizes: Vec<usize> =
7714                    counts.into_iter().filter(|&count| count != 0).collect();
7715                sizes.sort_unstable();
7716                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7717                println!(
7718                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7719                     m_e: min={} median={} mean={mean:.1} max={}",
7720                    sizes.len(),
7721                    n_expert,
7722                    sizes.first().copied().unwrap_or(0),
7723                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7724                    sizes.last().copied().unwrap_or(0),
7725                );
7726            }
7727            return Ok(moe_out);
7728        }
7729
7730        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
7731        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
7732        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
7733        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7734        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7735        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7736        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7737
7738        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7739        for (pair, &expert) in pair_ex.iter().enumerate() {
7740            by_expert[expert as usize].push(pair as i32);
7741        }
7742
7743        let pair_tok_d = e.htod_i32(&pair_tok)?;
7744        let pair_ex_d = e.htod_i32(&pair_ex)?;
7745        let pair_w_d = e.htod(w_all)?;
7746        let tok_off_d = e.htod_i32(&tok_off)?;
7747        let tok_ids_d = e.htod_i32(&tok_ids)?;
7748
7749        let matvec = |proj: i32,
7750                      pair_rows: &CudaSlice<i32>,
7751                      aq: &CudaSlice<i8>,
7752                      ad: &CudaSlice<f32>,
7753                      in_f: usize,
7754                      out_f: usize,
7755                      qtype: i32,
7756                      row_bytes: usize|
7757         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7758            e.moe_pairs_matvec_q8(
7759                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7760                row_bytes,
7761            )
7762        };
7763
7764        let (gate_row_bytes, up_row_bytes) = if gu_il {
7765            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7766            (combined, combined)
7767        } else {
7768            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7769        };
7770        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7771        let gate = matvec(
7772            0,
7773            &pair_tok_d,
7774            &zq,
7775            &zd,
7776            n_embd,
7777            n_ff_exp,
7778            m.gate_exps.qtype,
7779            gate_row_bytes,
7780        )?;
7781        let up = matvec(
7782            1,
7783            &pair_tok_d,
7784            &zq,
7785            &zd,
7786            n_embd,
7787            n_ff_exp,
7788            m.up_exps.qtype,
7789            up_row_bytes,
7790        )?;
7791        let mut act = e.uninit(n_pairs * n_ff_exp)?;
7792        Self::ffn_act_lim(
7793            e,
7794            cfg,
7795            &gate,
7796            &up,
7797            1.0,
7798            1.0,
7799            cfg.clamp_exp_at(il as u32),
7800            &mut act,
7801            n_pairs * n_ff_exp,
7802        )?;
7803        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7804        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7805        let pair_self_d = e.htod_i32(&pair_self)?;
7806        let down = matvec(
7807            2,
7808            &pair_self_d,
7809            &aq2,
7810            &ad2,
7811            n_ff_exp,
7812            n_embd,
7813            m.down_exps.qtype,
7814            m.down_exps.row_bytes,
7815        )?;
7816        let mut moe_out = e.uninit(t * n_embd)?;
7817        e.moe_pairs_scatter(
7818            &down,
7819            &pair_w_d,
7820            &tok_off_d,
7821            &tok_ids_d,
7822            &mut moe_out,
7823            t,
7824            n_embd,
7825        )?;
7826
7827        if std::env::var("MEMRA_MOE_STATS").is_ok() {
7828            let mut sizes: Vec<usize> = by_expert
7829                .iter()
7830                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7831                .collect();
7832            sizes.sort_unstable();
7833            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7834            println!(
7835                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7836                 m_e: min={} median={} mean={mean:.1} max={}",
7837                sizes.len(),
7838                n_expert,
7839                sizes.first().copied().unwrap_or(0),
7840                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7841                sizes.last().copied().unwrap_or(0),
7842            );
7843        }
7844        Ok(moe_out)
7845    }
7846
7847    fn moe_ffn_grouped_add_shared(
7848        e: &Engine,
7849        m: &MoeWeights,
7850        z: &CudaSlice<f32>,
7851        t: usize,
7852        cfg: &ModelConfig,
7853        il: u16,
7854        moe_out: &mut CudaSlice<f32>,
7855    ) -> Result<(), Box<dyn std::error::Error>> {
7856        let n_embd = cfg.n_embd as usize;
7857        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7858            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7859        {
7860            let n_ff_sh = gate_shexp.out_features();
7861            let sg_gate = e.matmul(gate_shexp, z, t)?;
7862            let sg_up = e.matmul(up_shexp, z, t)?;
7863            let mut sa = e.uninit(t * n_ff_sh)?;
7864            Self::ffn_act_lim(
7865                e,
7866                cfg,
7867                &sg_gate,
7868                &sg_up,
7869                1.0,
7870                1.0,
7871                cfg.clamp_shexp_at(il as u32),
7872                &mut sa,
7873                t * n_ff_sh,
7874            )?;
7875            let sh = e.matmul(down_shexp, &sa, t)?;
7876            let gate = match &m.gate_inp_shexp {
7877                Some(gate_inp_shexp) => {
7878                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7879                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7880                    } else {
7881                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7882                        let mut gate = e.uninit(t)?;
7883                        e.sigmoid(&raw, &mut gate, t)?;
7884                        gate
7885                    }
7886                }
7887                None => e.htod(&vec![1.0f32; t])?,
7888            };
7889            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7890        }
7891        Ok(())
7892    }
7893
7894    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
7895    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
7896    pub(crate) fn moe_ffn_grouped(
7897        e: &Engine,
7898        m: &MoeWeights,
7899        z: &CudaSlice<f32>,
7900        t: usize,
7901        cfg: &ModelConfig,
7902        il: u16,
7903        max_block: usize,
7904    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7905        let moe = cfg.moe.as_ref().unwrap();
7906        let n_embd = cfg.n_embd as usize;
7907        let n_expert = moe.expert_count as usize;
7908        let n_used = moe.expert_used_count as usize;
7909        let n_ff_exp = moe.expert_ff_length as usize;
7910        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
7911        let lim_exp = cfg.clamp_exp_at(il as u32);
7912
7913        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
7914        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
7915        // enters the softmax-only pairs/dev router.
7916        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7917        if let Some(sig) = cfg.sigmoid_router() {
7918            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7919        }
7920        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7921            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7922        } else {
7923            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7924        };
7925        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7926        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7927        Self::trace_moe_input(e, il, t, n_embd, z)?;
7928
7929        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
7930        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
7931        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
7932        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
7933        let no_exp_macros = m.gate_exps.macros.is_none()
7934            && m.up_exps.macros.is_none()
7935            && m.down_exps.macros.is_none();
7936        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7937            m.has_uniform_expert_layout()
7938                && no_exp_macros
7939                && moe_q8_enabled()
7940                && q8_expert_supported(m.gate_exps.qtype)
7941                && q8_expert_supported(m.up_exps.qtype)
7942                && q8_expert_supported(m.down_exps.qtype)
7943                && moe_slab_enabled()
7944                && dev.dev == e.ctx().ordinal()
7945        });
7946        if let Some(dev) = resident_q8 {
7947            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7948                e,
7949                m,
7950                z,
7951                t,
7952                cfg,
7953                il,
7954                &sel_all,
7955                &w_all,
7956                &dev.ptr_row,
7957                dev.gu_il,
7958            )?;
7959            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7960            return Ok(moe_out);
7961        }
7962
7963        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
7964        // For each expert e, we need: which tokens use it, their positions in z, their top-k
7965        // slot index (for bit-identical accumulation), and their weights.
7966        struct ExpertGroup {
7967            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
7968            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
7969            weights: Vec<f32>,      // renormalized weight for that token-expert pair
7970        }
7971        let mut groups: Vec<ExpertGroup> = (0..n_expert)
7972            .map(|_| ExpertGroup {
7973                tok_indices: Vec::new(),
7974                slot_indices: Vec::new(),
7975                weights: Vec::new(),
7976            })
7977            .collect();
7978
7979        for tok in 0..t {
7980            for j in 0..n_used {
7981                let ex = sel_all[tok * n_used + j] as usize;
7982                let w = w_all[tok * n_used + j];
7983                groups[ex].tok_indices.push(tok as i32);
7984                groups[ex].slot_indices.push(j as i32);
7985                groups[ex].weights.push(w);
7986            }
7987        }
7988
7989        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
7990        // Each token's 8 expert contributions land in their respective slots.
7991        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7992        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
7993
7994        // Expert weight dimensions (used in both cache and staging paths).
7995        let g_len = m.gate_exps.max_expert_bytes();
7996        let u_len = m.up_exps.max_expert_bytes();
7997        let d_len = m.down_exps.max_expert_bytes();
7998        let moe_q8 = m.has_uniform_expert_layout()
7999            && moe_q8_enabled()
8000            && q8_expert_supported(m.gate_exps.qtype)
8001            && q8_expert_supported(m.up_exps.qtype)
8002            && q8_expert_supported(m.down_exps.qtype);
8003        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
8004        // Interleaved GU slabs require the pointer-table fast path above.
8005        let slab_local = m
8006            .dev_exps
8007            .as_ref()
8008            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
8009        let use_cache =
8010            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
8011        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
8012        // also does: a local resident slab or a live SLRU dispatch.
8013        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
8014
8015        // GPU scratch for staging (only allocated without a local slab or cache).
8016        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
8017            (
8018                Some(e.alloc_u8(g_len)?),
8019                Some(e.alloc_u8(u_len)?),
8020                Some(e.alloc_u8(d_len)?),
8021            )
8022        } else {
8023            (None, None, None)
8024        };
8025
8026        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
8027        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
8028        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
8029        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
8030        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
8031        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
8032        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
8033        // at long prompts where every expert stages regardless. Order is FREE to change without
8034        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
8035        // regardless of expert processing order (the whole point of the slots).
8036        let mut order: Vec<usize> = (0..n_expert)
8037            .filter(|&ex| !groups[ex].tok_indices.is_empty())
8038            .collect();
8039        order.sort_by(|&a, &b| {
8040            groups[b]
8041                .tok_indices
8042                .len()
8043                .cmp(&groups[a].tok_indices.len())
8044                .then(a.cmp(&b))
8045        });
8046        let mut m_dist: Vec<usize> = Vec::new(); // for stats
8047        let page_window = moe_page_prefetch_window();
8048        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
8049        if worker_disk_prefetch {
8050            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
8051                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
8052            }
8053        }
8054        for (order_pos, &ex) in order.iter().enumerate() {
8055            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
8056                Self::moe_prefetch_host_expert(order[next], m);
8057            }
8058            if worker_disk_prefetch {
8059                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
8060                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8061                    let keep = [
8062                        BlockId::new(il, PROJ_GATE, ex as u16),
8063                        BlockId::new(il, PROJ_UP, ex as u16),
8064                        BlockId::new(il, PROJ_DOWN, ex as u16),
8065                    ];
8066                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
8067                }
8068            }
8069            let grp = &groups[ex];
8070            let m_e = grp.tok_indices.len();
8071            m_dist.push(m_e);
8072            let gl = m.gate_exps.expert_layout(ex);
8073            let ul = m.up_exps.expert_layout(ex);
8074            let dl = m.down_exps.expert_layout(ex);
8075
8076            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
8077            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
8078            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
8079            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
8080            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
8081            let dmac = m.down_exps.macro_scale(ex);
8082            let weight_d = if dmac == 1.0 {
8083                e.htod(&grp.weights)?
8084            } else {
8085                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
8086                e.htod(&scaled)?
8087            };
8088
8089            // GATHER: collect m_e activation rows from z into a contiguous buffer.
8090            let mut gathered = e.zeros(m_e * n_embd)?;
8091            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
8092            let gv = gathered.slice(0..m_e * n_embd);
8093
8094            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
8095            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
8096            let y = if let Some(dev) = slab_local {
8097                let gate_start = ex * m.gate_exps.expert_stride;
8098                let up_start = ex * m.up_exps.expert_stride;
8099                let down_start = ex * m.down_exps.expert_stride;
8100                if grouped_q8 {
8101                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8102                    let gate = e.qmatvec_expert_q8(
8103                        &dev.gate,
8104                        gate_start..gate_start + gl.len,
8105                        &zq,
8106                        &zd,
8107                        m_e,
8108                        m.gate_exps.in_f,
8109                        m.gate_exps.out_f,
8110                        gl.qtype,
8111                        gl.row_bytes,
8112                    )?;
8113                    let up = e.qmatvec_expert_q8(
8114                        &dev.up,
8115                        up_start..up_start + ul.len,
8116                        &zq,
8117                        &zd,
8118                        m_e,
8119                        m.up_exps.in_f,
8120                        m.up_exps.out_f,
8121                        ul.qtype,
8122                        ul.row_bytes,
8123                    )?;
8124                    let mut act = e.uninit(m_e * n_ff_exp)?;
8125                    Self::ffn_act_lim(
8126                        e,
8127                        cfg,
8128                        &gate,
8129                        &up,
8130                        m.gate_exps.macro_scale(ex),
8131                        m.up_exps.macro_scale(ex),
8132                        lim_exp,
8133                        &mut act,
8134                        m_e * n_ff_exp,
8135                    )?;
8136                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8137                    e.qmatvec_expert_q8(
8138                        &dev.down,
8139                        down_start..down_start + dl.len,
8140                        &aq2,
8141                        &ad2,
8142                        m_e,
8143                        m.down_exps.in_f,
8144                        m.down_exps.out_f,
8145                        dl.qtype,
8146                        dl.row_bytes,
8147                    )?
8148                } else {
8149                    let gate = e.qmatvec_view(
8150                        &dev.gate,
8151                        gate_start..gate_start + gl.len,
8152                        &gv,
8153                        m_e,
8154                        m.gate_exps.in_f,
8155                        m.gate_exps.out_f,
8156                        gl.qtype,
8157                        gl.row_bytes,
8158                    )?;
8159                    let up = e.qmatvec_view(
8160                        &dev.up,
8161                        up_start..up_start + ul.len,
8162                        &gv,
8163                        m_e,
8164                        m.up_exps.in_f,
8165                        m.up_exps.out_f,
8166                        ul.qtype,
8167                        ul.row_bytes,
8168                    )?;
8169                    let mut act = e.uninit(m_e * n_ff_exp)?;
8170                    Self::ffn_act_lim(
8171                        e,
8172                        cfg,
8173                        &gate,
8174                        &up,
8175                        m.gate_exps.macro_scale(ex),
8176                        m.up_exps.macro_scale(ex),
8177                        lim_exp,
8178                        &mut act,
8179                        m_e * n_ff_exp,
8180                    )?;
8181                    let actv = act.slice(0..m_e * n_ff_exp);
8182                    e.qmatvec_view(
8183                        &dev.down,
8184                        down_start..down_start + dl.len,
8185                        &actv,
8186                        m_e,
8187                        m.down_exps.in_f,
8188                        m.down_exps.out_f,
8189                        dl.qtype,
8190                        dl.row_bytes,
8191                    )?
8192                }
8193            } else if use_cache {
8194                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8195                if grouped_q8 {
8196                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8197                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8198                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8199                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8200                        eng.qmatvec_expert_q8(
8201                            cache.buf(slot),
8202                            0..gl.len,
8203                            &zq,
8204                            &zd,
8205                            m_e,
8206                            m.gate_exps.in_f,
8207                            m.gate_exps.out_f,
8208                            gl.qtype,
8209                            gl.row_bytes,
8210                        )
8211                    })?;
8212                    let up = e.with_moe_cache(max_block, |cache, eng| {
8213                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8214                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8215                        eng.qmatvec_expert_q8(
8216                            cache.buf(slot),
8217                            0..ul.len,
8218                            &zq,
8219                            &zd,
8220                            m_e,
8221                            m.up_exps.in_f,
8222                            m.up_exps.out_f,
8223                            ul.qtype,
8224                            ul.row_bytes,
8225                        )
8226                    })?;
8227                    let mut act = e.uninit(m_e * n_ff_exp)?;
8228                    Self::ffn_act_lim(
8229                        e,
8230                        cfg,
8231                        &gate,
8232                        &up,
8233                        m.gate_exps.macro_scale(ex),
8234                        m.up_exps.macro_scale(ex),
8235                        lim_exp,
8236                        &mut act,
8237                        m_e * n_ff_exp,
8238                    )?;
8239                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8240                    e.with_moe_cache(max_block, |cache, eng| {
8241                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8242                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8243                        eng.qmatvec_expert_q8(
8244                            cache.buf(slot),
8245                            0..dl.len,
8246                            &aq2,
8247                            &ad2,
8248                            m_e,
8249                            m.down_exps.in_f,
8250                            m.down_exps.out_f,
8251                            dl.qtype,
8252                            dl.row_bytes,
8253                        )
8254                    })?
8255                } else {
8256                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8257                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8258                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8259                        eng.qmatvec_view(
8260                            cache.buf(slot),
8261                            0..gl.len,
8262                            &gv,
8263                            m_e,
8264                            m.gate_exps.in_f,
8265                            m.gate_exps.out_f,
8266                            gl.qtype,
8267                            gl.row_bytes,
8268                        )
8269                    })?;
8270                    let up = e.with_moe_cache(max_block, |cache, eng| {
8271                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8272                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8273                        eng.qmatvec_view(
8274                            cache.buf(slot),
8275                            0..ul.len,
8276                            &gv,
8277                            m_e,
8278                            m.up_exps.in_f,
8279                            m.up_exps.out_f,
8280                            ul.qtype,
8281                            ul.row_bytes,
8282                        )
8283                    })?;
8284                    let mut act = e.uninit(m_e * n_ff_exp)?;
8285                    Self::ffn_act_lim(
8286                        e,
8287                        cfg,
8288                        &gate,
8289                        &up,
8290                        m.gate_exps.macro_scale(ex),
8291                        m.up_exps.macro_scale(ex),
8292                        lim_exp,
8293                        &mut act,
8294                        m_e * n_ff_exp,
8295                    )?;
8296                    let actv = act.slice(0..m_e * n_ff_exp);
8297                    e.with_moe_cache(max_block, |cache, eng| {
8298                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8299                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8300                        eng.qmatvec_view(
8301                            cache.buf(slot),
8302                            0..dl.len,
8303                            &actv,
8304                            m_e,
8305                            m.down_exps.in_f,
8306                            m.down_exps.out_f,
8307                            dl.qtype,
8308                            dl.row_bytes,
8309                        )
8310                    })?
8311                }
8312            } else {
8313                let sg = scratch_g.as_mut().unwrap();
8314                let su = scratch_u.as_mut().unwrap();
8315                let sd = scratch_d.as_mut().unwrap();
8316                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8317                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8318                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8319                if grouped_q8 {
8320                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8321                    let gate = e.qmatvec_expert_q8(
8322                        sg,
8323                        0..gl.len,
8324                        &zq,
8325                        &zd,
8326                        m_e,
8327                        m.gate_exps.in_f,
8328                        m.gate_exps.out_f,
8329                        gl.qtype,
8330                        gl.row_bytes,
8331                    )?;
8332                    let up = e.qmatvec_expert_q8(
8333                        su,
8334                        0..ul.len,
8335                        &zq,
8336                        &zd,
8337                        m_e,
8338                        m.up_exps.in_f,
8339                        m.up_exps.out_f,
8340                        ul.qtype,
8341                        ul.row_bytes,
8342                    )?;
8343                    let mut act = e.uninit(m_e * n_ff_exp)?;
8344                    Self::ffn_act_lim(
8345                        e,
8346                        cfg,
8347                        &gate,
8348                        &up,
8349                        m.gate_exps.macro_scale(ex),
8350                        m.up_exps.macro_scale(ex),
8351                        lim_exp,
8352                        &mut act,
8353                        m_e * n_ff_exp,
8354                    )?;
8355                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8356                    e.qmatvec_expert_q8(
8357                        sd,
8358                        0..dl.len,
8359                        &aq2,
8360                        &ad2,
8361                        m_e,
8362                        m.down_exps.in_f,
8363                        m.down_exps.out_f,
8364                        dl.qtype,
8365                        dl.row_bytes,
8366                    )?
8367                } else {
8368                    let gate = e.qmatvec_view(
8369                        sg,
8370                        0..gl.len,
8371                        &gv,
8372                        m_e,
8373                        m.gate_exps.in_f,
8374                        m.gate_exps.out_f,
8375                        gl.qtype,
8376                        gl.row_bytes,
8377                    )?;
8378                    let up = e.qmatvec_view(
8379                        su,
8380                        0..ul.len,
8381                        &gv,
8382                        m_e,
8383                        m.up_exps.in_f,
8384                        m.up_exps.out_f,
8385                        ul.qtype,
8386                        ul.row_bytes,
8387                    )?;
8388                    let mut act = e.uninit(m_e * n_ff_exp)?;
8389                    Self::ffn_act_lim(
8390                        e,
8391                        cfg,
8392                        &gate,
8393                        &up,
8394                        m.gate_exps.macro_scale(ex),
8395                        m.up_exps.macro_scale(ex),
8396                        lim_exp,
8397                        &mut act,
8398                        m_e * n_ff_exp,
8399                    )?;
8400                    let actv = act.slice(0..m_e * n_ff_exp);
8401                    e.qmatvec_view(
8402                        sd,
8403                        0..dl.len,
8404                        &actv,
8405                        m_e,
8406                        m.down_exps.in_f,
8407                        m.down_exps.out_f,
8408                        dl.qtype,
8409                        dl.row_bytes,
8410                    )?
8411                }
8412            };
8413
8414            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
8415            e.scatter_slot(
8416                &y,
8417                &tok_idx_d,
8418                &slot_idx_d,
8419                &weight_d,
8420                &mut slot_buf,
8421                &mut wbuf,
8422                n_embd,
8423                n_used,
8424                m_e,
8425            )?;
8426        }
8427
8428        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
8429        let mut moe_out = e.zeros(t * n_embd)?;
8430        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8431
8432        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
8433        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8434            m_dist.sort_unstable();
8435            let active = m_dist.len();
8436            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8437            let median = m_dist[active / 2];
8438            let max_m = *m_dist.last().unwrap();
8439            let min_m = m_dist[0];
8440            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8441            println!(
8442                "moe-grouped il={il} t={t} active={active}/{n_expert} \
8443                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8444                      above_gemm_threshold(>=16)={above16}/{active}"
8445            );
8446        }
8447
8448        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8449        Ok(moe_out)
8450    }
8451
8452    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
8453    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
8454    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
8455    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
8456    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
8457    /// expert-sum order identical to the sequential path.
8458    pub(crate) fn moe_ffn_lockstep(
8459        &self,
8460        e: &Engine,
8461        m: &MoeWeights,
8462        zbatch: &CudaSlice<f32>,
8463        mrows: usize,
8464        il: u16,
8465        max_block: usize,
8466    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8467        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8468        let cfg = &self.cfg;
8469        let moe = cfg.moe.as_ref().unwrap();
8470        let n_embd = cfg.n_embd as usize;
8471        let n_expert = moe.expert_count as usize;
8472        let n_used = moe.expert_used_count as usize;
8473        let n_ff_exp = moe.expert_ff_length as usize;
8474        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
8475        let lim_exp = cfg.clamp_exp_at(il as u32);
8476        let lim_shexp = cfg.clamp_shexp_at(il as u32);
8477
8478        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8479        if let Some(sig) = cfg.sigmoid_router() {
8480            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8481        }
8482        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8483            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8484        } else {
8485            Self::moe_route_cfg(
8486                e,
8487                &logits,
8488                mrows,
8489                n_expert,
8490                n_used,
8491                m.active_experts.as_deref(),
8492            )?
8493        };
8494        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8495
8496        // Residency split at whole-expert granularity against the (frozen) cache.
8497        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8498            Ok((0..n_expert)
8499                .map(|ex| {
8500                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8501                        .into_iter()
8502                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8503                })
8504                .collect())
8505        })?;
8506
8507        struct Group {
8508            rows: Vec<i32>,
8509            slots: Vec<i32>,
8510            weights: Vec<f32>,
8511        }
8512        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8513        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8514        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8515            Default::default();
8516        for row in 0..mrows {
8517            for j in 0..n_used {
8518                let ex = sel_all[row * n_used + j] as usize;
8519                let w = w_all[row * n_used + j];
8520                if resident_expert[ex] {
8521                    let group = groups.entry(ex).or_insert_with(|| Group {
8522                        rows: Vec::new(),
8523                        slots: Vec::new(),
8524                        weights: Vec::new(),
8525                    });
8526                    group.rows.push(row as i32);
8527                    group.slots.push(j as i32);
8528                    group.weights.push(w);
8529                } else {
8530                    crate::cpu_experts::record_incomplete_gpu_residency(0);
8531                    cpu_rows[row].push((ex, w));
8532                    cpu_by_expert.entry(ex).or_default().push((row, w));
8533                }
8534            }
8535        }
8536
8537        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
8538        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
8539        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
8540        // order per row differs from the sequential single-call chunk — part of the
8541        // documented lockstep numeric class.
8542        let host_rows = e.dtoh(zbatch)?;
8543        let rows_ok = crate::cpu_experts::rows_supported();
8544        enum CpuPart {
8545            Single { row: usize },
8546            Rows { rows: Vec<usize> },
8547        }
8548        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8549        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8550        if rows_ok {
8551            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8552                .into_iter()
8553                .filter(|(_, rows)| rows.len() >= 2)
8554                .collect();
8555            shared.sort_by_key(|(ex, _)| *ex);
8556            for (ex, mut row_weights) in shared {
8557                row_weights.sort_by_key(|(row, _)| *row);
8558                let inputs: Vec<(&[f32], f32)> = row_weights
8559                    .iter()
8560                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8561                    .collect();
8562                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8563                    .map_err(std::io::Error::other)?;
8564                for &(row, _) in &row_weights {
8565                    rows_served.insert((row, ex));
8566                }
8567                tickets.push((
8568                    CpuPart::Rows {
8569                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
8570                    },
8571                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8572                ));
8573            }
8574        }
8575        for (row, selected) in cpu_rows.iter().enumerate() {
8576            let leftover: Vec<(usize, f32)> = selected
8577                .iter()
8578                .copied()
8579                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8580                .collect();
8581            if leftover.is_empty() {
8582                continue;
8583            }
8584            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8585            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8586                .map_err(std::io::Error::other)?;
8587            tickets.push((
8588                CpuPart::Single { row },
8589                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8590            ));
8591        }
8592
8593        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8594        let mut wbuf = e.zeros(mrows * n_used)?;
8595        let mut order: Vec<usize> = groups.keys().copied().collect();
8596        order.sort_by(|&a, &b| {
8597            groups[&b]
8598                .rows
8599                .len()
8600                .cmp(&groups[&a].rows.len())
8601                .then(a.cmp(&b))
8602        });
8603        for &ex in &order {
8604            let group = &groups[&ex];
8605            let m_e = group.rows.len();
8606            let gl = m.gate_exps.expert_layout(ex);
8607            let ul = m.up_exps.expert_layout(ex);
8608            let dl = m.down_exps.expert_layout(ex);
8609            let row_idx_d = e.htod_i32(&group.rows)?;
8610            let slot_idx_d = e.htod_i32(&group.slots)?;
8611            let dmac = m.down_exps.macro_scale(ex);
8612            let weight_d = if dmac == 1.0 {
8613                e.htod(&group.weights)?
8614            } else {
8615                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8616                e.htod(&scaled)?
8617            };
8618            let mut gathered = e.zeros(m_e * n_embd)?;
8619            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8620            let gv = gathered.slice(0..m_e * n_embd);
8621            let gate = e.with_moe_cache(max_block, |c, eng| {
8622                let slot = c
8623                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8624                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8625                eng.qmatvec_view(
8626                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8627                    0..gl.len,
8628                    &gv,
8629                    m_e,
8630                    m.gate_exps.in_f,
8631                    m.gate_exps.out_f,
8632                    gl.qtype,
8633                    gl.row_bytes,
8634                )
8635            })?;
8636            let up = e.with_moe_cache(max_block, |c, eng| {
8637                let slot = c
8638                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
8639                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8640                eng.qmatvec_view(
8641                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8642                    0..ul.len,
8643                    &gv,
8644                    m_e,
8645                    m.up_exps.in_f,
8646                    m.up_exps.out_f,
8647                    ul.qtype,
8648                    ul.row_bytes,
8649                )
8650            })?;
8651            let mut act = e.zeros(m_e * n_ff_exp)?;
8652            Self::ffn_act_lim(
8653                e,
8654                cfg,
8655                &gate,
8656                &up,
8657                m.gate_exps.macro_scale(ex),
8658                m.up_exps.macro_scale(ex),
8659                lim_exp,
8660                &mut act,
8661                m_e * n_ff_exp,
8662            )?;
8663            let actv = act.slice(0..m_e * n_ff_exp);
8664            let y = e.with_moe_cache(max_block, |c, eng| {
8665                let slot = c
8666                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8667                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8668                eng.qmatvec_view(
8669                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8670                    0..dl.len,
8671                    &actv,
8672                    m_e,
8673                    m.down_exps.in_f,
8674                    m.down_exps.out_f,
8675                    dl.qtype,
8676                    dl.row_bytes,
8677                )
8678            })?;
8679            e.scatter_slot(
8680                &y,
8681                &row_idx_d,
8682                &slot_idx_d,
8683                &weight_d,
8684                &mut slot_buf,
8685                &mut wbuf,
8686                n_embd,
8687                n_used,
8688                m_e,
8689            )?;
8690        }
8691        let mut moe_out = e.zeros(mrows * n_embd)?;
8692        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8693
8694        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
8695        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8696        for (part, ticket) in tickets {
8697            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8698            let mut add_row = |row: usize, chunk: &[f32]| {
8699                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8700                for (accumulator, value) in sum.iter_mut().zip(chunk) {
8701                    *accumulator += value;
8702                }
8703            };
8704            match part {
8705                CpuPart::Single { row } => add_row(row, &cpu_output),
8706                CpuPart::Rows { rows } => {
8707                    for (slot, row) in rows.into_iter().enumerate() {
8708                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8709                    }
8710                }
8711            }
8712        }
8713        for (row, sum) in row_sums.into_iter().enumerate() {
8714            let Some(sum) = sum else { continue };
8715            let cpu_output = e.htod(&sum)?;
8716            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8717            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8718        }
8719
8720        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8721            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8722        {
8723            let n_ff_sh = gate_shexp.out_features();
8724            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8725            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8726            let mut sa = e.zeros(mrows * n_ff_sh)?;
8727            Self::ffn_act_lim(
8728                e,
8729                cfg,
8730                &sg_gate,
8731                &sg_up,
8732                1.0,
8733                1.0,
8734                lim_shexp,
8735                &mut sa,
8736                mrows * n_ff_sh,
8737            )?;
8738            let sh = e.matmul(down_shexp, &sa, mrows)?;
8739            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
8740            // decode matches the single-sequence decode chain bit-for-bit.
8741            let g = match &m.gate_inp_shexp {
8742                Some(gate_inp_shexp) => {
8743                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8744                }
8745                None => e.htod(&vec![1.0f32; mrows])?,
8746            };
8747            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8748        }
8749
8750        Ok(moe_out)
8751    }
8752}
8753
8754// ============================ gemma4 (R8 verified wiring) ==================================
8755// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
8756// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
8757// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
8758// gemma variants after the correctness gate).
8759impl HybridModel {
8760    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
8761    /// Rotary width the gemma-4 config DECLARES for layer `il`, per SWA/global class — the
8762    /// `n_rot` argument of the fused `rms_norm_qkv_rope*` kernels
8763    /// (`Engine::full_width_rope_only`, lane/graph-s-key-exactness-20260819).
8764    ///
8765    /// It exists so the fusion's full-width assumption is checked against the CONFIG instead of
8766    /// against itself: passing `hd` twice would guard nothing. On every gemma-4 artifact both
8767    /// routes set `rope_dims_{global,swa} == key_length_{global,swa}` (GGUF reads
8768    /// `rope.dimension_count(_swa)`, the safetensors route derives them from
8769    /// `global_head_dim`/`head_dim`), which is exactly WHY the fusion is legal here — gemma's
8770    /// "partial" rotary is per-dimension frequency scaling through the `rope_freqs` factors
8771    /// tensor, not an n_rot truncation. An arch that truncates instead (qwen3.5: 64 of 256;
8772    /// step35 full-attn: 64 of 128) must not reach these kernels, and now cannot silently.
8773    pub(crate) fn gemma4_rope_dims(&self, il: usize) -> usize {
8774        let g = self
8775            .cfg
8776            .gemma4
8777            .as_ref()
8778            .expect("gemma4_rope_dims on a non-gemma4 config");
8779        if g.swa_pattern[il] {
8780            g.rope_dims_swa as usize
8781        } else {
8782            g.rope_dims_global as usize
8783        }
8784    }
8785
8786    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8787        let g = self.cfg.gemma4.as_ref().unwrap();
8788        let swa = g.swa_pattern[il];
8789        let hd = if swa {
8790            g.key_length_swa
8791        } else {
8792            g.key_length_global
8793        } as usize;
8794        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
8795        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
8796        // rows exact (softmax over one element) while every later position drifted).
8797        (
8798            hd,
8799            g.head_count_kv[il] as usize,
8800            self.cfg.n_head as usize,
8801            if swa {
8802                g.rope_base_swa
8803            } else {
8804                g.rope_base_global
8805            },
8806            1.0,
8807            swa,
8808        )
8809    }
8810
8811    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
8812    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
8813    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
8814    pub(crate) fn gemma4_suppress(
8815        &self,
8816        e: &Engine,
8817        ld: &mut CudaSlice<f32>,
8818        t: usize,
8819    ) -> Result<(), Box<dyn std::error::Error>> {
8820        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8821            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
8822            // stage as primary, and this tail runs only after the last stage). The assert turns
8823            // that argued invariant into a checked one: any topology violating primary==head
8824            // trips here in debug instead of silently peer-reading a device-0 buffer.
8825            #[cfg(debug_assertions)]
8826            crate::debug_assert_tensor_stream_device(
8827                ids,
8828                &e.stream(),
8829                "gemma4_suppress.suppress_d",
8830            );
8831            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8832        }
8833        Ok(())
8834    }
8835
8836    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
8837    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
8838    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
8839    /// only (v0): attends within `tokens` via the f32 sdpa.
8840    #[allow(clippy::too_many_arguments)]
8841    /// ONE-PROGRAM seam (lane/gemma-restore-exactness-20260819, DEFAULT OFF): route dense
8842    /// gemma SWA prefill attention through the WINDOWED arm at every prompt length instead of
8843    /// switching program at `t > sliding_window`. The door is the measured cause of the
8844    /// prefix-restore exactness violation: below it the SWA layers attend unwindowed on f32
8845    /// operands, above it windowed on bf16 operands, so rows[0..932) of a 932-row prime and of
8846    /// a 1048-row prime differ in 59 of 60 layers. `=1` collapses it to one program so the
8847    /// published prefix KV stops depending on the total prompt length. Off by default because
8848    /// collapsing changes the bytes short prompts get today; the flip is an owner call.
8849    fn gemma_fa_one_program() -> bool {
8850        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8851        *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_FA_ONE_PROGRAM").as_deref() == Ok("1"))
8852    }
8853
8854    fn gemma4_attn_prime(
8855        &self,
8856        e: &Engine,
8857        fa: &crate::hybrid::FullAttnLayer,
8858        il: usize,
8859        h: &CudaSlice<f32>,
8860        pos_d: &CudaSlice<i32>,
8861        t: usize,
8862        cache: Option<&mut Cache>,
8863        island: Option<&CudaSlice<i32>>,
8864    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8865        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8866        let eps = self.cfg.rms_eps;
8867        let aux = self.gemma4_aux.as_ref().unwrap();
8868        let ones = aux.ones(e);
8869        #[cfg(debug_assertions)]
8870        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8871
8872        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
8873        // (h stays borrowed across the triple, so the cache key can't go stale).
8874        e.mmq_act_begin();
8875        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
8876        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8877            let v = e.dtoh(&q0)?;
8878            let nan = v.iter().filter(|x| x.is_nan()).count();
8879            let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8880            eprintln!(
8881                "[g4-prime-trace] L0 q0: nan={nan}/{} amax={amax:.3}",
8882                v.len()
8883            );
8884        }
8885        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
8886        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
8887        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
8888        let v0 = if swa {
8889            e.matmul(&fa.wv, h, t)?
8890        } else {
8891            e.clone_dtod(&k0)?
8892        };
8893        if il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
8894            for (tag, buf) in [("k0", &k0), ("v0", &v0)] {
8895                let v = e.dtoh(buf)?;
8896                let nan = v.iter().filter(|x| x.is_nan()).count();
8897                let amax = v.iter().fold(0f32, |a, x| a.max(x.abs()));
8898                eprintln!(
8899                    "[g4-prime-trace] L0 {tag}: nan={nan}/{} amax={amax:.3}",
8900                    v.len()
8901                );
8902            }
8903        }
8904
8905        let mut q = e.uninit(t * nh * hd)?;
8906        let mut k = e.uninit(t * nkv * hd)?;
8907        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
8908        let mut v = e.uninit(t * nkv * hd)?;
8909        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
8910        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
8911        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
8912        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8913        // Island primes take the mask-capable naive kernel below; keep the operands f32
8914        // (the bf16 FA emit path has no island consumer). Text-only keeps emit unchanged.
8915        let emit = island.is_none()
8916            && t >= 16
8917            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8918            && *EMIT.get_or_init(|| {
8919                std::env::var("MEMRA_FA_EMIT")
8920                    .map(|s| s != "0")
8921                    .unwrap_or(true)
8922            });
8923        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8924        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8925        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8926        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
8927        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
8928        let v_f16 = emit
8929            && crate::fa_f16pv_on()
8930            && match hd {
8931                512 => true,
8932                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8933                _ => false,
8934            };
8935        if emit {
8936            e.rms_norm_qkv_w4b(
8937                &q0,
8938                &k0,
8939                &v0,
8940                fa.q_norm.float_data(),
8941                fa.k_norm.float_data(),
8942                ones,
8943                &mut q,
8944                &mut k,
8945                &mut v,
8946                &mut vb,
8947                hd,
8948                nh * t,
8949                nkv * t,
8950                eps,
8951                v_f16,
8952            )?;
8953        } else {
8954            e.rms_norm_qkv(
8955                &q0,
8956                &k0,
8957                &v0,
8958                fa.q_norm.float_data(),
8959                fa.k_norm.float_data(),
8960                ones,
8961                &mut q,
8962                &mut k,
8963                &mut v,
8964                hd,
8965                nh * t,
8966                nkv * t,
8967                eps,
8968            )?;
8969        }
8970
8971        let ff = if swa {
8972            None
8973        } else {
8974            Some(
8975                aux.rope_freqs(e)
8976                    .expect("gemma4 global rope needs rope_freqs.weight"),
8977            )
8978        };
8979        #[cfg(debug_assertions)]
8980        if let Some(ff) = ff {
8981            crate::debug_assert_tensor_stream_device(
8982                ff,
8983                &e.stream(),
8984                "gemma4_attn_prime.rope_freqs",
8985            );
8986        }
8987        if emit {
8988            e.rope_neox2_bf16e(
8989                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8990            )?;
8991        } else {
8992            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8993        }
8994
8995        if let Some(cache) = cache {
8996            let kvl = cache.kv[il].as_mut().unwrap();
8997            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8998            e.append_kv_quantized_rows(
8999                &k,
9000                &v,
9001                &mut kvl.k,
9002                &mut kvl.v,
9003                kvl.len,
9004                t,
9005                kvl.kv_dim_k,
9006                kvl.kv_dim_v,
9007                kvl.k_tok_bytes,
9008                kvl.v_tok_bytes,
9009                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
9010            )?;
9011            kvl.len += t;
9012        }
9013        let mut attn = e.zeros(t * nh * hd)?;
9014        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
9015        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
9016        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
9017        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
9018        if let Some(span) = island {
9019            // Masked-prefill arm: every layer routes through the island-aware naive
9020            // kernel (correctness-first, same posture as the vision tower v1). The
9021            // window argument keeps the R6 shortcut: 0 while the prompt fits the
9022            // window, the real window beyond it.
9023            let w = if swa && t > win { win } else { 0 };
9024            e.sdpa_naive_island(&q, &k, &v, &mut attn, span, hd, nh, nkv, t, t, scale, w)?;
9025        } else if swa && (t > win || Self::gemma_fa_one_program()) {
9026            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
9027                if emit {
9028                    e.fa_prefill_w_pre(
9029                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
9030                    )?;
9031                } else {
9032                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9033                }
9034            } else {
9035                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9036            }
9037        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
9038            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9039        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
9040            if emit {
9041                e.fa_prefill_hd512_pre(
9042                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
9043                )?;
9044            } else {
9045                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9046            }
9047        } else {
9048            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9049        }
9050        Ok(e.matmul(&fa.wo, &attn, t)?)
9051    }
9052
9053    /// Back-compat wrapper (pure prefill, no cache).
9054    fn gemma4_attn(
9055        &self,
9056        e: &Engine,
9057        fa: &crate::hybrid::FullAttnLayer,
9058        il: usize,
9059        h: &CudaSlice<f32>,
9060        pos_d: &CudaSlice<i32>,
9061        t: usize,
9062    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9063        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None, None)
9064    }
9065
9066    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
9067    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
9068    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
9069    /// the q8z epilogue is quantize_q8_1 verbatim).
9070    fn gemma4_moe_q8(
9071        &self,
9072        e: &Engine,
9073        m: &crate::hybrid::MoeWeights,
9074        bits: &crate::hybrid::Gemma4MoeBits,
9075        mq: &(CudaSlice<i8>, CudaSlice<f32>),
9076        router_in: &CudaSlice<f32>,
9077        t: usize,
9078    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9079        let cfg = &self.cfg;
9080        let moe = cfg.moe.as_ref().unwrap();
9081        let n_embd = cfg.n_embd as usize;
9082        let n_expert = moe.expert_count as usize;
9083        let n_used = moe.expert_used_count as usize;
9084        let n_ff_exp = moe.expert_ff_length as usize;
9085        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
9086        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
9087        // the pair's 12us is kernel time, not launch gaps.
9088        let logits = if crate::router_kernel_on() {
9089            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9090        } else {
9091            e.matmul(&m.gate_inp, router_in, t)?
9092        };
9093        let dev = m.dev_exps.as_ref().unwrap();
9094        let (sel_d, w_d) =
9095            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9096        let (zq, zd) = mq;
9097        if t == 1 {
9098            let selv = sel_d.slice(0..n_used);
9099            let wv = w_d.slice(0..n_used);
9100            let act = e.moe_gate_up_gelu8_dev_q8(
9101                &dev.ptr_row,
9102                &selv,
9103                zq,
9104                zd,
9105                n_embd,
9106                n_ff_exp,
9107                n_used,
9108                n_expert,
9109                m.gate_exps.qtype,
9110                m.up_exps.qtype,
9111                m.gate_exps.row_bytes,
9112                m.up_exps.row_bytes,
9113            )?;
9114            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9115            let mut moe_out = e.uninit(n_embd)?;
9116            e.moe_down8_fma_dev_q8(
9117                &dev.ptr_row,
9118                &selv,
9119                &wv,
9120                &aq2,
9121                &ad2,
9122                &mut moe_out.slice_mut(0..n_embd),
9123                n_ff_exp,
9124                n_embd,
9125                n_used,
9126                n_expert,
9127                m.down_exps.qtype,
9128                m.down_exps.row_bytes,
9129            )?;
9130            return Ok(moe_out);
9131        }
9132        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9133        let act = if csr {
9134            e.moe_gate_up_gelu8_dev_q8_csr(
9135                &dev.ptr_row,
9136                &sel_d,
9137                zq,
9138                zd,
9139                t * n_used,
9140                n_embd,
9141                n_ff_exp,
9142                n_used,
9143                n_expert,
9144                m.gate_exps.qtype,
9145                m.up_exps.qtype,
9146                m.gate_exps.row_bytes,
9147                m.up_exps.row_bytes,
9148            )?
9149        } else {
9150            e.moe_gate_up_gelu8_dev_q8_rows(
9151                &dev.ptr_row,
9152                &sel_d,
9153                zq,
9154                zd,
9155                t,
9156                n_embd,
9157                n_ff_exp,
9158                n_used,
9159                n_expert,
9160                m.gate_exps.qtype,
9161                m.up_exps.qtype,
9162                m.gate_exps.row_bytes,
9163                m.up_exps.row_bytes,
9164            )?
9165        };
9166        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9167        let mut moe_out = e.uninit(t * n_embd)?;
9168        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
9169        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
9170        e.moe_down8_fma_dev_q8_rows_g(
9171            &dev.ptr_row,
9172            &sel_d,
9173            &w_d,
9174            &aq2,
9175            &ad2,
9176            &mut moe_out,
9177            t,
9178            n_ff_exp,
9179            n_embd,
9180            n_used,
9181            n_expert,
9182            m.down_exps.qtype,
9183            m.down_exps.row_bytes,
9184        )?;
9185        Ok(moe_out)
9186    }
9187
9188    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
9189    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
9190    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
9191    fn gemma4_moe(
9192        &self,
9193        e: &Engine,
9194        m: &crate::hybrid::MoeWeights,
9195        bits: &crate::hybrid::Gemma4MoeBits,
9196        moe_in: &CudaSlice<f32>,
9197        router_in: &CudaSlice<f32>,
9198        t: usize,
9199    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9200        let cfg = &self.cfg;
9201        let moe = cfg.moe.as_ref().unwrap();
9202        let n_embd = cfg.n_embd as usize;
9203        let n_expert = moe.expert_count as usize;
9204        let n_used = moe.expert_used_count as usize;
9205        let n_ff_exp = moe.expert_ff_length as usize;
9206
9207        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
9208        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
9209        // batched matmul only at real prefill.
9210        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
9211            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
9212        } else {
9213            e.matmul(&m.gate_inp, router_in, t)?
9214        };
9215
9216        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
9217        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
9218        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
9219        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
9220        if t < PRIME_MIN_T
9221            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9222            && expert_dp4a_supported(m.gate_exps.qtype)
9223            && expert_dp4a_supported(m.up_exps.qtype)
9224            && expert_dp4a_supported(m.down_exps.qtype)
9225            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9226        {
9227            let dev = m.dev_exps.as_ref().unwrap();
9228            let (sel_d, w_d) =
9229                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
9230            if t == 1 {
9231                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
9232                let selv = sel_d.slice(0..n_used);
9233                let wv = w_d.slice(0..n_used);
9234                let act = e.moe_gate_up_gelu8_dev_q8(
9235                    &dev.ptr_row,
9236                    &selv,
9237                    &zq,
9238                    &zd,
9239                    n_embd,
9240                    n_ff_exp,
9241                    n_used,
9242                    n_expert,
9243                    m.gate_exps.qtype,
9244                    m.up_exps.qtype,
9245                    m.gate_exps.row_bytes,
9246                    m.up_exps.row_bytes,
9247                )?;
9248                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
9249                let mut moe_out = e.uninit(n_embd)?;
9250                e.moe_down8_fma_dev_q8(
9251                    &dev.ptr_row,
9252                    &selv,
9253                    &wv,
9254                    &aq2,
9255                    &ad2,
9256                    &mut moe_out.slice_mut(0..n_embd),
9257                    n_ff_exp,
9258                    n_embd,
9259                    n_used,
9260                    n_expert,
9261                    m.down_exps.qtype,
9262                    m.down_exps.row_bytes,
9263                )?;
9264                return Ok(moe_out);
9265            }
9266            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
9267            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
9268            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
9269            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
9270            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9271            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
9272            let act = if csr {
9273                e.moe_gate_up_gelu8_dev_q8_csr(
9274                    &dev.ptr_row,
9275                    &sel_d,
9276                    &zq,
9277                    &zd,
9278                    t * n_used,
9279                    n_embd,
9280                    n_ff_exp,
9281                    n_used,
9282                    n_expert,
9283                    m.gate_exps.qtype,
9284                    m.up_exps.qtype,
9285                    m.gate_exps.row_bytes,
9286                    m.up_exps.row_bytes,
9287                )?
9288            } else {
9289                e.moe_gate_up_gelu8_dev_q8_rows(
9290                    &dev.ptr_row,
9291                    &sel_d,
9292                    &zq,
9293                    &zd,
9294                    t,
9295                    n_embd,
9296                    n_ff_exp,
9297                    n_used,
9298                    n_expert,
9299                    m.gate_exps.qtype,
9300                    m.up_exps.qtype,
9301                    m.gate_exps.row_bytes,
9302                    m.up_exps.row_bytes,
9303                )?
9304            };
9305            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9306            let mut moe_out = e.uninit(t * n_embd)?;
9307            e.moe_down8_fma_dev_q8_rows_g(
9308                &dev.ptr_row,
9309                &sel_d,
9310                &w_d,
9311                &aq2,
9312                &ad2,
9313                &mut moe_out,
9314                t,
9315                n_ff_exp,
9316                n_embd,
9317                n_used,
9318                n_expert,
9319                m.down_exps.qtype,
9320                m.down_exps.row_bytes,
9321            )?;
9322            return Ok(moe_out);
9323        }
9324
9325        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9326        for (i, &sx) in sel_all.iter().enumerate() {
9327            w_all[i] *= bits.per_expert_scale[sx as usize];
9328        }
9329
9330        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
9331        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
9332        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
9333        if t >= PRIME_MIN_T
9334            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9335            && expert_dp4a_supported(m.gate_exps.qtype)
9336            && expert_dp4a_supported(m.up_exps.qtype)
9337            && expert_dp4a_supported(m.down_exps.qtype)
9338            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9339        {
9340            let dev = m.dev_exps.as_ref().unwrap();
9341            let n_pairs = t * n_used;
9342            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9343            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9344            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9345            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9346            let pt = e.htod_i32(&pair_tok)?;
9347            let pw = e.htod(&w_all)?;
9348            let toff = e.htod_i32(&tok_off)?;
9349            let tids = e.htod_i32(&tok_ids)?;
9350            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9351            for p in 0..n_pairs {
9352                by_ex[pair_ex[p] as usize].push(p as i32);
9353            }
9354            let mut ex_ids: Vec<i32> = Vec::new();
9355            let mut ex_off: Vec<i32> = vec![0];
9356            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9357            for (ex, list) in by_ex.iter().enumerate() {
9358                if list.is_empty() {
9359                    continue;
9360                }
9361                ex_ids.push(ex as i32);
9362                ex_pairs.extend_from_slice(list);
9363                ex_off.push(ex_pairs.len() as i32);
9364            }
9365            let n_active = ex_ids.len();
9366            let exi = e.htod_i32(&ex_ids)?;
9367            let exo = e.htod_i32(&ex_off)?;
9368            let exp_d = e.htod_i32(&ex_pairs)?;
9369            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
9370            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
9371            // end-to-end (gelu is elementwise), one row permute before the scatter. The
9372            // ragged down k (704) needs no padding here — cublas takes any k.
9373            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
9374            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
9375            // Hopper default — see moe_f16g_gemma_on.
9376            if crate::moe_f16g_gemma_on()
9377                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9378                && f16g_proj_ok(m.up_exps.qtype, n_embd)
9379                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9380            {
9381                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9382                let csr_tok_d = e.htod_i32(&csr_tok)?;
9383                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9384                let g_csr = e.moe_f16_grouped(
9385                    &dev.ptr_row,
9386                    0,
9387                    n_expert,
9388                    &exi,
9389                    &ex_off,
9390                    &exo,
9391                    &z_f16,
9392                    &z_s,
9393                    n_embd,
9394                    n_ff_exp,
9395                    n_active,
9396                    n_pairs,
9397                    m.gate_exps.qtype,
9398                    m.gate_exps.row_bytes,
9399                )?;
9400                let u_csr = e.moe_f16_grouped(
9401                    &dev.ptr_row,
9402                    1,
9403                    n_expert,
9404                    &exi,
9405                    &ex_off,
9406                    &exo,
9407                    &z_f16,
9408                    &z_s,
9409                    n_embd,
9410                    n_ff_exp,
9411                    n_active,
9412                    n_pairs,
9413                    m.up_exps.qtype,
9414                    m.up_exps.row_bytes,
9415                )?;
9416                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9417                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9418                let d_csr = e.moe_f16_grouped(
9419                    &dev.ptr_row,
9420                    2,
9421                    n_expert,
9422                    &exi,
9423                    &ex_off,
9424                    &exo,
9425                    &a_f16,
9426                    &a_s,
9427                    n_ff_exp,
9428                    n_embd,
9429                    n_active,
9430                    n_pairs,
9431                    m.down_exps.qtype,
9432                    m.down_exps.row_bytes,
9433                )?;
9434                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9435                let mut moe_out = e.uninit(t * n_embd)?;
9436                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9437                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9438                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9439                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9440                    eprintln!(
9441                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
9442                        scan(&yd),
9443                        scan(&mo)
9444                    );
9445                }
9446                return Ok(moe_out);
9447            }
9448            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
9449            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
9450            let mma =
9451                n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9452            let (gate, up) = if mma {
9453                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9454                (
9455                    e.mmq_iq_experts(
9456                        &dev.ptr_row,
9457                        0,
9458                        n_expert,
9459                        &exi,
9460                        &exo,
9461                        &exp_d,
9462                        &pt,
9463                        &z_scr,
9464                        n_embd,
9465                        n_ff_exp,
9466                        n_active,
9467                        n_pairs,
9468                        t,
9469                        m.gate_exps.qtype,
9470                        m.gate_exps.row_bytes,
9471                    )?,
9472                    e.mmq_iq_experts(
9473                        &dev.ptr_row,
9474                        1,
9475                        n_expert,
9476                        &exi,
9477                        &exo,
9478                        &exp_d,
9479                        &pt,
9480                        &z_scr,
9481                        n_embd,
9482                        n_ff_exp,
9483                        n_active,
9484                        n_pairs,
9485                        t,
9486                        m.up_exps.qtype,
9487                        m.up_exps.row_bytes,
9488                    )?,
9489                )
9490            } else {
9491                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9492                (
9493                    e.moe_pairs_matvec_q8_dec(
9494                        &dev.ptr_row,
9495                        0,
9496                        &exi,
9497                        &exo,
9498                        &exp_d,
9499                        &pt,
9500                        &zq,
9501                        &zd,
9502                        n_embd,
9503                        n_ff_exp,
9504                        n_expert,
9505                        n_active,
9506                        n_pairs,
9507                        m.gate_exps.qtype,
9508                        m.gate_exps.row_bytes,
9509                    )?,
9510                    e.moe_pairs_matvec_q8_dec(
9511                        &dev.ptr_row,
9512                        1,
9513                        &exi,
9514                        &exo,
9515                        &exp_d,
9516                        &pt,
9517                        &zq,
9518                        &zd,
9519                        n_embd,
9520                        n_ff_exp,
9521                        n_expert,
9522                        n_active,
9523                        n_pairs,
9524                        m.up_exps.qtype,
9525                        m.up_exps.row_bytes,
9526                    )?,
9527                )
9528            };
9529            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9530            let pself = e.htod_i32(&pair_self)?;
9531            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
9532            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
9533            // to the 256-val superblock (768) while the act quantizer's zero padding
9534            // makes every padded-k product exactly zero (weight overread bytes multiply
9535            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
9536            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
9537            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
9538            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
9539            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
9540            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
9541            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
9542            let y_down = if mma {
9543                let in_pad = n_ff_exp.div_ceil(256) * 256;
9544                let a_scr = if crate::moe_fuse_actq_on() {
9545                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9546                } else {
9547                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9548                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9549                };
9550                e.mmq_iq_experts(
9551                    &dev.ptr_row,
9552                    2,
9553                    n_expert,
9554                    &exi,
9555                    &exo,
9556                    &exp_d,
9557                    &pself,
9558                    &a_scr,
9559                    in_pad,
9560                    n_embd,
9561                    n_active,
9562                    n_pairs,
9563                    n_pairs,
9564                    m.down_exps.qtype,
9565                    m.down_exps.row_bytes,
9566                )?
9567            } else {
9568                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9569                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9570                e.moe_pairs_matvec_q8_dec(
9571                    &dev.ptr_row,
9572                    2,
9573                    &exi,
9574                    &exo,
9575                    &exp_d,
9576                    &pself,
9577                    &aq2,
9578                    &ad2,
9579                    n_ff_exp,
9580                    n_embd,
9581                    n_expert,
9582                    n_active,
9583                    n_pairs,
9584                    m.down_exps.qtype,
9585                    m.down_exps.row_bytes,
9586                )?
9587            };
9588            let mut moe_out = e.uninit(t * n_embd)?;
9589            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9590            return Ok(moe_out);
9591        }
9592
9593        let g_len = m.gate_exps.expert_stride;
9594        let u_len = m.up_exps.expert_stride;
9595        let d_len = m.down_exps.expert_stride;
9596        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
9597        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
9598        // the spill fallback.
9599        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9600        let (mut sg, mut su, mut sd) = if dev.is_some() {
9601            (None, None, None)
9602        } else {
9603            (
9604                Some(e.alloc_u8_uninit(g_len)?),
9605                Some(e.alloc_u8_uninit(u_len)?),
9606                Some(e.alloc_u8_uninit(d_len)?),
9607            )
9608        };
9609        let mut moe_out = e.zeros(t * n_embd)?;
9610        for tok in 0..t {
9611            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9612            let w = &w_all[tok * n_used..(tok + 1) * n_used];
9613            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9614            for (j, &ex) in sel.iter().enumerate() {
9615                let ex = ex as usize;
9616                let gate = match dev {
9617                    Some(d) => e.qmatvec_view(
9618                        &d.gate,
9619                        ex * g_len..(ex + 1) * g_len,
9620                        &zt,
9621                        1,
9622                        m.gate_exps.in_f,
9623                        m.gate_exps.out_f,
9624                        m.gate_exps.qtype,
9625                        m.gate_exps.row_bytes,
9626                    )?,
9627                    None => {
9628                        let sg = sg.as_mut().unwrap();
9629                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9630                        e.qmatvec_view(
9631                            sg,
9632                            0..g_len,
9633                            &zt,
9634                            1,
9635                            m.gate_exps.in_f,
9636                            m.gate_exps.out_f,
9637                            m.gate_exps.qtype,
9638                            m.gate_exps.row_bytes,
9639                        )?
9640                    }
9641                };
9642                let up = match dev {
9643                    Some(d) => e.qmatvec_view(
9644                        &d.up,
9645                        ex * u_len..(ex + 1) * u_len,
9646                        &zt,
9647                        1,
9648                        m.up_exps.in_f,
9649                        m.up_exps.out_f,
9650                        m.up_exps.qtype,
9651                        m.up_exps.row_bytes,
9652                    )?,
9653                    None => {
9654                        let su = su.as_mut().unwrap();
9655                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9656                        e.qmatvec_view(
9657                            su,
9658                            0..u_len,
9659                            &zt,
9660                            1,
9661                            m.up_exps.in_f,
9662                            m.up_exps.out_f,
9663                            m.up_exps.qtype,
9664                            m.up_exps.row_bytes,
9665                        )?
9666                    }
9667                };
9668                let mut act = e.uninit(n_ff_exp)?;
9669                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9670                let actv = act.slice(0..n_ff_exp);
9671                let y = match dev {
9672                    Some(d) => e.qmatvec_view(
9673                        &d.down,
9674                        ex * d_len..(ex + 1) * d_len,
9675                        &actv,
9676                        1,
9677                        m.down_exps.in_f,
9678                        m.down_exps.out_f,
9679                        m.down_exps.qtype,
9680                        m.down_exps.row_bytes,
9681                    )?,
9682                    None => {
9683                        let sd = sd.as_mut().unwrap();
9684                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9685                        e.qmatvec_view(
9686                            sd,
9687                            0..d_len,
9688                            &actv,
9689                            1,
9690                            m.down_exps.in_f,
9691                            m.down_exps.out_f,
9692                            m.down_exps.qtype,
9693                            m.down_exps.row_bytes,
9694                        )?
9695                    }
9696                };
9697                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9698                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9699            }
9700        }
9701        Ok(moe_out)
9702    }
9703
9704    /// One gemma4 trunk layer (R8): x -> x_next.
9705    fn gemma4_layer(
9706        &self,
9707        e: &Engine,
9708        il: usize,
9709        layer: &crate::hybrid::HybridLayer,
9710        x: &CudaSlice<f32>,
9711        pos_d: &CudaSlice<i32>,
9712        t: usize,
9713    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9714        let n_embd = self.cfg.n_embd as usize;
9715        let eps = self.cfg.rms_eps;
9716
9717        let mut h = e.zeros(t * n_embd)?;
9718        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9719        let Mixer::Full(fa) = &layer.mixer else {
9720            panic!("gemma4 layer {il} not full-attn")
9721        };
9722        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9723        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
9724        let mut cur = e.zeros(t * n_embd)?;
9725        e.rms_norm(
9726            &o,
9727            layer.post_attn_norm.float_data(),
9728            &mut cur,
9729            n_embd,
9730            t,
9731            eps,
9732        )?;
9733        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9734    }
9735
9736    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
9737    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
9738    /// layer scale — shared verbatim by the prefill, decode and verify paths.
9739    fn gemma4_layer_tail_add(
9740        &self,
9741        e: &Engine,
9742        layer: &crate::hybrid::HybridLayer,
9743        cur: &CudaSlice<f32>,
9744        x: &CudaSlice<f32>,
9745        t: usize,
9746    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9747        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9748    }
9749
9750    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
9751    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
9752    fn gemma4_layer_tail_add_n(
9753        &self,
9754        e: &Engine,
9755        layer: &crate::hybrid::HybridLayer,
9756        cur: &CudaSlice<f32>,
9757        x: &CudaSlice<f32>,
9758        t: usize,
9759        next_norm: Option<&CudaSlice<f32>>,
9760    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9761        let n_embd = self.cfg.n_embd as usize;
9762        let bits = layer.gemma4.as_ref().unwrap();
9763        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9764        let mut xn = e.uninit(t * n_embd)?;
9765        match next_norm {
9766            Some(w) => {
9767                let mut hn = e.uninit(t * n_embd)?;
9768                e.add_scale_rms_norm(
9769                    &sn,
9770                    &attn_out,
9771                    bits.layer_scale,
9772                    w,
9773                    &mut xn,
9774                    &mut hn,
9775                    n_embd,
9776                    t,
9777                    self.cfg.rms_eps,
9778                )?;
9779                Ok((xn, Some(hn)))
9780            }
9781            None => {
9782                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9783                Ok((xn, None))
9784            }
9785        }
9786    }
9787
9788    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
9789    /// norm — returns (sn, attn_out) for the closing add+scale variants.
9790    fn gemma4_layer_tail_core(
9791        &self,
9792        e: &Engine,
9793        layer: &crate::hybrid::HybridLayer,
9794        cur: &CudaSlice<f32>,
9795        x: &CudaSlice<f32>,
9796        t: usize,
9797    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9798        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9799    }
9800
9801    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
9802    /// means `cur` is the RAW attention output and the dense entry runs
9803    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
9804    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
9805    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
9806    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
9807    fn gemma4_layer_tail_core_pn(
9808        &self,
9809        e: &Engine,
9810        layer: &crate::hybrid::HybridLayer,
9811        cur: &CudaSlice<f32>,
9812        x: &CudaSlice<f32>,
9813        t: usize,
9814        pre_norm: Option<&CudaSlice<f32>>,
9815        defer_post_norm: bool,
9816    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9817        let n_embd = self.cfg.n_embd as usize;
9818        let eps = self.cfg.rms_eps;
9819        let bits = layer.gemma4.as_ref().unwrap();
9820
9821        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
9822        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
9823        let Some(mbits) = bits.moe_bits.as_ref() else {
9824            let crate::hybrid::Ffn::Dense {
9825                ffn_gate,
9826                ffn_up,
9827                ffn_down,
9828            } = &layer.ffn
9829            else {
9830                panic!("gemma4 dense layer without Dense ffn")
9831            };
9832            let mut attn_out = e.uninit(t * n_embd)?;
9833            let mut zsh = e.uninit(t * n_embd)?;
9834            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
9835            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
9836            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9837            match pre_norm {
9838                Some(wa) if t == 1 => {
9839                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
9840                        cur,
9841                        wa,
9842                        x,
9843                        bits.ffn_norm.float_data(),
9844                        &mut attn_out,
9845                        &mut zsh,
9846                        n_embd,
9847                        t,
9848                        eps,
9849                    )?);
9850                }
9851                Some(wa) => e.rms_pre_add_rms_norm(
9852                    cur,
9853                    wa,
9854                    x,
9855                    bits.ffn_norm.float_data(),
9856                    &mut attn_out,
9857                    &mut zsh,
9858                    n_embd,
9859                    t,
9860                    eps,
9861                )?,
9862                None => e.add_rms_norm(
9863                    cur,
9864                    x,
9865                    bits.ffn_norm.float_data(),
9866                    &mut attn_out,
9867                    &mut zsh,
9868                    n_embd,
9869                    t,
9870                    eps,
9871                )?,
9872            }
9873            let n_ff = ffn_gate.out_features();
9874            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
9875            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
9876            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
9877            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
9878            // rescue segment C — the megakernel front is closed for the dense tail.
9879            let (gate, up) = if t == 1 {
9880                let (zq, zd) = match zpair {
9881                    Some(p) => p,
9882                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9883                };
9884                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9885                    Some(p) => p,
9886                    // NVFP4mix dense trunk: gate/up are both NVFP4 (down stays Q8_0).
9887                    None => match e.matmul_nvfp4_fused2(ffn_gate, ffn_up, &zq, &zd, 1)? {
9888                        Some(p) => p,
9889                        None => (
9890                            e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9891                            e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9892                        ),
9893                    },
9894                }
9895            } else {
9896                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
9897                // launch for the verify's gate+up — the up segment's blocks fill SMs as
9898                // the gate segment drains (the launch-tail mechanism behind the b-tier
9899                // plateau; first positive after six falsified in-kernel variants).
9900                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9901                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9902                let fused = if f2b {
9903                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9904                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9905                } else {
9906                    None
9907                };
9908                match fused {
9909                    Some(p) => p,
9910                    None => {
9911                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
9912                        e.mmq_act_begin();
9913                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9914                    }
9915                }
9916            };
9917            let mut act = e.uninit(t * n_ff)?;
9918            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
9919            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
9920            let f0 = if e.uses_q8_1_fast(ffn_down) {
9921                let upv = e.view(&up, t * n_ff);
9922                let up_all = upv.slice(0..t * n_ff);
9923                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9924                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9925            } else {
9926                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9927                e.matmul(ffn_down, &act, t)?
9928            };
9929            if defer_post_norm {
9930                return Ok((f0, attn_out));
9931            }
9932            let mut sn = e.uninit(t * n_embd)?;
9933            e.rms_norm(
9934                &f0,
9935                bits.post_ffw_norm.float_data(),
9936                &mut sn,
9937                n_embd,
9938                t,
9939                eps,
9940            )?;
9941            return Ok((sn, attn_out));
9942        };
9943
9944        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9945        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
9946        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
9947        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
9948        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
9949        let mut attn_out = e.uninit(t * n_embd)?;
9950        let mut router_in = e.uninit(t * n_embd)?;
9951        let fast_moe = match &layer.ffn {
9952            crate::hybrid::Ffn::Moe(m) => {
9953                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9954                    && expert_dp4a_supported(m.gate_exps.qtype)
9955                    && expert_dp4a_supported(m.up_exps.qtype)
9956                    && expert_dp4a_supported(m.down_exps.qtype)
9957                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9958            }
9959            _ => false,
9960        };
9961        let q8z = t < PRIME_MIN_T && fast_moe;
9962        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9963            let (z0, m2) = e.add_rms_norm3_q8z(
9964                cur,
9965                x,
9966                bits.ffn_norm.float_data(),
9967                &mbits.router_scale_pre,
9968                mbits.pre_ffw_norm_2.float_data(),
9969                &mut attn_out,
9970                &mut router_in,
9971                n_embd,
9972                t,
9973                eps,
9974            )?;
9975            (None, Some(z0), Some(m2))
9976        } else {
9977            let mut zsh = e.uninit(t * n_embd)?;
9978            let mut moe_in = e.uninit(t * n_embd)?;
9979            e.add_rms_norm3(
9980                cur,
9981                x,
9982                bits.ffn_norm.float_data(),
9983                &mbits.router_scale_pre,
9984                mbits.pre_ffw_norm_2.float_data(),
9985                &mut attn_out,
9986                &mut zsh,
9987                &mut router_in,
9988                &mut moe_in,
9989                n_embd,
9990                t,
9991                eps,
9992            )?;
9993            (Some((zsh, moe_in)), None, None)
9994        };
9995        let attn_out2 = attn_out;
9996        #[allow(unused_variables)]
9997        let attn_out = &attn_out2;
9998        let n_ff = mbits.shared_gate.out_features();
9999        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
10000            if t == 1 {
10001                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
10002                    Some(p) => p,
10003                    None => match e.matmul_nvfp4_fused2(
10004                        &mbits.shared_gate,
10005                        &mbits.shared_up,
10006                        zq,
10007                        zd,
10008                        1,
10009                    )? {
10010                        Some(p) => p,
10011                        None => {
10012                            let h0 = e.zeros(0)?;
10013                            (
10014                                e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
10015                                e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
10016                            )
10017                        }
10018                    },
10019                }
10020            } else {
10021                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
10022                let h0 = e.zeros(0)?;
10023                (
10024                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
10025                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
10026                )
10027            }
10028        } else {
10029            let (zsh, _) = zsh_f32.as_ref().unwrap();
10030            (
10031                e.matmul(&mbits.shared_gate, zsh, t)?,
10032                e.matmul(&mbits.shared_up, zsh, t)?,
10033            )
10034        };
10035        let mut act = e.uninit(t * n_ff)?;
10036        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
10037        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
10038        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
10039            panic!("gemma4 layer not MoE")
10040        };
10041        let moe0 = match (&moe_q8, &zsh_f32) {
10042            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
10043            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
10044            _ => unreachable!(),
10045        };
10046        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
10047        let mut mlp = e.uninit(t * n_embd)?;
10048        let mut moe = e.uninit(t * n_embd)?;
10049        e.rms_norm2x(
10050            &mlp0,
10051            &moe0,
10052            mbits.post_ffw_norm_1.float_data(),
10053            mbits.post_ffw_norm_2.float_data(),
10054            &mut mlp,
10055            &mut moe,
10056            n_embd,
10057            t,
10058            eps,
10059        )?;
10060
10061        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
10062        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
10063        let mut sum = e.uninit(t * n_embd)?;
10064        let mut sn = e.uninit(t * n_embd)?;
10065        e.add_rms_norm(
10066            &mlp,
10067            &moe,
10068            bits.post_ffw_norm.float_data(),
10069            &mut sum,
10070            &mut sn,
10071            n_embd,
10072            t,
10073            eps,
10074        )?;
10075        Ok((sn, attn_out2))
10076    }
10077
10078    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
10079    /// pn-fold front of `gemma4_layer_tail_add_nq` (MEMRA_G4_PNFOLD, GAP-DIAGNOSIS
10080    /// verdict 7): takes the RAW attention output `o` and, on the dense trunk, folds
10081    /// post_attn_norm into the tail entry and post_ffw_norm into the residual exit —
10082    /// the E4B glue chain (13594) backported to the decode/verify trio. Kills the
10083    /// standalone rms_norm(o) + rms_norm(f0) (+ffn quantize) launches per layer.
10084    /// BITS-CHANGING vs the two-launch chain; every caller rides this ONE front so
10085    /// decode == verify == graph parity holds by construction at either seam value.
10086    /// Non-dense (26B MoE) and seam-off fall to the unfused chain unchanged.
10087    pub(crate) fn gemma4_layer_tail_add_nq_pn(
10088        &self,
10089        e: &Engine,
10090        layer: &crate::hybrid::HybridLayer,
10091        o: &CudaSlice<f32>,
10092        x: &CudaSlice<f32>,
10093        t: usize,
10094        next_norm: Option<&CudaSlice<f32>>,
10095    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10096    {
10097        let n_embd = self.cfg.n_embd as usize;
10098        let eps = self.cfg.rms_eps;
10099        let bits = layer.gemma4.as_ref().unwrap();
10100        if Engine::g4_pnfold_on() && matches!(layer.ffn, crate::hybrid::Ffn::Dense { .. }) {
10101            let (f0, attn_out) = self.gemma4_layer_tail_core_pn(
10102                e,
10103                layer,
10104                o,
10105                x,
10106                t,
10107                Some(layer.post_attn_norm.float_data()),
10108                true,
10109            )?;
10110            let mut xn = e.uninit(t * n_embd)?;
10111            return match next_norm {
10112                Some(w) => {
10113                    let pair = e.rms_pre_add_scale_rms_norm_q8_1(
10114                        &f0,
10115                        bits.post_ffw_norm.float_data(),
10116                        &attn_out,
10117                        bits.layer_scale,
10118                        w,
10119                        &mut xn,
10120                        n_embd,
10121                        t,
10122                        eps,
10123                    )?;
10124                    Ok((xn, Some(pair)))
10125                }
10126                None => {
10127                    let mut sn = e.uninit(t * n_embd)?;
10128                    e.rms_norm(
10129                        &f0,
10130                        bits.post_ffw_norm.float_data(),
10131                        &mut sn,
10132                        n_embd,
10133                        t,
10134                        eps,
10135                    )?;
10136                    e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10137                    Ok((xn, None))
10138                }
10139            };
10140        }
10141        let mut cur = e.uninit(t * n_embd)?;
10142        e.rms_norm(
10143            o,
10144            layer.post_attn_norm.float_data(),
10145            &mut cur,
10146            n_embd,
10147            t,
10148            eps,
10149        )?;
10150        self.gemma4_layer_tail_add_nq(e, layer, &cur, x, t, next_norm)
10151    }
10152
10153    pub(crate) fn gemma4_layer_tail_add_nq(
10154        &self,
10155        e: &Engine,
10156        layer: &crate::hybrid::HybridLayer,
10157        cur: &CudaSlice<f32>,
10158        x: &CudaSlice<f32>,
10159        t: usize,
10160        next_norm: Option<&CudaSlice<f32>>,
10161    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
10162    {
10163        let n_embd = self.cfg.n_embd as usize;
10164        let bits = layer.gemma4.as_ref().unwrap();
10165        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
10166        let mut xn = e.uninit(t * n_embd)?;
10167        match next_norm {
10168            Some(w) => {
10169                let pair = e.add_scale_rms_norm_q8_1(
10170                    &sn,
10171                    &attn_out,
10172                    bits.layer_scale,
10173                    w,
10174                    &mut xn,
10175                    n_embd,
10176                    t,
10177                    self.cfg.rms_eps,
10178                )?;
10179                Ok((xn, Some(pair)))
10180            }
10181            None => {
10182                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
10183                Ok((xn, None))
10184            }
10185        }
10186    }
10187
10188    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
10189    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
10190    fn gemma4_forward(
10191        &self,
10192        e: &Engine,
10193        tokens: &[u32],
10194        last_only: bool,
10195    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
10196        // E4B routes to its own forward regardless of the caller's entry point (forward /
10197        // forward_last / prime paths all funnel here for gemma4).
10198        if self.is_gemma4_e4b() {
10199            return self.gemma4_e4b_forward(e, tokens, last_only);
10200        }
10201        let n_embd = self.cfg.n_embd as usize;
10202        let t = tokens.len();
10203        let pos: Vec<i32> = (0..t as i32).collect();
10204        let pos_d = e.htod_i32(&pos)?;
10205
10206        let mut x = self.embed(e, tokens)?;
10207        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10208        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
10209        // the bring-up bisect vs llama-eval-callback node stats.
10210        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
10211        let stat =
10212            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
10213                let h = e.dtoh(x)?;
10214                let bad = h.iter().filter(|v| !v.is_finite()).count();
10215                let mx = h
10216                    .iter()
10217                    .filter(|v| v.is_finite())
10218                    .fold(0.0f32, |m, v| m.max(v.abs()));
10219                eprintln!(
10220                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
10221                    &h[..3]
10222                );
10223                Ok(())
10224            };
10225        if probe {
10226            stat(e, &x, "embed")?;
10227        }
10228        for (il, layer) in self.layers.iter().enumerate() {
10229            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
10230            if probe {
10231                stat(e, &x, &format!("L{il}"))?;
10232            }
10233        }
10234        let mut hn = e.zeros(t * n_embd)?;
10235        e.rms_norm(
10236            &x,
10237            self.output_norm.float_data(),
10238            &mut hn,
10239            n_embd,
10240            t,
10241            self.cfg.rms_eps,
10242        )?;
10243        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10244        let n_vocab = self.output.out_features();
10245        let logits = if last_only {
10246            let hv = e.view(&hn, t * n_embd);
10247            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
10248            let mut hlast = e.zeros(n_embd)?;
10249            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
10250            let mut ld = e.matmul(&self.output, &hlast, 1)?;
10251            e.softcap(&mut ld, cap, n_vocab)?;
10252            self.gemma4_suppress(e, &mut ld, 1)?;
10253            e.dtoh(&ld)?
10254        } else {
10255            let mut ld = e.matmul(&self.output, &hn, t)?;
10256            e.softcap(&mut ld, cap, t * n_vocab)?;
10257            self.gemma4_suppress(e, &mut ld, t)?;
10258            e.dtoh(&ld)?
10259        };
10260        Ok(logits)
10261    }
10262
10263    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
10264    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
10265    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
10266    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
10267    pub(crate) fn gemma4_prime(
10268        &self,
10269        e: &Engine,
10270        tokens: &[u32],
10271        cache: &mut Cache,
10272        overlay: Option<&crate::vision::EmbedOverlay>,
10273    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10274        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
10275        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
10276        // whole worker process on this line. The worker now primes gemma4 monolithically and
10277        // routes continuation suffixes tokenwise; this is the per-request backstop.
10278        if cache.pos != 0 {
10279            return Err(
10280                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
10281                        — prime the full prompt in one call or decode tokenwise"
10282                    .into(),
10283            );
10284        }
10285        let n_embd = self.cfg.n_embd as usize;
10286        let eps = self.cfg.rms_eps;
10287        let t = tokens.len();
10288        let pos: Vec<i32> = (0..t as i32).collect();
10289        let pos_d = e.htod_i32(&pos)?;
10290        let mut x = self.embed(e, tokens)?;
10291        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
10292        // Masked-prefill arm (lane/gemma-vision): splice the tower rows AFTER the
10293        // sqrt(n_embd) text scale — the reference scales token batches only
10294        // (llama.cpp gemma4.cpp:182 `ubatch.token ? sqrtf(n_embd) : 1.0f`), so image
10295        // embeddings enter the trunk raw. Each span becomes an attention ISLAND:
10296        // bidirectional within itself, causal+SWA everywhere else, matching the
10297        // reference's llama_set_causal_attn(false) image batch exactly.
10298        let island: Option<CudaSlice<i32>> = match overlay {
10299            Some(ov) => {
10300                let mut span_id = vec![-1i32; t];
10301                for (i, &(pos, row_off, n_rows)) in ov.spans.iter().enumerate() {
10302                    if pos + n_rows > t {
10303                        return Err(format!(
10304                            "gemma4 overlay span {i} [{pos}, {}) exceeds the prompt ({t})",
10305                            pos + n_rows
10306                        )
10307                        .into());
10308                    }
10309                    let view = ov.rows.slice(row_off * n_embd..(row_off + n_rows) * n_embd);
10310                    e.copy_view_into(&mut x, pos * n_embd, &view, n_rows * n_embd)?;
10311                    for s in span_id.iter_mut().skip(pos).take(n_rows) {
10312                        *s = i as i32;
10313                    }
10314                }
10315                // MEMRA_GV_FORCE_CAUSAL=1: WRONG-ARM probe seam — splice the rows but
10316                // keep the plain causal mask. Exists only so the decisive probe can show
10317                // the island mask itself changes the answer; never on in serving.
10318                if std::env::var("MEMRA_GV_FORCE_CAUSAL").as_deref() == Ok("1") {
10319                    None
10320                } else {
10321                    Some(e.htod_i32(&span_id)?)
10322                }
10323            }
10324            None => None,
10325        };
10326        for (il, layer) in self.layers.iter().enumerate() {
10327            let mut h = e.zeros(t * n_embd)?;
10328            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
10329            let Mixer::Full(fa) = &layer.mixer else {
10330                panic!("gemma4 layer not full-attn")
10331            };
10332            let trace = il == 0 && std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1");
10333            if trace {
10334                let v = e.dtoh(&h)?;
10335                let nan = v.iter().filter(|x| x.is_nan()).count();
10336                eprintln!("[g4-prime-trace] L0 post-attn_norm: nan={nan}/{}", v.len());
10337            }
10338            let o =
10339                self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache), island.as_ref())?;
10340            if trace {
10341                let v = e.dtoh(&o)?;
10342                let nan = v.iter().filter(|x| x.is_nan()).count();
10343                eprintln!("[g4-prime-trace] L0 post-attn: nan={nan}/{}", v.len());
10344            }
10345            let mut cur = e.zeros(t * n_embd)?;
10346            e.rms_norm(
10347                &o,
10348                layer.post_attn_norm.float_data(),
10349                &mut cur,
10350                n_embd,
10351                t,
10352                eps,
10353            )?;
10354            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
10355            self.dflash_tap(e, cache, il, &x, t)?;
10356            // MEMRA_G4_PRIME_TRACE=1: per-layer NaN hunt (recipe-probe diagnostic; off = zero cost)
10357            if std::env::var("MEMRA_G4_PRIME_TRACE").as_deref() == Ok("1") {
10358                let h = e.dtoh(&x)?;
10359                let nan = h.iter().filter(|v| v.is_nan()).count();
10360                let amax = h.iter().fold(0f32, |a, v| a.max(v.abs()));
10361                eprintln!(
10362                    "[g4-prime-trace] layer {il}: nan={nan}/{} amax={amax:.3}",
10363                    h.len()
10364                );
10365                if nan > 0 {
10366                    return Err(format!("g4-prime-trace: first NaN at layer {il}").into());
10367                }
10368            }
10369        }
10370        cache.pos += t;
10371        let hiddens = e.clone_dtod(&x)?;
10372        let xv = e.view(&x, t * n_embd);
10373        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
10374        let mut h_seed = e.zeros(n_embd)?;
10375        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
10376        let mut hn = e.uninit(n_embd)?;
10377        e.rms_norm(
10378            &h_seed,
10379            self.output_norm.float_data(),
10380            &mut hn,
10381            n_embd,
10382            1,
10383            eps,
10384        )?;
10385        let mut ld = e.matmul(&self.output, &hn, 1)?;
10386        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
10387        e.softcap(&mut ld, cap, self.output.out_features())?;
10388        self.gemma4_suppress(e, &mut ld, 1)?;
10389        let logits = e.dtoh(&ld)?;
10390        Ok((logits, h_seed, hiddens))
10391    }
10392
10393    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
10394    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
10395    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
10396    /// fused norm emits q8 directly — the f32 h never materializes).
10397    fn gemma4_decode_attn(
10398        &self,
10399        e: &Engine,
10400        fa: &crate::hybrid::FullAttnLayer,
10401        il: usize,
10402        hq: &CudaSlice<i8>,
10403        hdq: &CudaSlice<f32>,
10404        pos_d: &CudaSlice<i32>,
10405        cache: &mut Cache,
10406    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10407        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10408        let eps = self.cfg.rms_eps;
10409        let aux = self.gemma4_aux.as_ref().unwrap();
10410        let ones = aux.ones(e);
10411        #[cfg(debug_assertions)]
10412        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
10413        let (hq, hdq) = (hq, hdq);
10414        let h0 = e.zeros(0)?;
10415        let h = &h0;
10416        let (q0, k0, v0) = if swa {
10417            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
10418                Some(t3) => t3,
10419                // NVFP4mix trio is MIXED-type (wv stays Q8_0), so fused3 can never
10420                // match — fuse the uniform (q,k) pair and take v as its own single.
10421                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10422                    Some((q0, k0)) => {
10423                        let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?;
10424                        (q0, k0, v0)
10425                    }
10426                    None => (
10427                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10428                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10429                        e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10430                    ),
10431                },
10432            }
10433        } else {
10434            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10435                Some(p) => p,
10436                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, &hq, &hdq, 1)? {
10437                    Some(p) => p,
10438                    None => (
10439                        e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10440                        e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10441                    ),
10442                },
10443            };
10444            let v0 = e.clone_dtod(&k0)?;
10445            (q0, k0, v0)
10446        };
10447        let mut q = e.uninit(nh * hd)?;
10448        let mut k = e.uninit(nkv * hd)?;
10449        let mut v = e.uninit(nkv * hd)?;
10450        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
10451        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
10452        let ff = if swa {
10453            None
10454        } else {
10455            Some(
10456                aux.rope_freqs(e)
10457                    .expect("gemma4 global rope needs rope_freqs.weight"),
10458            )
10459        };
10460        #[cfg(debug_assertions)]
10461        if let Some(ff) = ff {
10462            crate::debug_assert_tensor_stream_device(
10463                ff,
10464                &e.stream(),
10465                "gemma4_decode_attn.rope_freqs",
10466            );
10467        }
10468        let kvl = cache.kv[il].as_mut().unwrap();
10469        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10470        if crate::Engine::qkv_append_on() {
10471            // append fold, eager completion (zoo-fusion arc): the dc arms took this fold
10472            // 2026-07-23; the eager step kept the 2-launch pair it doesn't need. Host-len
10473            // twin of the dc fold — bit-identical bodies, one launch per layer.
10474            e.rms_norm_qkv_rope_append(
10475                &q0,
10476                &k0,
10477                &v0,
10478                fa.q_norm.float_data(),
10479                fa.k_norm.float_data(),
10480                ones,
10481                &mut q,
10482                &mut k,
10483                &mut v,
10484                hd,
10485                self.gemma4_rope_dims(il),
10486                nh,
10487                nkv,
10488                pos_d,
10489                nh,
10490                nkv,
10491                base,
10492                1.0,
10493                ff,
10494                eps,
10495                &mut kvl.k,
10496                &mut kvl.v,
10497                kvl.len,
10498                kvl.k_tok_bytes,
10499                kvl.v_tok_bytes,
10500                kv_fp8,
10501            )?;
10502        } else {
10503            e.rms_norm_qkv_rope(
10504                &q0,
10505                &k0,
10506                &v0,
10507                fa.q_norm.float_data(),
10508                fa.k_norm.float_data(),
10509                ones,
10510                &mut q,
10511                &mut k,
10512                &mut v,
10513                hd,
10514                self.gemma4_rope_dims(il),
10515                nh,
10516                nkv,
10517                pos_d,
10518                nh,
10519                nkv,
10520                base,
10521                1.0,
10522                ff,
10523                eps,
10524            )?;
10525            e.append_kv_quantized(
10526                &k,
10527                &v,
10528                &mut kvl.k,
10529                &mut kvl.v,
10530                kvl.len,
10531                kvl.kv_dim_k,
10532                kvl.kv_dim_v,
10533                kvl.k_tok_bytes,
10534                kvl.v_tok_bytes,
10535                kv_fp8,
10536            )?;
10537        }
10538        kvl.len += 1;
10539        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
10540        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
10541        // positional). Globals attend the full history.
10542        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10543        let mut attn = e.uninit(nh * hd)?;
10544        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
10545        if !swa
10546            && hd == 512
10547            && kvl.len >= crate::fa512_min_tkv()
10548            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10549        {
10550            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10551            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10552            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
10553            let base = kvl.len as i32;
10554            e.i32_set_k(&mut kvl.len_d, base)?;
10555            e.fa_decode_rows(
10556                &q,
10557                &kp,
10558                &vp,
10559                &mut attn,
10560                hd,
10561                nh,
10562                nkv,
10563                kvl.len - 1,
10564                1,
10565                scale,
10566                kvl.k_tok_bytes,
10567                kvl.v_tok_bytes,
10568                Some((&kvl.len_d, -1)),
10569                false,
10570                false,
10571                None,
10572            )?;
10573            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10574        }
10575        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
10576        if swa
10577            && kvl.len > win
10578            && hd == 256
10579            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10580        {
10581            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10582            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10583            let base = kvl.len as i32;
10584            e.i32_set_k(&mut kvl.len_d, base)?;
10585            e.fa_decode_rows_w(
10586                &q,
10587                &kp,
10588                &vp,
10589                &mut attn,
10590                hd,
10591                nh,
10592                nkv,
10593                &kvl.len_d,
10594                -1,
10595                1,
10596                scale,
10597                win,
10598                kvl.k_tok_bytes,
10599                kvl.v_tok_bytes,
10600                None,
10601            )?;
10602            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10603        }
10604        let (off_tok, t_kv) = if swa && kvl.len > win {
10605            (kvl.len - win, win)
10606        } else {
10607            (0, kvl.len)
10608        };
10609        let k_view = e.view_u8_range(
10610            &kvl.k,
10611            off_tok * kvl.k_tok_bytes,
10612            (off_tok + t_kv) * kvl.k_tok_bytes,
10613        );
10614        let v_view = e.view_u8_range(
10615            &kvl.v,
10616            off_tok * kvl.v_tok_bytes,
10617            (off_tok + t_kv) * kvl.v_tok_bytes,
10618        );
10619        e.fa_decode_kvmod(
10620            &q,
10621            &k_view,
10622            &v_view,
10623            &mut attn,
10624            hd,
10625            nh,
10626            nkv,
10627            t_kv,
10628            scale,
10629            kvl.k_tok_bytes,
10630            kvl.v_tok_bytes,
10631            swa && crate::Engine::wkv_on(),
10632        )?;
10633        Ok(e.matmul(&fa.wo, &attn, 1)?)
10634    }
10635
10636    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
10637    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
10638    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
10639    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
10640    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
10641    /// in-graph; the driver gates).
10642    #[allow(clippy::too_many_arguments)]
10643    pub fn gemma4_decode_step_dc(
10644        &self,
10645        e: &Engine,
10646        token_d: &CudaSlice<u32>,
10647        pos_d: &mut CudaSlice<i32>,
10648        embd_gpu: &CudaSlice<u8>,
10649        embd_qt: i32,
10650        embd_rb: usize,
10651        cache: &mut Cache,
10652        n_vocab: usize,
10653        cap_bucket_max: Option<(usize, usize)>,
10654    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10655        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10656        self.gemma4_decode_step_dc_into(
10657            e,
10658            token_d,
10659            pos_d,
10660            embd_gpu,
10661            embd_qt,
10662            embd_rb,
10663            cache,
10664            n_vocab,
10665            cap_bucket_max,
10666            &mut tok_out,
10667        )?;
10668        Ok(tok_out)
10669    }
10670
10671    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
10672    /// every replay; pass `token_d` itself for the self-feeding graph loop).
10673    #[allow(clippy::too_many_arguments)]
10674    pub fn gemma4_decode_step_dc_into(
10675        &self,
10676        e: &Engine,
10677        token_d: &CudaSlice<u32>,
10678        pos_d: &mut CudaSlice<i32>,
10679        embd_gpu: &CudaSlice<u8>,
10680        embd_qt: i32,
10681        embd_rb: usize,
10682        cache: &mut Cache,
10683        n_vocab: usize,
10684        cap_bucket_max: Option<(usize, usize)>,
10685        tok_out: &mut CudaSlice<u32>,
10686    ) -> Result<(), Box<dyn std::error::Error>> {
10687        let n_embd = self.cfg.n_embd as usize;
10688        let eps = self.cfg.rms_eps;
10689        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10690        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10691        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10692        let n_layers = self.layers.len();
10693        for (il, layer) in self.layers.iter().enumerate() {
10694            let (hq, hdq) = match h_carry.take() {
10695                Some(p) => p,
10696                None => {
10697                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10698                }
10699            };
10700            let Mixer::Full(fa) = &layer.mixer else {
10701                panic!("gemma4 layer {il} not full-attn")
10702            };
10703            let o =
10704                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10705            let next_norm = if il + 1 < n_layers {
10706                Some(self.layers[il + 1].attn_norm.float_data())
10707            } else {
10708                None
10709            };
10710            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
10711            x = xn;
10712            h_carry = hn;
10713        }
10714        let mut hn = e.uninit(n_embd)?;
10715        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10716        let mut logits = e.matmul(&self.output, &hn, 1)?;
10717        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
10718        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10719        e.inc_seqlen(pos_d)?;
10720        if cap_bucket_max.is_none() {
10721            cache.pos += 1;
10722        }
10723        Ok(())
10724    }
10725
10726    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
10727    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
10728    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
10729    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
10730
10731    /// Build the slot set (call OUTSIDE any capture).
10732    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10733        let n_embd = self.cfg.n_embd as usize;
10734        let n_vocab = self.output.out_features();
10735        let n_layers = self.layers.len();
10736        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10737        for il in 0..n_layers {
10738            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10739            qmax = qmax.max(nh * hd);
10740            kvmax = kvmax.max(nkv * hd);
10741            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10742                ffmax = ffmax.max(ffn_gate.out_features());
10743            }
10744        }
10745        Ok(G4DcSlots {
10746            x: e.uninit(n_embd)?,
10747            xn: e.uninit(n_embd)?,
10748            cur: e.uninit(n_embd)?,
10749            hq: e.alloc_i8_uninit(n_embd)?,
10750            hd_: e.uninit(n_embd / 32)?,
10751            q0: e.uninit(qmax)?,
10752            k0: e.uninit(kvmax)?,
10753            v0: e.uninit(kvmax)?,
10754            q: e.uninit(qmax)?,
10755            k: e.uninit(kvmax)?,
10756            v: e.uninit(kvmax)?,
10757            attn: e.uninit(qmax)?,
10758            o: e.uninit(n_embd)?,
10759            attn_out: e.uninit(n_embd)?,
10760            zsh: e.uninit(n_embd)?,
10761            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
10762            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
10763            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10764            zd: e.uninit(n_embd.max(qmax) / 32)?,
10765            gate: e.uninit(ffmax)?,
10766            up: e.uninit(ffmax)?,
10767            act: e.uninit(ffmax)?,
10768            actq: e.alloc_i8_uninit(ffmax)?,
10769            actd: e.uninit(ffmax / 32)?,
10770            f0: e.uninit(n_embd)?,
10771            sn: e.uninit(n_embd)?,
10772            hn: e.uninit(n_embd)?,
10773            logits: e.uninit(n_vocab)?,
10774        })
10775    }
10776
10777    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
10778    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
10779    fn g4_matvec_m1_into(
10780        &self,
10781        e: &Engine,
10782        w: &crate::model::GpuTensor,
10783        aq: &CudaSlice<i8>,
10784        ad: &CudaSlice<f32>,
10785        y: &mut CudaSlice<f32>,
10786    ) -> Result<(), Box<dyn std::error::Error>> {
10787        use crate::model::GpuTensor;
10788        let (bytes, qtype, row_bytes, scale, rp) = match w {
10789            GpuTensor::Quant {
10790                bytes,
10791                qtype,
10792                row_bytes,
10793                scale,
10794                rp,
10795                ..
10796            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10797            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10798        };
10799        let (mbytes, mrp) = match w {
10800            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10801            _ => (bytes, rp),
10802        };
10803        e.qmatvec_mmvq_into(
10804            mbytes,
10805            aq,
10806            ad,
10807            1,
10808            w.in_features(),
10809            w.out_features(),
10810            qtype,
10811            row_bytes,
10812            scale,
10813            mrp,
10814            y,
10815        )
10816    }
10817
10818    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
10819    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
10820    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
10821    #[allow(clippy::too_many_arguments)]
10822    pub fn gemma4_decode_step_dc_slotted(
10823        &self,
10824        e: &Engine,
10825        token_d: &CudaSlice<u32>,
10826        pos_d: &mut CudaSlice<i32>,
10827        embd_gpu: &CudaSlice<u8>,
10828        embd_qt: i32,
10829        embd_rb: usize,
10830        cache: &mut Cache,
10831        n_vocab: usize,
10832        cap_bucket_max: Option<(usize, usize)>,
10833        sl: &mut G4DcSlots,
10834        tok_out: &mut CudaSlice<u32>,
10835        ring: Option<(&mut CudaSlice<u32>, usize)>,
10836    ) -> Result<(), Box<dyn std::error::Error>> {
10837        let n_embd = self.cfg.n_embd as usize;
10838        let eps = self.cfg.rms_eps;
10839        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10840        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10841        let n_layers = self.layers.len();
10842        let mut has_carry = false;
10843        for il in 0..n_layers {
10844            if !has_carry {
10845                e.rms_norm_q8_1_into(
10846                    &sl.x,
10847                    self.layers[il].attn_norm.float_data(),
10848                    n_embd,
10849                    1,
10850                    eps,
10851                    &mut sl.hq,
10852                    &mut sl.hd_,
10853                )?;
10854            }
10855            has_carry = true;
10856            let layer = &self.layers[il];
10857            let Mixer::Full(fa) = &layer.mixer else {
10858                panic!("gemma4 layer {il} not full-attn")
10859            };
10860            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10861            // pn-fold: the tail's entry launch consumes RAW sl.o (post_attn folded there);
10862            // the standalone norm only survives on the unfused seam arm.
10863            if !Engine::g4_pnfold_on() {
10864                e.rms_norm(
10865                    &sl.o,
10866                    layer.post_attn_norm.float_data(),
10867                    &mut sl.cur,
10868                    n_embd,
10869                    1,
10870                    eps,
10871                )?;
10872            }
10873            let next_norm = if il + 1 < n_layers {
10874                Some(self.layers[il + 1].attn_norm.float_data())
10875            } else {
10876                None
10877            };
10878            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10879            std::mem::swap(&mut sl.x, &mut sl.xn);
10880        }
10881        e.rms_norm(
10882            &sl.x,
10883            self.output_norm.float_data(),
10884            &mut sl.hn,
10885            n_embd,
10886            1,
10887            eps,
10888        )?;
10889        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10890        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
10891        {
10892            let (zq, zd) = (&sl.zq, &sl.zd);
10893            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10894            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10895            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10896        }
10897        self.gemma4_suppress(e, &mut sl.logits, 1)?;
10898        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10899        if let Some((ring, base)) = ring {
10900            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
10901            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
10902            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
10903            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10904        }
10905        e.inc_seqlen(pos_d)?;
10906        if cap_bucket_max.is_none() {
10907            cache.pos += 1;
10908        }
10909        Ok(())
10910    }
10911
10912    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
10913    #[allow(clippy::too_many_arguments)]
10914    fn gemma4_decode_attn_dc_slotted(
10915        &self,
10916        e: &Engine,
10917        fa: &crate::hybrid::FullAttnLayer,
10918        il: usize,
10919        pos_d: &CudaSlice<i32>,
10920        cache: &mut Cache,
10921        cap_bucket_max: Option<(usize, usize)>,
10922        sl: &mut G4DcSlots,
10923    ) -> Result<(), Box<dyn std::error::Error>> {
10924        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10925        let eps = self.cfg.rms_eps;
10926        let aux = self.gemma4_aux.as_ref().unwrap();
10927        let ones = aux.ones(e);
10928        #[cfg(debug_assertions)]
10929        crate::debug_assert_tensor_stream_device(
10930            ones,
10931            &e.stream(),
10932            "gemma4_decode_attn_dc_slotted.ones",
10933        );
10934        {
10935            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10936            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10937            if swa {
10938                if !e.matmul_q4_fused3_into(
10939                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10940                )? {
10941                    // NVFP4mix trio is MIXED-type (wv stays Q8_0): fuse the uniform
10942                    // (q,k) pair, v through the generic m1 slot matvec — the same two
10943                    // bodies the eager arm runs, slot-fed for the alloc-free capture.
10944                    if e.matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10945                    {
10946                        self.g4_matvec_m1_into(e, &fa.wv, hq, hdq, &mut sl.v0)?;
10947                    } else {
10948                        return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10949                    }
10950                }
10951            } else {
10952                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10953                    && !e
10954                        .matmul_nvfp4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)?
10955                {
10956                    return Err("slotted step: fused2 unavailable".into());
10957                }
10958                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10959                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10960            }
10961        }
10962        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
10963        // kernel-for-kernel (graph stream-identity gate).
10964        let ff = if swa {
10965            None
10966        } else {
10967            Some(
10968                aux.rope_freqs(e)
10969                    .expect("gemma4 global rope needs rope_freqs.weight"),
10970            )
10971        };
10972        #[cfg(debug_assertions)]
10973        if let Some(ff) = ff {
10974            crate::debug_assert_tensor_stream_device(
10975                ff,
10976                &e.stream(),
10977                "gemma4_decode_attn_dc_slotted.rope_freqs",
10978            );
10979        }
10980        let kvl = cache.kv[il].as_mut().unwrap();
10981        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10982        if crate::Engine::qkv_append_on() {
10983            // append fold (2026-07-23): mirrors dc_into.
10984            e.rms_norm_qkv_rope_append_dc(
10985                &sl.q0,
10986                &sl.k0,
10987                &sl.v0,
10988                fa.q_norm.float_data(),
10989                fa.k_norm.float_data(),
10990                ones,
10991                &mut sl.q,
10992                &mut sl.k,
10993                &mut sl.v,
10994                hd,
10995                self.gemma4_rope_dims(il),
10996                nh,
10997                nkv,
10998                pos_d,
10999                nh,
11000                nkv,
11001                base,
11002                1.0,
11003                ff,
11004                eps,
11005                &mut kvl.k,
11006                &mut kvl.v,
11007                &kvl.len_d,
11008                kvl.k_tok_bytes,
11009                kvl.v_tok_bytes,
11010                kv_fp8,
11011            )?;
11012        } else {
11013            e.rms_norm_qkv_rope(
11014                &sl.q0,
11015                &sl.k0,
11016                &sl.v0,
11017                fa.q_norm.float_data(),
11018                fa.k_norm.float_data(),
11019                ones,
11020                &mut sl.q,
11021                &mut sl.k,
11022                &mut sl.v,
11023                hd,
11024                self.gemma4_rope_dims(il),
11025                nh,
11026                nkv,
11027                pos_d,
11028                nh,
11029                nkv,
11030                base,
11031                1.0,
11032                ff,
11033                eps,
11034            )?;
11035            e.append_kv_quantized_dc(
11036                &sl.k,
11037                &sl.v,
11038                &mut kvl.k,
11039                &mut kvl.v,
11040                &kvl.len_d,
11041                kvl.kv_dim_k,
11042                kvl.kv_dim_v,
11043                kvl.k_tok_bytes,
11044                kvl.v_tok_bytes,
11045                kv_fp8,
11046            )?;
11047        }
11048        e.inc_seqlen(&mut kvl.len_d)?;
11049        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
11050        let k_view = e.view_u8(&kvl.k, kvl.k.len());
11051        let v_view = e.view_u8(&kvl.v, kvl.v.len());
11052        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11053        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11054        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
11055        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
11056        // the dc_into arm branch-for-branch (stream gate).
11057        let mut fa_q8 = false;
11058        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11059            e.fa_decode_rows(
11060                &sl.q,
11061                &k_view,
11062                &v_view,
11063                &mut sl.attn,
11064                hd,
11065                nh,
11066                nkv,
11067                b_glob - 1,
11068                1,
11069                scale,
11070                kvl.k_tok_bytes,
11071                kvl.v_tok_bytes,
11072                Some((&kvl.len_d, -1)),
11073                false,
11074                false,
11075                Some((&mut sl.zq, &mut sl.zd)),
11076            )?;
11077            fa_q8 = true;
11078        } else if swa && b_swa > win && hd == 256 && rows_on {
11079            e.fa_decode_rows_w(
11080                &sl.q,
11081                &k_view,
11082                &v_view,
11083                &mut sl.attn,
11084                hd,
11085                nh,
11086                nkv,
11087                &kvl.len_d,
11088                -1,
11089                1,
11090                scale,
11091                win,
11092                kvl.k_tok_bytes,
11093                kvl.v_tok_bytes,
11094                Some((&mut sl.zq, &mut sl.zd)),
11095            )?;
11096            fa_q8 = true;
11097        } else {
11098            let b = if swa { b_swa } else { b_glob };
11099            e.fa_decode_dc(
11100                &sl.q,
11101                &k_view,
11102                &v_view,
11103                &mut sl.attn,
11104                hd,
11105                nh,
11106                nkv,
11107                &kvl.len_d,
11108                b,
11109                scale,
11110                kvl.k_tok_bytes,
11111                kvl.v_tok_bytes,
11112                swa && crate::Engine::wkv_on(),
11113            )?;
11114        }
11115        if !fa_q8 {
11116            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
11117            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
11118        }
11119        {
11120            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11121            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11122            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
11123        }
11124        Ok(())
11125    }
11126
11127    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
11128    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
11129    fn gemma4_layer_tail_slotted(
11130        &self,
11131        e: &Engine,
11132        layer: &crate::hybrid::HybridLayer,
11133        next_norm: Option<&CudaSlice<f32>>,
11134        sl: &mut G4DcSlots,
11135    ) -> Result<(), Box<dyn std::error::Error>> {
11136        let n_embd = self.cfg.n_embd as usize;
11137        let eps = self.cfg.rms_eps;
11138        let bits = layer.gemma4.as_ref().unwrap();
11139        let crate::hybrid::Ffn::Dense {
11140            ffn_gate,
11141            ffn_up,
11142            ffn_down,
11143        } = &layer.ffn
11144        else {
11145            return Err("slotted tail: dense ffn only".into());
11146        };
11147        let pnfold = Engine::g4_pnfold_on();
11148        if pnfold {
11149            // pn-fold entry (E4B glue, slot-fed): rms(o, post_attn) + residual + ffn_norm
11150            // + q8 emit in ONE launch — replaces rms_norm + add_rms_norm + quantize.
11151            let or = unsafe { &*(&sl.o as *const CudaSlice<f32>) };
11152            let xr = unsafe { &*(&sl.x as *const CudaSlice<f32>) };
11153            e.rms_pre_add_rms_norm_q8z_into(
11154                or,
11155                layer.post_attn_norm.float_data(),
11156                xr,
11157                bits.ffn_norm.float_data(),
11158                &mut sl.attn_out,
11159                &mut sl.zsh,
11160                n_embd,
11161                1,
11162                eps,
11163                &mut sl.zq,
11164                &mut sl.zd,
11165            )?;
11166        } else {
11167            e.add_rms_norm(
11168                &sl.cur,
11169                &sl.x,
11170                bits.ffn_norm.float_data(),
11171                &mut sl.attn_out,
11172                &mut sl.zsh,
11173                n_embd,
11174                1,
11175                eps,
11176            )?;
11177        }
11178        let n_ff = ffn_gate.out_features();
11179        if !pnfold {
11180            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
11181            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
11182        }
11183        {
11184            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
11185            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
11186            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)?
11187                && !e.matmul_nvfp4_fused2_into(
11188                    ffn_gate,
11189                    ffn_up,
11190                    zq,
11191                    zd,
11192                    &mut sl.gate,
11193                    &mut sl.up,
11194                )?
11195            {
11196                return Err("slotted tail: ffn fused2 unavailable".into());
11197            }
11198        }
11199        debug_assert!(e.uses_q8_1_fast(ffn_down));
11200        {
11201            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
11202            let upv = e.view(upr, n_ff);
11203            let up_all = upv.slice(0..n_ff);
11204            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
11205            e.gelu_tanh_mul_q8_1_into(
11206                gr,
11207                &up_all,
11208                &mut sl.act,
11209                n_ff,
11210                1,
11211                &mut sl.actq,
11212                &mut sl.actd,
11213            )?;
11214        }
11215        {
11216            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
11217            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
11218            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
11219        }
11220        if pnfold {
11221            // pn-fold exit: rms(f0, post_ffw) + scaled residual + next attn_norm + q8
11222            // emit in ONE launch — replaces rms_norm + add_scale_rms_norm_q8_1.
11223            if let Some(w) = next_norm {
11224                let f0r = unsafe { &*(&sl.f0 as *const CudaSlice<f32>) };
11225                let aor = unsafe { &*(&sl.attn_out as *const CudaSlice<f32>) };
11226                e.rms_pre_add_scale_rms_norm_q8_1_into(
11227                    f0r,
11228                    bits.post_ffw_norm.float_data(),
11229                    aor,
11230                    bits.layer_scale,
11231                    w,
11232                    &mut sl.xn,
11233                    n_embd,
11234                    1,
11235                    eps,
11236                    &mut sl.hq,
11237                    &mut sl.hd_,
11238                )?;
11239                return Ok(());
11240            }
11241        }
11242        e.rms_norm(
11243            &sl.f0,
11244            bits.post_ffw_norm.float_data(),
11245            &mut sl.sn,
11246            n_embd,
11247            1,
11248            eps,
11249        )?;
11250        match next_norm {
11251            Some(w) => {
11252                e.add_scale_rms_norm_q8_1_into(
11253                    &sl.sn,
11254                    &sl.attn_out,
11255                    bits.layer_scale,
11256                    w,
11257                    &mut sl.xn,
11258                    n_embd,
11259                    1,
11260                    eps,
11261                    &mut sl.hq,
11262                    &mut sl.hd_,
11263                )?;
11264            }
11265            None => {
11266                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
11267            }
11268        }
11269        Ok(())
11270    }
11271
11272    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
11273    #[allow(clippy::too_many_arguments)]
11274    fn gemma4_decode_attn_dc(
11275        &self,
11276        e: &Engine,
11277        fa: &crate::hybrid::FullAttnLayer,
11278        il: usize,
11279        hq: &CudaSlice<i8>,
11280        hdq: &CudaSlice<f32>,
11281        pos_d: &CudaSlice<i32>,
11282        cache: &mut Cache,
11283        cap_bucket_max: Option<(usize, usize)>,
11284    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11285        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11286        let eps = self.cfg.rms_eps;
11287        let aux = self.gemma4_aux.as_ref().unwrap();
11288        let ones = aux.ones(e);
11289        #[cfg(debug_assertions)]
11290        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
11291        let (q0, k0, v0) = if swa {
11292            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
11293                Some(t3) => t3,
11294                // NVFP4mix: wv is Q8_0 (mixed trio) — fuse the (q,k) pair, v single.
11295                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11296                    Some((q0, k0)) => {
11297                        let h0 = e.zeros(0)?;
11298                        let v0 = e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?;
11299                        (q0, k0, v0)
11300                    }
11301                    None => {
11302                        let h0 = e.zeros(0)?;
11303                        (
11304                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11305                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11306                            e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
11307                        )
11308                    }
11309                },
11310            }
11311        } else {
11312            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
11313                Some(p) => p,
11314                None => match e.matmul_nvfp4_fused2(&fa.wq, &fa.wk, hq, hdq, 1)? {
11315                    Some(p) => p,
11316                    None => {
11317                        let h0 = e.zeros(0)?;
11318                        (
11319                            e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
11320                            e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
11321                        )
11322                    }
11323                },
11324            };
11325            let v0 = e.clone_dtod(&k0)?;
11326            (q0, k0, v0)
11327        };
11328        let mut q = e.uninit(nh * hd)?;
11329        let mut k = e.uninit(nkv * hd)?;
11330        let mut v = e.uninit(nkv * hd)?;
11331        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
11332        let ff = if swa {
11333            None
11334        } else {
11335            Some(
11336                aux.rope_freqs(e)
11337                    .expect("gemma4 global rope needs rope_freqs.weight"),
11338            )
11339        };
11340        #[cfg(debug_assertions)]
11341        if let Some(ff) = ff {
11342            crate::debug_assert_tensor_stream_device(
11343                ff,
11344                &e.stream(),
11345                "gemma4_decode_attn_dc.rope_freqs",
11346            );
11347        }
11348        let kvl = cache.kv[il].as_mut().unwrap();
11349        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
11350        if crate::Engine::qkv_append_on() {
11351            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
11352            e.rms_norm_qkv_rope_append_dc(
11353                &q0,
11354                &k0,
11355                &v0,
11356                fa.q_norm.float_data(),
11357                fa.k_norm.float_data(),
11358                ones,
11359                &mut q,
11360                &mut k,
11361                &mut v,
11362                hd,
11363                self.gemma4_rope_dims(il),
11364                nh,
11365                nkv,
11366                pos_d,
11367                nh,
11368                nkv,
11369                base,
11370                1.0,
11371                ff,
11372                eps,
11373                &mut kvl.k,
11374                &mut kvl.v,
11375                &kvl.len_d,
11376                kvl.k_tok_bytes,
11377                kvl.v_tok_bytes,
11378                kv_fp8,
11379            )?;
11380        } else {
11381            e.rms_norm_qkv_rope(
11382                &q0,
11383                &k0,
11384                &v0,
11385                fa.q_norm.float_data(),
11386                fa.k_norm.float_data(),
11387                ones,
11388                &mut q,
11389                &mut k,
11390                &mut v,
11391                hd,
11392                self.gemma4_rope_dims(il),
11393                nh,
11394                nkv,
11395                pos_d,
11396                nh,
11397                nkv,
11398                base,
11399                1.0,
11400                ff,
11401                eps,
11402            )?;
11403            e.append_kv_quantized_dc(
11404                &k,
11405                &v,
11406                &mut kvl.k,
11407                &mut kvl.v,
11408                &kvl.len_d,
11409                kvl.kv_dim_k,
11410                kvl.kv_dim_v,
11411                kvl.k_tok_bytes,
11412                kvl.v_tok_bytes,
11413                kv_fp8,
11414            )?;
11415        }
11416        e.inc_seqlen(&mut kvl.len_d)?;
11417        let mut attn = e.uninit(nh * hd)?;
11418        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
11419        // rides g4_matvec_m1_into instead of matmul's internal quantize.
11420        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11421        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
11422        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
11423        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
11424        // (gemma4_e4b_attn, +0.65% valid window).
11425        match cap_bucket_max {
11426            None => {
11427                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
11428                // decode (SWA layers attend the last `sliding_window` keys); the device
11429                // counters carry only the append slot + the graph seam.
11430                kvl.len += 1;
11431                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11432                if !swa
11433                    && hd == 512
11434                    && kvl.len >= crate::fa512_min_tkv()
11435                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11436                {
11437                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
11438                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
11439                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11440                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11441                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11442                    e.fa_decode_rows(
11443                        &q,
11444                        &kp,
11445                        &vp,
11446                        &mut attn,
11447                        hd,
11448                        nh,
11449                        nkv,
11450                        kvl.len - 1,
11451                        1,
11452                        scale,
11453                        kvl.k_tok_bytes,
11454                        kvl.v_tok_bytes,
11455                        Some((&kvl.len_d, -1)),
11456                        false,
11457                        false,
11458                        Some((&mut aq8, &mut ad8)),
11459                    )?;
11460                    fa_q8 = Some((aq8, ad8));
11461                } else if swa
11462                    && kvl.len > win
11463                    && hd == 256
11464                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11465                {
11466                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
11467                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
11468                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
11469                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11470                    e.fa_decode_rows_w(
11471                        &q,
11472                        &kp,
11473                        &vp,
11474                        &mut attn,
11475                        hd,
11476                        nh,
11477                        nkv,
11478                        &kvl.len_d,
11479                        -1,
11480                        1,
11481                        scale,
11482                        win,
11483                        kvl.k_tok_bytes,
11484                        kvl.v_tok_bytes,
11485                        Some((&mut aq8, &mut ad8)),
11486                    )?;
11487                    fa_q8 = Some((aq8, ad8));
11488                } else {
11489                    let (off_tok, t_kv) = if swa && kvl.len > win {
11490                        (kvl.len - win, win)
11491                    } else {
11492                        (0, kvl.len)
11493                    };
11494                    let k_view = e.view_u8_range(
11495                        &kvl.k,
11496                        off_tok * kvl.k_tok_bytes,
11497                        (off_tok + t_kv) * kvl.k_tok_bytes,
11498                    );
11499                    let v_view = e.view_u8_range(
11500                        &kvl.v,
11501                        off_tok * kvl.v_tok_bytes,
11502                        (off_tok + t_kv) * kvl.v_tok_bytes,
11503                    );
11504                    e.fa_decode_kvmod(
11505                        &q,
11506                        &k_view,
11507                        &v_view,
11508                        &mut attn,
11509                        hd,
11510                        nh,
11511                        nkv,
11512                        t_kv,
11513                        scale,
11514                        kvl.k_tok_bytes,
11515                        kvl.v_tok_bytes,
11516                        swa && crate::Engine::wkv_on(),
11517                    )?;
11518                }
11519            }
11520            Some((b_swa, b_glob)) => {
11521                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
11522                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
11523                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
11524                // the RUNG max for the rows family (kernels derive per-replay splits from
11525                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
11526                let k_view = e.view_u8(&kvl.k, kvl.k.len());
11527                let v_view = e.view_u8(&kvl.v, kvl.v.len());
11528                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
11529                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11530                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
11531                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11532                    e.fa_decode_rows(
11533                        &q,
11534                        &k_view,
11535                        &v_view,
11536                        &mut attn,
11537                        hd,
11538                        nh,
11539                        nkv,
11540                        b_glob - 1,
11541                        1,
11542                        scale,
11543                        kvl.k_tok_bytes,
11544                        kvl.v_tok_bytes,
11545                        Some((&kvl.len_d, -1)),
11546                        false,
11547                        false,
11548                        Some((&mut aq8, &mut ad8)),
11549                    )?;
11550                    fa_q8 = Some((aq8, ad8));
11551                } else if swa && b_swa > win && hd == 256 && rows_on {
11552                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11553                    e.fa_decode_rows_w(
11554                        &q,
11555                        &k_view,
11556                        &v_view,
11557                        &mut attn,
11558                        hd,
11559                        nh,
11560                        nkv,
11561                        &kvl.len_d,
11562                        -1,
11563                        1,
11564                        scale,
11565                        win,
11566                        kvl.k_tok_bytes,
11567                        kvl.v_tok_bytes,
11568                        Some((&mut aq8, &mut ad8)),
11569                    )?;
11570                    fa_q8 = Some((aq8, ad8));
11571                } else {
11572                    let b = if swa { b_swa } else { b_glob };
11573                    e.fa_decode_dc(
11574                        &q,
11575                        &k_view,
11576                        &v_view,
11577                        &mut attn,
11578                        hd,
11579                        nh,
11580                        nkv,
11581                        &kvl.len_d,
11582                        b,
11583                        scale,
11584                        kvl.k_tok_bytes,
11585                        kvl.v_tok_bytes,
11586                        swa && crate::Engine::wkv_on(),
11587                    )?;
11588                }
11589            }
11590        }
11591        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
11592        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
11593        if let Some((aq8, ad8)) = fa_q8 {
11594            let mut y = e.uninit(fa.wo.out_features())?;
11595            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11596            return Ok(y);
11597        }
11598        Ok(e.matmul(&fa.wo, &attn, 1)?)
11599    }
11600
11601    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
11602    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
11603    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
11604    /// views in-graph); caller gates and falls back to the dc-eager loop.
11605    pub fn gemma4_generate_graph(
11606        &self,
11607        e: &Engine,
11608        prompt_pos: usize,
11609        first_token: u32,
11610        cache: &mut Cache,
11611        max_new: usize,
11612        eos: &[u32],
11613        mut on_token: impl FnMut(u32) -> bool,
11614    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11615        if self.is_gemma4_e4b() {
11616            return Err(
11617                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11618                    .into(),
11619            );
11620        }
11621        use crate::decode::StopReason;
11622        let n_vocab = self.output.out_features();
11623        let n_embd = self.cfg.n_embd as usize;
11624        let embd_gpu = self
11625            .embd_gpu
11626            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11627        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11628        for kvl in cache.kv.iter_mut().flatten() {
11629            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11630        }
11631        let mut token_d = e.stream().clone_htod(&[first_token])?;
11632        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11633        let g4 = self.cfg.gemma4.as_ref().unwrap();
11634        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11635        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
11636        let nkv_s = g4
11637            .head_count_kv
11638            .iter()
11639            .zip(g4.swa_pattern.iter())
11640            .find(|p| *p.1)
11641            .map(|p| *p.0 as usize)
11642            .unwrap_or(8);
11643        let nkv_g = g4
11644            .head_count_kv
11645            .iter()
11646            .zip(g4.swa_pattern.iter())
11647            .find(|p| !*p.1)
11648            .map(|p| *p.0 as usize)
11649            .unwrap_or(2);
11650        let mut graphs: std::collections::HashMap<
11651            ((bool, usize), (bool, usize), bool, bool),
11652            (
11653                cudarc::driver::CudaGraph,
11654                Vec<Box<dyn std::any::Any + Send>>,
11655            ),
11656        > = Default::default();
11657        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
11658        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
11659        let mut slots = self.g4_dc_slots(e)?;
11660        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
11661        // baked at the door entry (the modulo keeps every capture valid indefinitely).
11662        const RING: usize = 64;
11663        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
11664        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
11665        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
11666        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
11667        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
11668        const DRAIN: usize = 1;
11669        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11670        let ring_base = prompt_pos;
11671        let mut out = Vec::with_capacity(max_new);
11672        let mut reason = StopReason::MaxNew;
11673        let mut next = first_token;
11674        let mut captures = 0usize;
11675        for _ in 0..max_new {
11676            out.push(next);
11677            if eos.contains(&next) {
11678                reason = StopReason::Eos;
11679                break;
11680            }
11681            if !on_token(next) {
11682                reason = StopReason::Callback;
11683                break;
11684            }
11685            let t_kv = cache.pos + 1;
11686            // Bucket key per ARM (graph arc step 3):
11687            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
11688            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
11689            //    the component collapses to a single marker).
11690            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
11691            //    at/above it — the kernel derives splits from len_d per replay, so buckets
11692            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
11693            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11694            let f512 = crate::fa512_min_tkv();
11695            let key_s = if t_kv > win {
11696                (true, usize::MAX)
11697            } else {
11698                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11699            };
11700            let (key_g, rung_end) = if t_kv >= f512 {
11701                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
11702                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
11703                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11704                ((true, end), end)
11705            } else {
11706                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11707            };
11708            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11709            if !graphs.contains_key(&key) {
11710                let bucket_max = (t_kv, rung_end);
11711                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
11712                let snap = cache.snapshot(e)?;
11713                let pos_save = e.dtoh_i32_one(&pos_d)?;
11714                let len_save: Vec<Option<i32>> = cache
11715                    .kv
11716                    .iter()
11717                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11718                    .collect();
11719                let tok_save = e.dtoh_u32_one(&token_d)?;
11720                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
11721                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
11722                // regression class, and this door's measured -8.8%. The keeper pins warmup
11723                // transients so the captured graph holds kernel nodes only.
11724                let graph = {
11725                    let tok_ref = &mut token_d;
11726                    let pos_ref = &mut pos_d;
11727                    let cache_ref = &mut *cache;
11728                    let slots_ref = &mut slots;
11729                    let ring_ref = &mut ring;
11730                    e.capture_graph_retained_flags(
11731                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11732                        |e| {
11733                        // self-feeding: the argmax writes token_d itself.
11734                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11735                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11736                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11737                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11738                                                           cache_ref, n_vocab, Some(bucket_max),
11739                                                           sl, tok_ref, Some((rg, ring_base)))
11740                    })?
11741                };
11742                cache.rollback(e, &snap, 0)?;
11743                e.set_i32_one(&mut pos_d, pos_save)?;
11744                for (il, ls) in len_save.iter().enumerate() {
11745                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11746                        e.set_i32_one(&mut kvl.len_d, *v)?;
11747                    }
11748                }
11749                e.set_u32_one(&mut token_d, tok_save)?;
11750                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11751                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11752                        eprintln!("[graph-census] {c:?}");
11753                    }
11754                }
11755                graphs.insert(key, graph);
11756                captures += 1;
11757            }
11758            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
11759            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
11760            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
11761            // the budget; capture warmups already emitted their tokens through the ring.
11762            let mut chunk = 1usize;
11763            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11764                .ok()
11765                .and_then(|v| v.parse().ok())
11766                .unwrap_or(DRAIN);
11767            while chunk < drain_cap && out.len() + chunk < max_new {
11768                let t_next = cache.pos + 1 + chunk;
11769                let key_s2 = if t_next > win {
11770                    (true, usize::MAX)
11771                } else {
11772                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11773                };
11774                let key_g2 = if t_next >= f512 {
11775                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11776                } else {
11777                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11778                };
11779                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11780                    break;
11781                }
11782                chunk += 1;
11783            }
11784            let g = &graphs.get(&key).unwrap().0;
11785            for _ in 0..chunk {
11786                g.launch()?;
11787            }
11788            e.stream().synchronize()?;
11789            let ringh = e.dtoh_u32(&ring)?;
11790            for j in 0..chunk {
11791                let pos_j = cache.pos + j;
11792                let tok_j = ringh[(pos_j - ring_base) % RING];
11793                cache.pos += 0; // advanced below in one shot
11794                if j + 1 == chunk {
11795                    next = tok_j;
11796                } else {
11797                    out.push(tok_j);
11798                    if eos.contains(&tok_j) || !on_token(tok_j) {
11799                        reason = if eos.contains(&tok_j) {
11800                            StopReason::Eos
11801                        } else {
11802                            StopReason::Callback
11803                        };
11804                        // roll device/host state back to the stop point.
11805                        let keep = cache.pos + j + 1;
11806                        e.set_i32_one(&mut pos_d, keep as i32)?;
11807                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11808                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11809                            kvl.len = keep;
11810                        }
11811                        cache.pos = keep;
11812                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11813                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11814                        }
11815                        return Ok((out, reason));
11816                    }
11817                }
11818            }
11819            cache.pos += chunk;
11820            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11821                kvl.len += chunk;
11822            }
11823        }
11824        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11825            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11826        }
11827        Ok((out, reason))
11828    }
11829
11830    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
11831    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
11832    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
11833    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
11834    /// logits (host) + advances cache.pos by t.
11835    pub(crate) fn gemma4_decode_step_t(
11836        &self,
11837        e: &Engine,
11838        tokens: &[u32],
11839        pos0: usize,
11840        cache: &mut Cache,
11841    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11842        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11843    }
11844
11845    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
11846    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
11847    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
11848    pub(crate) fn gemma4_decode_step_t_am(
11849        &self,
11850        e: &Engine,
11851        tokens: &[u32],
11852        pos0: usize,
11853        cache: &mut Cache,
11854    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11855        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11856        let t = tokens.len();
11857        let n_vocab = self.output.out_features();
11858        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11859        for i in 0..t {
11860            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11861        }
11862        Ok((e.dtoh_u32(&toks)?, hn))
11863    }
11864
11865    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
11866    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
11867    pub(crate) fn gemma4_decode_step_t_am_dev(
11868        &self,
11869        e: &Engine,
11870        tok_d: &CudaSlice<u32>,
11871        t: usize,
11872        pos0: usize,
11873        cache: &mut Cache,
11874    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11875        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11876        let n_vocab = self.output.out_features();
11877        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11878        for i in 0..t {
11879            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11880        }
11881        Ok((vam, hn))
11882    }
11883
11884    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
11885    /// llama's h_nextn convention).
11886    pub(crate) fn gemma4_decode_step_t_h(
11887        &self,
11888        e: &Engine,
11889        tokens: &[u32],
11890        pos0: usize,
11891        cache: &mut Cache,
11892    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11893        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11894        let t = tokens.len();
11895        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11896        e.softcap(&mut ld, cap, t * self.output.out_features())?;
11897        Ok((e.dtoh(&ld)?, hn))
11898    }
11899
11900    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
11901    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
11902    pub(crate) fn verify_stream_scratch(
11903        &self,
11904        e: &Engine,
11905        cap: usize,
11906    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11907        Ok(VerifyStreamScratch {
11908            pos_d: e.htod_i32(&vec![0i32; cap])?,
11909            row_ctrs: (0..cap)
11910                .map(|_| e.htod_i32(&[0]))
11911                .collect::<Result<_, _>>()?,
11912        })
11913    }
11914
11915    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
11916    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
11917    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
11918    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
11919    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
11920    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
11921    /// sync, exactly the turnaround the burst exists to remove.
11922    pub(crate) fn gemma4_verify_t_am_stream(
11923        &self,
11924        e: &Engine,
11925        tok_d: &CudaSlice<u32>,
11926        t: usize,
11927        ctr: &CudaSlice<i32>,
11928        hint: usize,
11929        cache: &mut Cache,
11930        scr: &mut VerifyStreamScratch,
11931    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11932        let n_embd = self.cfg.n_embd as usize;
11933        let eps = self.cfg.rms_eps;
11934        assert!(t <= scr.row_ctrs.len() && t <= 64);
11935        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11936        for i in 0..t {
11937            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11938        }
11939        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11940        let embd_gpu = self
11941            .embd_gpu
11942            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11943        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11944        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11945        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11946        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11947        let n_layers = self.layers.len();
11948        for (il, layer) in self.layers.iter().enumerate() {
11949            let (hq, hdq) = match h_carry.take() {
11950                Some(p) => p,
11951                None => {
11952                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11953                }
11954            };
11955            let Mixer::Full(fa) = &layer.mixer else {
11956                panic!("gemma4 layer {il} not full-attn")
11957            };
11958            let o = self
11959                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11960            let next_norm = if il + 1 < n_layers {
11961                Some(self.layers[il + 1].attn_norm.float_data())
11962            } else {
11963                None
11964            };
11965            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
11966            x = xn;
11967            h_carry = hn;
11968            self.dflash_tap(e, cache, il, &x, t)?;
11969        }
11970        let mut hn = e.uninit(t * n_embd)?;
11971        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11972        let ld = e.matmul(&self.output, &hn, t)?;
11973        let n_vocab = self.output.out_features();
11974        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11975        for i in 0..t {
11976            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11977        }
11978        Ok((vam, hn))
11979    }
11980
11981    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
11982    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
11983    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
11984    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
11985    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
11986    /// kernel later if it shows in the profile).
11987    pub(crate) fn dflash_tap(
11988        &self,
11989        e: &Engine,
11990        cache: &mut Cache,
11991        il: usize,
11992        x: &CudaSlice<f32>,
11993        t: usize,
11994    ) -> Result<(), Box<dyn std::error::Error>> {
11995        let Some(taps) = cache.dflash_taps.as_mut() else {
11996            return Ok(());
11997        };
11998        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11999            return Ok(());
12000        };
12001        let h = taps.hidden;
12002        let n_taps = taps.layer_ids.len();
12003        let base = taps.base;
12004        debug_assert!(
12005            base + t <= taps.t,
12006            "tap window {base}+{t} exceeds sink {}",
12007            taps.t
12008        );
12009        let xv = e.view(x, t * h);
12010        for r in 0..t {
12011            let row = xv.slice(r * h..(r + 1) * h);
12012            e.copy_view_into(&mut taps.buf, (base + r) * n_taps * h + slot * h, &row, h)?;
12013        }
12014        Ok(())
12015    }
12016
12017    fn gemma4_verify_trunk(
12018        &self,
12019        e: &Engine,
12020        tokens: &[u32],
12021        pos0: usize,
12022        cache: &mut Cache,
12023        tok_dev: Option<&CudaSlice<u32>>,
12024    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12025        let n_embd = self.cfg.n_embd as usize;
12026        let eps = self.cfg.rms_eps;
12027        let t = tokens.len();
12028        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
12029        let pos_d = e.htod_i32(&pos)?;
12030        let mut x = match tok_dev {
12031            Some(td) => {
12032                let embd_gpu = self
12033                    .embd_gpu
12034                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
12035                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
12036                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
12037            }
12038            None => e.htod(&self.embd.gather(n_embd, tokens))?,
12039        };
12040        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
12041        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12042        let n_layers = self.layers.len();
12043        for (il, layer) in self.layers.iter().enumerate() {
12044            let (hq, hdq) = match h_carry.take() {
12045                Some(p) => p,
12046                None => {
12047                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
12048                }
12049            };
12050            let Mixer::Full(fa) = &layer.mixer else {
12051                panic!("gemma4 layer {il} not full-attn")
12052            };
12053            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
12054            let next_norm = if il + 1 < n_layers {
12055                Some(self.layers[il + 1].attn_norm.float_data())
12056            } else {
12057                None
12058            };
12059            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, t, next_norm)?;
12060            x = xn;
12061            h_carry = hn;
12062            self.dflash_tap(e, cache, il, &x, t)?;
12063        }
12064        let mut hn = e.uninit(t * n_embd)?;
12065        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
12066        let mut ld = e.matmul(&self.output, &hn, t)?;
12067        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
12068        cache.pos += t;
12069        Ok((ld, hn))
12070    }
12071
12072    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
12073    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
12074    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
12075    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
12076    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
12077    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
12078    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
12079    #[allow(clippy::too_many_arguments)]
12080    fn gemma4_verify_attn_stream(
12081        &self,
12082        e: &Engine,
12083        fa: &crate::hybrid::FullAttnLayer,
12084        il: usize,
12085        hq: &CudaSlice<i8>,
12086        hdq: &CudaSlice<f32>,
12087        pos_d: &CudaSlice<i32>,
12088        t: usize,
12089        cache: &mut Cache,
12090        hint: usize,
12091        row_ctrs: &[CudaSlice<i32>],
12092    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12093        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12094        let eps = self.cfg.rms_eps;
12095        let aux = self.gemma4_aux.as_ref().unwrap();
12096        let ones = aux.ones(e);
12097        #[cfg(debug_assertions)]
12098        crate::debug_assert_tensor_stream_device(
12099            ones,
12100            &e.stream(),
12101            "gemma4_verify_attn_stream.ones",
12102        );
12103        let h0 = e.zeros(0)?;
12104        let h = &h0;
12105        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
12106        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
12107        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12108        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12109        let fused_qkv = if f2b {
12110            if swa {
12111                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12112                    .map(|(a, b, c)| (a, b, Some(c)))
12113            } else {
12114                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12115                    .map(|(a, b)| (a, b, None))
12116            }
12117        } else {
12118            None
12119        };
12120        let (q0, k0, v0) = match fused_qkv {
12121            Some((a, b, cv)) => {
12122                let v = match cv {
12123                    Some(c) => c,
12124                    None => e.clone_dtod(&b)?,
12125                };
12126                (a, b, v)
12127            }
12128            None => {
12129                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12130                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12131                let v0 = if swa {
12132                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12133                } else {
12134                    e.clone_dtod(&k0)?
12135                };
12136                (q0, k0, v0)
12137            }
12138        };
12139        let mut q = e.uninit(t * nh * hd)?;
12140        let mut k = e.uninit(t * nkv * hd)?;
12141        let mut v = e.uninit(t * nkv * hd)?;
12142        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
12143        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
12144        let ff = if swa {
12145            None
12146        } else {
12147            Some(
12148                aux.rope_freqs(e)
12149                    .expect("gemma4 global rope needs rope_freqs.weight"),
12150            )
12151        };
12152        #[cfg(debug_assertions)]
12153        if let Some(ff) = ff {
12154            crate::debug_assert_tensor_stream_device(
12155                ff,
12156                &e.stream(),
12157                "gemma4_verify_attn_stream.rope_freqs",
12158            );
12159        }
12160        e.rms_norm_qkv_rope(
12161            &q0,
12162            &k0,
12163            &v0,
12164            fa.q_norm.float_data(),
12165            fa.k_norm.float_data(),
12166            ones,
12167            &mut q,
12168            &mut k,
12169            &mut v,
12170            hd,
12171            self.gemma4_rope_dims(il),
12172            nh * t,
12173            nkv * t,
12174            pos_d,
12175            nh,
12176            nkv,
12177            base,
12178            1.0,
12179            ff,
12180            eps,
12181        )?;
12182        let kvl = cache.kv[il].as_mut().unwrap();
12183        // append at the DEVICE slot; the counter advances by t on-device.
12184        e.append_kv_quantized_rows_dc(
12185            &k,
12186            &v,
12187            &mut kvl.k,
12188            &mut kvl.v,
12189            &kvl.len_d,
12190            t,
12191            kvl.kv_dim_k,
12192            kvl.kv_dim_v,
12193            kvl.k_tok_bytes,
12194            kvl.v_tok_bytes,
12195            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12196        )?;
12197        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
12198        // the sole len writer after this round's attention (base stays = old len, plus = 0).
12199        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12200        let mut attn = e.uninit(t * nh * hd)?;
12201        let k_view = e.view_u8(&kvl.k, kvl.k.len());
12202        let v_view = e.view_u8(&kvl.v, kvl.v.len());
12203        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
12204        // and a stable window regime — the same rung/regime keys as the draft graph).
12205        if swa && hint + 1 >= win {
12206            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
12207            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
12208            e.fa_decode_rows_w(
12209                &q,
12210                &k_view,
12211                &v_view,
12212                &mut attn,
12213                hd,
12214                nh,
12215                nkv,
12216                &kvl.len_d,
12217                0,
12218                t,
12219                scale,
12220                win,
12221                kvl.k_tok_bytes,
12222                kvl.v_tok_bytes,
12223                None,
12224            )?;
12225        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
12226            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
12227            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
12228            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
12229            // Burst entry gates the horizon onto one side of the crossover, so hint decides
12230            // for every row.
12231            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
12232            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
12233            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
12234            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
12235            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
12236            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
12237            // any bucket >= the live length is exact.
12238            let bucket = (hint + t + 2)
12239                .next_power_of_two()
12240                .min(crate::fa512_min_tkv().saturating_sub(1));
12241            let qv = e.view(&q, t * nh * hd);
12242            for i in 0..t {
12243                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
12244                let mut q_one = e.uninit(nh * hd)?;
12245                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12246                let mut a_one = e.uninit(nh * hd)?;
12247                e.fa_decode_dc(
12248                    &q_one,
12249                    &k_view,
12250                    &v_view,
12251                    &mut a_one,
12252                    hd,
12253                    nh,
12254                    nkv,
12255                    &row_ctrs[i],
12256                    bucket,
12257                    scale,
12258                    kvl.k_tok_bytes,
12259                    kvl.v_tok_bytes,
12260                    false,
12261                )?;
12262                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12263            }
12264        } else if hd == 512 {
12265            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
12266            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
12267            e.fa_decode_rows(
12268                &q,
12269                &k_view,
12270                &v_view,
12271                &mut attn,
12272                hd,
12273                nh,
12274                nkv,
12275                hint,
12276                t,
12277                scale,
12278                kvl.k_tok_bytes,
12279                kvl.v_tok_bytes,
12280                Some((&kvl.len_d, 0)),
12281                false,
12282                false,
12283                None,
12284            )?;
12285        } else {
12286            // hd256 under-window: v4 device-len rows twin.
12287            e.fa_decode_rows_dc(
12288                &q,
12289                &k_view,
12290                &v_view,
12291                &mut attn,
12292                hd,
12293                nh,
12294                nkv,
12295                &kvl.len_d,
12296                hint + t,
12297                t,
12298                scale,
12299                kvl.k_tok_bytes,
12300                kvl.v_tok_bytes,
12301                0,
12302                swa && crate::Engine::wkv_on(),
12303            )?;
12304        }
12305        Ok(e.matmul(&fa.wo, &attn, t)?)
12306    }
12307
12308    fn gemma4_verify_attn(
12309        &self,
12310        e: &Engine,
12311        fa: &crate::hybrid::FullAttnLayer,
12312        il: usize,
12313        hq: &CudaSlice<i8>,
12314        hdq: &CudaSlice<f32>,
12315        pos_d: &CudaSlice<i32>,
12316        t: usize,
12317        cache: &mut Cache,
12318    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12319        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
12320        let eps = self.cfg.rms_eps;
12321        let aux = self.gemma4_aux.as_ref().unwrap();
12322        let ones = aux.ones(e);
12323        #[cfg(debug_assertions)]
12324        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
12325        let n_embd = self.cfg.n_embd as usize;
12326        let _ = n_embd;
12327
12328        let h0 = e.zeros(0)?;
12329        let h = &h0;
12330        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
12331        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
12332        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12333        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
12334        let fused_qkv = if f2b {
12335            if swa {
12336                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
12337                    .map(|(a, b, c)| (a, b, Some(c)))
12338            } else {
12339                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
12340                    .map(|(a, b)| (a, b, None))
12341            }
12342        } else {
12343            None
12344        };
12345        let (q0, k0, v0) = match fused_qkv {
12346            Some((a, b, cv)) => {
12347                let v = match cv {
12348                    Some(c) => c,
12349                    None => e.clone_dtod(&b)?,
12350                };
12351                (a, b, v)
12352            }
12353            None => {
12354                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
12355                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
12356                let v0 = if swa {
12357                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
12358                } else {
12359                    e.clone_dtod(&k0)?
12360                };
12361                (q0, k0, v0)
12362            }
12363        };
12364        let mut q = e.uninit(t * nh * hd)?;
12365        let mut k = e.uninit(t * nkv * hd)?;
12366        let mut v = e.uninit(t * nkv * hd)?;
12367        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
12368        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
12369        let ff = if swa {
12370            None
12371        } else {
12372            Some(
12373                aux.rope_freqs(e)
12374                    .expect("gemma4 global rope needs rope_freqs.weight"),
12375            )
12376        };
12377        #[cfg(debug_assertions)]
12378        if let Some(ff) = ff {
12379            crate::debug_assert_tensor_stream_device(
12380                ff,
12381                &e.stream(),
12382                "gemma4_verify_attn.rope_freqs",
12383            );
12384        }
12385        e.rms_norm_qkv_rope(
12386            &q0,
12387            &k0,
12388            &v0,
12389            fa.q_norm.float_data(),
12390            fa.k_norm.float_data(),
12391            ones,
12392            &mut q,
12393            &mut k,
12394            &mut v,
12395            hd,
12396            self.gemma4_rope_dims(il),
12397            nh * t,
12398            nkv * t,
12399            pos_d,
12400            nh,
12401            nkv,
12402            base,
12403            1.0,
12404            ff,
12405            eps,
12406        )?;
12407        let kvl = cache.kv[il].as_mut().unwrap();
12408        let base_len = kvl.len;
12409        e.append_kv_quantized_rows(
12410            &k,
12411            &v,
12412            &mut kvl.k,
12413            &mut kvl.v,
12414            base_len,
12415            t,
12416            kvl.kv_dim_k,
12417            kvl.kv_dim_v,
12418            kvl.k_tok_bytes,
12419            kvl.v_tok_bytes,
12420            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
12421        )?;
12422        kvl.len += t;
12423        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
12424        let mut attn = e.uninit(t * nh * hd)?;
12425        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
12426        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
12427        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
12428            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
12429            // decode rides the SAME symbol at t=1 (parity law).
12430            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
12431        if rows_ok && (!swa || base_len + t <= win) {
12432            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12433            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12434            if hd == 512 {
12435                // device-len twin: sync the counter to the verify base (async arg-store).
12436                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12437                e.fa_decode_rows(
12438                    &q,
12439                    &k_view,
12440                    &v_view,
12441                    &mut attn,
12442                    hd,
12443                    nh,
12444                    nkv,
12445                    base_len,
12446                    t,
12447                    scale,
12448                    kvl.k_tok_bytes,
12449                    kvl.v_tok_bytes,
12450                    Some((&kvl.len_d, 0)),
12451                    false,
12452                    swa && crate::Engine::wkv_on(),
12453                    None,
12454                )?;
12455            } else {
12456                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
12457                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
12458                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
12459                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12460                e.fa_decode_rows_dc(
12461                    &q,
12462                    &k_view,
12463                    &v_view,
12464                    &mut attn,
12465                    hd,
12466                    nh,
12467                    nkv,
12468                    &kvl.len_d,
12469                    base_len + t,
12470                    t,
12471                    scale,
12472                    kvl.k_tok_bytes,
12473                    kvl.v_tok_bytes,
12474                    0,
12475                    swa && crate::Engine::wkv_on(),
12476                )?;
12477            }
12478            return Ok(e.matmul(&fa.wo, &attn, t)?);
12479        }
12480        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
12481        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
12482        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
12483        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
12484        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
12485        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
12486        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
12487        if hd == 256
12488            && swa
12489            && base_len + 1 >= win
12490            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12491        {
12492            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
12493            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
12494            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
12495            e.fa_decode_rows_w(
12496                &q,
12497                &k_view,
12498                &v_view,
12499                &mut attn,
12500                hd,
12501                nh,
12502                nkv,
12503                &kvl.len_d,
12504                0,
12505                t,
12506                scale,
12507                win,
12508                kvl.k_tok_bytes,
12509                kvl.v_tok_bytes,
12510                None,
12511            )?;
12512            return Ok(e.matmul(&fa.wo, &attn, t)?);
12513        }
12514        for i in 0..t {
12515            let avail = base_len + i + 1;
12516            let (off_tok, t_kv) = if swa && avail > win {
12517                (avail - win, win)
12518            } else {
12519                (0, avail)
12520            };
12521            let k_view = e.view_u8_range(
12522                &kvl.k,
12523                off_tok * kvl.k_tok_bytes,
12524                (off_tok + t_kv) * kvl.k_tok_bytes,
12525            );
12526            let v_view = e.view_u8_range(
12527                &kvl.v,
12528                off_tok * kvl.v_tok_bytes,
12529                (off_tok + t_kv) * kvl.v_tok_bytes,
12530            );
12531            let qi = e.view(&q, t * nh * hd);
12532            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12533            let mut q_one = e.uninit(nh * hd)?;
12534            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12535            let mut a_one = e.uninit(nh * hd)?;
12536            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
12537            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
12538            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
12539            if swa
12540                && avail > win
12541                && hd == 256
12542                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12543            {
12544                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12545                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12546                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12547                e.fa_decode_rows_w(
12548                    &q_one,
12549                    &kp,
12550                    &vp,
12551                    &mut a_one,
12552                    hd,
12553                    nh,
12554                    nkv,
12555                    &kvl.len_d,
12556                    0,
12557                    1,
12558                    scale,
12559                    win,
12560                    kvl.k_tok_bytes,
12561                    kvl.v_tok_bytes,
12562                    None,
12563                )?;
12564            } else if !swa
12565                && hd == 512
12566                && avail >= crate::fa512_min_tkv()
12567                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12568            {
12569                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12570                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12571                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12572                e.fa_decode_rows(
12573                    &q_one,
12574                    &kp,
12575                    &vp,
12576                    &mut a_one,
12577                    hd,
12578                    nh,
12579                    nkv,
12580                    avail - 1,
12581                    1,
12582                    scale,
12583                    kvl.k_tok_bytes,
12584                    kvl.v_tok_bytes,
12585                    Some((&kvl.len_d, 0)),
12586                    false,
12587                    false,
12588                    None,
12589                )?;
12590            } else {
12591                e.fa_decode_kvmod(
12592                    &q_one,
12593                    &k_view,
12594                    &v_view,
12595                    &mut a_one,
12596                    hd,
12597                    nh,
12598                    nkv,
12599                    t_kv,
12600                    scale,
12601                    kvl.k_tok_bytes,
12602                    kvl.v_tok_bytes,
12603                    swa && crate::Engine::wkv_on(),
12604                )?;
12605            }
12606            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12607        }
12608        Ok(e.matmul(&fa.wo, &attn, t)?)
12609    }
12610
12611    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
12612    /// h_seed = pre-output_norm hidden). Advances cache.pos.
12613    pub(crate) fn gemma4_decode_step_h(
12614        &self,
12615        e: &Engine,
12616        token: u32,
12617        cache: &mut Cache,
12618    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12619        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
12620        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
12621        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
12622        // unsplit rather than guessing a fence.
12623        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12624            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12625        }
12626        if crate::pp::pp_cuts(self.layers.len()).is_some() {
12627            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12628        }
12629        let n_embd = self.cfg.n_embd as usize;
12630        let eps = self.cfg.rms_eps;
12631        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12632        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12633        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12634        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
12635        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
12636        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12637        let n_layers = self.layers.len();
12638        for (il, layer) in self.layers.iter().enumerate() {
12639            let (hq, hdq) = match h_carry.take() {
12640                Some(p) => p,
12641                None => {
12642                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12643                }
12644            };
12645            let Mixer::Full(fa) = &layer.mixer else {
12646                panic!("gemma4 layer {il} not full-attn")
12647            };
12648            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12649            let next_norm = if il + 1 < n_layers {
12650                Some(self.layers[il + 1].attn_norm.float_data())
12651            } else {
12652                None
12653            };
12654            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12655            x = xn;
12656            h_carry = hn;
12657        }
12658        let mut hn = e.uninit(n_embd)?;
12659        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12660        let h_seed = e.clone_dtod(&x)?;
12661        let mut ld = e.matmul(&self.output, &hn, 1)?;
12662        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12663        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
12664        self.gemma4_suppress(e, &mut ld, 1)?;
12665        let logits = e.dtoh(&ld)?;
12666        cache.pos += 1;
12667        Ok((logits, h_seed))
12668    }
12669
12670    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
12671    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
12672    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
12673    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
12674    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
12675    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
12676    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
12677    fn gemma4_decode_layers(
12678        &self,
12679        e: &Engine,
12680        mut x: CudaSlice<f32>,
12681        lo: usize,
12682        hi: usize,
12683        pos_d: &CudaSlice<i32>,
12684        cache: &mut Cache,
12685    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12686        let n_embd = self.cfg.n_embd as usize;
12687        let eps = self.cfg.rms_eps;
12688        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12689        for il in lo..hi {
12690            let layer = &self.layers[il];
12691            let (hq, hdq) = match h_carry.take() {
12692                Some(p) => p,
12693                // range head: il == lo — norm against THIS layer's attn_norm.
12694                None => {
12695                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12696                }
12697            };
12698            let Mixer::Full(fa) = &layer.mixer else {
12699                panic!("gemma4 layer {il} not full-attn")
12700            };
12701            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12702            let next_norm = if il + 1 < hi {
12703                Some(self.layers[il + 1].attn_norm.float_data())
12704            } else {
12705                None
12706            };
12707            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, 1, next_norm)?;
12708            x = xn;
12709            h_carry = hn;
12710        }
12711        Ok(x)
12712    }
12713
12714    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
12715    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
12716    /// boundary handoff — same choreography as the generic arm (decode.rs), same
12717    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
12718    /// stage 1 = layers [split, n) + output_norm + softcapped head.
12719    /// Each stage uploads its own copy of the step's position scalar on its own stream.
12720    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
12721    fn gemma4_decode_step_h_pp2(
12722        &self,
12723        e: &Engine,
12724        token: u32,
12725        cache: &mut Cache,
12726        split: usize,
12727    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12728        if crate::pp::pp2_streams_off() {
12729            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12730        }
12731        let rt = crate::pp::Pp2Rt::get(e)?;
12732        let e0 = rt.engine(0, e);
12733        let e1 = rt.engine(1, e);
12734        let n_embd = self.cfg.n_embd as usize;
12735        let eps = self.cfg.rms_eps;
12736        let pos = cache.pos as i32;
12737
12738        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
12739        let slot = {
12740            let _st0 = rt.enter(0);
12741            let pos_d = e0.htod_i32(&[pos])?;
12742            #[cfg(debug_assertions)]
12743            crate::debug_assert_tensor_stream_device(
12744                &pos_d,
12745                &e0.stream(),
12746                "gemma4_decode_step_h_pp2.stage0.pos_d",
12747            );
12748            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12749            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12750            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12751            rt.tx(0, &x, n_embd)?
12752        };
12753
12754        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
12755        let _st1 = rt.enter(1);
12756        let pos_d = e1.htod_i32(&[pos])?;
12757        #[cfg(debug_assertions)]
12758        crate::debug_assert_tensor_stream_device(
12759            &pos_d,
12760            &e1.stream(),
12761            "gemma4_decode_step_h_pp2.stage1.pos_d",
12762        );
12763        let x = rt.rx(0, slot, n_embd)?;
12764        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12765
12766        let mut hn = e1.uninit(n_embd)?;
12767        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12768        let h_seed = e1.clone_dtod(&x)?;
12769        let mut ld = e1.matmul(&self.output, &hn, 1)?;
12770        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12771        e1.softcap(&mut ld, cap, self.output.out_features())?;
12772        self.gemma4_suppress(e1, &mut ld, 1)?;
12773        let logits = e1.dtoh(&ld)?;
12774        cache.pos += 1;
12775        Ok((logits, h_seed))
12776    }
12777
12778    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
12779    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
12780    fn gemma4_decode_step_h_pp2_samestream(
12781        &self,
12782        e: &Engine,
12783        token: u32,
12784        cache: &mut Cache,
12785        split: usize,
12786    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12787        let n_embd = self.cfg.n_embd as usize;
12788        let eps = self.cfg.rms_eps;
12789        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12790
12791        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
12792        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12793        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12794        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12795
12796        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
12797        let boundary_tx = e.clone_dtod(&x)?;
12798        let boundary_rx = e.clone_dtod(&boundary_tx)?;
12799
12800        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
12801        let x =
12802            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12803
12804        let mut hn = e.uninit(n_embd)?;
12805        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12806        let h_seed = e.clone_dtod(&x)?;
12807        let mut ld = e.matmul(&self.output, &hn, 1)?;
12808        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12809        e.softcap(&mut ld, cap, self.output.out_features())?;
12810        self.gemma4_suppress(e, &mut ld, 1)?;
12811        let logits = e.dtoh(&ld)?;
12812        cache.pos += 1;
12813        Ok((logits, h_seed))
12814    }
12815}
12816
12817// ============================ step35 (Step-3.7-Flash) ==================================
12818// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
12819// FAMILY and not a few branches inside the generic `full_attn*` chain:
12820//
12821//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
12822//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
12823//      shapes and the FA head counts would be wrong on 33 of 45 layers.
12824//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
12825//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
12826//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
12827//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
12828//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
12829//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
12830//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
12831//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
12832//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
12833//
12834// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
12835impl HybridModel {
12836    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
12837    /// synthesize a drafter or trunk layer from a neighboring class.
12838    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12839        let geometry = self
12840            .cfg
12841            .layer_geometry(il as u32)
12842            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12843        debug_assert_eq!(
12844            geometry.attention_gate,
12845            memra_gguf::config::AttentionGateKind::SeparateHead
12846        );
12847        geometry
12848    }
12849
12850    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
12851    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
12852    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
12853    ///
12854    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
12855    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
12856    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
12857    /// `cache`:
12858    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
12859    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
12860    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
12861    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
12862    ///     contract, lane/chunkinv-flip).
12863    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
12864    ///     q/k/v, no cache side effect.
12865    ///
12866    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
12867    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
12868    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
12869    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
12870    /// still contains must be masked per query. memra's window convention
12871    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
12872    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
12873    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
12874    ///
12875    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
12876    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
12877    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
12878    ///
12879    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
12880    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
12881    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
12882    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
12883    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
12884    /// hidden rows, and the generated text — a function of the chunk size:
12885    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
12886    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
12887    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
12888    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
12889    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
12890    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
12891    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
12892    ///   one-token change in a documented machine-config knob changed the answer.
12893    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
12894    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
12895    /// the same rows moves the logits by ~1.8.
12896    ///
12897    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
12898    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
12899    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
12900    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
12901    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
12902    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
12903    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
12904    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
12905    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
12906    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
12907    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
12908    /// those with t_kv <= win = 512.
12909    #[allow(clippy::too_many_arguments)]
12910    fn step35_attn_pre_wo(
12911        &self,
12912        e: &Engine,
12913        fa: &FullAttnLayer,
12914        mut g3: Vec<CudaSlice<f32>>,
12915        hg: Option<&CudaSlice<f32>>,
12916        gt_pre: Option<&CudaSlice<f32>>,
12917        pos_d: &CudaSlice<i32>,
12918        t: usize,
12919        cache: Option<&mut Cache>,
12920        il: usize,
12921        seq_end: usize,
12922    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12923        let geometry = self.step35_geom(il);
12924        let hd = geometry.head_dim_k as usize;
12925        let nkv = geometry.n_head_kv as usize;
12926        let nh = geometry.n_head as usize;
12927        let rbase = geometry.rope_base;
12928        let scale = geometry.attention_scale();
12929        let swa = geometry.window.is_some();
12930        let eps = self.cfg.rms_eps;
12931        let win = geometry.window.unwrap_or(0) as usize;
12932        let n_rot = geometry.n_rot as usize;
12933
12934        let v = g3.pop().unwrap();
12935        let k0 = g3.pop().unwrap();
12936        let q0 = g3.pop().unwrap();
12937
12938        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
12939        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
12940        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
12941        let mut q = e.uninit(t * nh * hd)?;
12942        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12943        let mut k = e.uninit(t * nkv * hd)?;
12944        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12945        let ff = if geometry.rope_factors {
12946            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12947        } else {
12948            None
12949        };
12950        #[cfg(debug_assertions)]
12951        if let Some(ff) = ff {
12952            crate::debug_assert_tensor_stream_device(
12953                ff,
12954                &e.stream(),
12955                "step35_attn_pre_wo.rope_freqs",
12956            );
12957        }
12958        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12959
12960        let mut attn = e.uninit(t * nh * hd)?;
12961        match cache {
12962            Some(cache) => {
12963                let base_len = cache.kv[il].as_ref().unwrap().len;
12964                // Read per layer call, never in a measured default.
12965                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12966                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12967                let off = if swa {
12968                    let raw = base_len.saturating_sub(win - 1);
12969                    if legacy_tkv || legacy_calllocal {
12970                        raw
12971                    } else {
12972                        raw & !31usize
12973                    }
12974                } else {
12975                    0
12976                };
12977                {
12978                    let kvl = cache.kv[il].as_mut().unwrap();
12979                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12980                    let write_row = e.prepare_kv_append(kvl, off, t)?;
12981                    e.append_kv_quantized_rows(
12982                        &k,
12983                        &v,
12984                        &mut kvl.k,
12985                        &mut kvl.v,
12986                        write_row,
12987                        t,
12988                        kvl.kv_dim_k,
12989                        kvl.kv_dim_v,
12990                        kvl.k_tok_bytes,
12991                        kvl.v_tok_bytes,
12992                        crate::Engine::kv_fp8_on(),
12993                    )?;
12994                    kvl.len += t;
12995                    let new_len = kvl.len as i32;
12996                    e.set_i32_one(&mut kvl.len_d, new_len)?;
12997                }
12998                let kvl = cache.kv[il].as_ref().unwrap();
12999                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
13000                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
13001                // unaligned view offset here. Both halves are load-bearing for the canaries:
13002                // on the FA default the predicate arms agree bitwise wherever they can differ
13003                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
13004                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
13005                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
13006                // on the current FA path: its tile grid starts at the chunk/call boundary.
13007                // SWA: trim the view to the oldest key any query in this chunk can reach —
13008                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
13009                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
13010                // kernel's online-softmax recurrence groups keys into BK tiles relative to
13011                // the VIEW START — so an unaligned off regroups the same absolute keys into
13012                // different tiles at different chunk sizes = different (m,l) rounding =
13013                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
13014                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
13015                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
13016                // size; the <=31 extra leading keys are older than EVERY query's window
13017                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
13018                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
13019                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
13020                // the floor arm's bits do not move either (gated: G2f, battery 2).
13021                let t_kv = base_len + t - off;
13022                let physical = kvl.physical_rows(off, off + t_kv)?;
13023                let k_view = e.view_u8_range(
13024                    &kvl.k,
13025                    physical.start * kvl.k_tok_bytes,
13026                    physical.end * kvl.k_tok_bytes,
13027                );
13028                let v_view = e.view_u8_range(
13029                    &kvl.v,
13030                    physical.start * kvl.v_tok_bytes,
13031                    physical.end * kvl.v_tok_bytes,
13032                );
13033                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
13034                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
13035                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
13036                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
13037                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
13038                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
13039                // construction, so the invariance assertion MUST break under it (the seam whose
13040                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
13041                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
13042                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
13043                // cached (probes flip it in-process). Never on in a measured default run.
13044                let swa_naive = if legacy_tkv {
13045                    t_kv > win
13046                } else {
13047                    seq_end > win
13048                };
13049                if swa && swa_naive {
13050                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
13051                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
13052                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
13053                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
13054                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
13055                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
13056                    // identically to the unwindowed one modulo the mask, which is the point.
13057                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
13058                    // selected on `seq_end` like every arm here, so the class is uniform for
13059                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
13060                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
13061                    // the f32 floor (the previous numeric config, kept as the A/B seam).
13062                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
13063                        e.sdpa_naive_w_quantized_view(
13064                            &q,
13065                            &k_view,
13066                            &v_view,
13067                            &mut attn,
13068                            hd,
13069                            nh,
13070                            nkv,
13071                            t,
13072                            t_kv,
13073                            scale,
13074                            true,
13075                            win,
13076                            kvl.k_tok_bytes,
13077                            kvl.v_tok_bytes,
13078                        )?;
13079                    } else {
13080                        e.fa_prefill_view_ws_w_hd128(
13081                            &q,
13082                            &k_view,
13083                            &v_view,
13084                            &mut attn,
13085                            hd,
13086                            nh,
13087                            nkv,
13088                            t,
13089                            t_kv,
13090                            scale,
13091                            true,
13092                            win,
13093                            kvl.k_tok_bytes,
13094                            kvl.v_tok_bytes,
13095                        )?;
13096                    }
13097                } else if std::env::var("MEMRA_NOFA").is_ok() {
13098                    e.sdpa_naive_quantized_view(
13099                        &q,
13100                        &k_view,
13101                        &v_view,
13102                        &mut attn,
13103                        hd,
13104                        nh,
13105                        nkv,
13106                        t,
13107                        t_kv,
13108                        scale,
13109                        true,
13110                        kvl.k_tok_bytes,
13111                        kvl.v_tok_bytes,
13112                    )?;
13113                } else {
13114                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
13115                    // reach past the window, so the window mask is a no-op under causal and every
13116                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
13117                    // request either way, which is what makes the chunk size arithmetic-free.
13118                    e.fa_prefill_view_ws(
13119                        &q,
13120                        &k_view,
13121                        &v_view,
13122                        &mut attn,
13123                        hd,
13124                        nh,
13125                        nkv,
13126                        t,
13127                        t_kv,
13128                        scale,
13129                        true,
13130                        kvl.k_tok_bytes,
13131                        kvl.v_tok_bytes,
13132                        crate::Engine::kv_fp8_on(),
13133                    )?;
13134                }
13135            }
13136            None => {
13137                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
13138                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
13139                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
13140                // seq_end here too or it re-opens the same door.
13141                debug_assert_eq!(
13142                    seq_end, t,
13143                    "step35 cacheless prefill is monolithic (seq_end == t)"
13144                );
13145                if swa && seq_end > win {
13146                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13147                } else if std::env::var("MEMRA_NOFA").is_ok() {
13148                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13149                } else {
13150                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13151                }
13152            }
13153        }
13154
13155        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
13156        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
13157        let gw = fa
13158            .attn_gate
13159            .as_ref()
13160            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13161        let gt_owned = if gt_pre.is_none() {
13162            Some(e.matmul(
13163                gw,
13164                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
13165                t,
13166            )?)
13167        } else {
13168            None
13169        };
13170        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
13171        let mut ag = e.uninit(t * nh * hd)?;
13172        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
13173        Ok(ag)
13174    }
13175
13176    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
13177    /// `forward_last`, t2probe). Post-`wo`.
13178    pub(crate) fn step35_attn(
13179        &self,
13180        e: &Engine,
13181        fa: &FullAttnLayer,
13182        h: &CudaSlice<f32>,
13183        pos_d: &CudaSlice<i32>,
13184        t: usize,
13185        il: usize,
13186    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13187        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
13188        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
13189        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
13190        Ok(e.matmul(&fa.wo, &ag, t)?)
13191    }
13192
13193    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
13194    /// resident quantized cache, attend through the cache view). Post-`wo`.
13195    ///
13196    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
13197    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
13198    /// own extent.
13199    #[allow(clippy::too_many_arguments)]
13200    pub(crate) fn step35_attn_prime(
13201        &self,
13202        e: &Engine,
13203        fa: &FullAttnLayer,
13204        h: &CudaSlice<f32>,
13205        hx: Option<&CudaSlice<u8>>,
13206        pos_d: &CudaSlice<i32>,
13207        t: usize,
13208        cache: &mut Cache,
13209        il: usize,
13210        seq_end: usize,
13211    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13212        let g3 = match hx {
13213            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
13214            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
13215        };
13216        let ag =
13217            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
13218        Ok(e.matmul(&fa.wo, &ag, t)?)
13219    }
13220
13221    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
13222    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
13223    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
13224    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
13225    /// requiring `attn_gate`).
13226    ///
13227    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
13228    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
13229    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
13230    #[allow(clippy::too_many_arguments)]
13231    pub(crate) fn step35_decode_attn(
13232        &self,
13233        e: &Engine,
13234        fa: &FullAttnLayer,
13235        il: usize,
13236        h: &CudaSlice<f32>,
13237        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
13238        pos_d: &CudaSlice<i32>,
13239        cache: &mut Cache,
13240    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13241        let geometry = self.step35_geom(il);
13242        let hd = geometry.head_dim_k as usize;
13243        let nkv = geometry.n_head_kv as usize;
13244        let nh = geometry.n_head as usize;
13245        let rbase = geometry.rope_base;
13246        let scale = geometry.attention_scale();
13247        let swa = geometry.window.is_some();
13248        let eps = self.cfg.rms_eps;
13249        let win = geometry.window.unwrap_or(0) as usize;
13250        let n_rot = geometry.n_rot as usize;
13251        let n_embd = self.cfg.n_embd as usize;
13252        let gw = fa
13253            .attn_gate
13254            .as_ref()
13255            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
13256
13257        let (q0, k0, v0, gt) = match pre_q {
13258            Some((hq, hdq)) => {
13259                debug_assert!(
13260                    e.uses_q8_1_fast(gw),
13261                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
13262                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
13263                );
13264                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
13265                    Some(t3) => t3,
13266                    None => (
13267                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
13268                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
13269                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
13270                    ),
13271                };
13272                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
13273                (a, b, c, gt)
13274            }
13275            None => {
13276                if e.uses_q8_1_fast(&fa.wq)
13277                    && e.uses_q8_1_fast(&fa.wk)
13278                    && e.uses_q8_1_fast(&fa.wv)
13279                    && e.uses_q8_1_fast(gw)
13280                {
13281                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
13282                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
13283                        Some(t3) => t3,
13284                        None => (
13285                            e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
13286                            e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
13287                            e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
13288                        ),
13289                    };
13290                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
13291                    (a, b, c, gt)
13292                } else {
13293                    (
13294                        e.matmul(&fa.wq, h, 1)?,
13295                        e.matmul(&fa.wk, h, 1)?,
13296                        e.matmul(&fa.wv, h, 1)?,
13297                        e.matmul(gw, h, 1)?,
13298                    )
13299                }
13300            }
13301        };
13302
13303        let mut q = e.uninit(nh * hd)?;
13304        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
13305        let mut k = e.uninit(nkv * hd)?;
13306        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
13307        let ff = if swa {
13308            None
13309        } else {
13310            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
13311        };
13312        #[cfg(debug_assertions)]
13313        if let Some(ff) = ff {
13314            crate::debug_assert_tensor_stream_device(
13315                ff,
13316                &e.stream(),
13317                "step35_decode_attn.rope_freqs",
13318            );
13319        }
13320        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
13321
13322        if std::env::var("MEMRA_NOFA").is_ok() {
13323            return Err(
13324                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
13325                        cache; unset MEMRA_NOFA to use fa_decode"
13326                    .into(),
13327            );
13328        }
13329        let kvl = cache.kv[il].as_mut().unwrap();
13330        let next_len = kvl.len + 1;
13331        let (off, t_kv) = if swa && next_len > win {
13332            (next_len - win, win)
13333        } else {
13334            (0, next_len)
13335        };
13336        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
13337        e.append_kv_quantized(
13338            &k,
13339            &v0,
13340            &mut kvl.k,
13341            &mut kvl.v,
13342            write_row,
13343            kvl.kv_dim_k,
13344            kvl.kv_dim_v,
13345            kvl.k_tok_bytes,
13346            kvl.v_tok_bytes,
13347            crate::Engine::kv_fp8_on(),
13348        )?;
13349        kvl.len = next_len;
13350        let physical = kvl.physical_rows(off, off + t_kv)?;
13351        let k_view = e.view_u8_range(
13352            &kvl.k,
13353            physical.start * kvl.k_tok_bytes,
13354            physical.end * kvl.k_tok_bytes,
13355        );
13356        let v_view = e.view_u8_range(
13357            &kvl.v,
13358            physical.start * kvl.v_tok_bytes,
13359            physical.end * kvl.v_tok_bytes,
13360        );
13361        let mut attn = e.uninit(nh * hd)?;
13362        e.fa_decode_kvmod(
13363            &q,
13364            &k_view,
13365            &v_view,
13366            &mut attn,
13367            hd,
13368            nh,
13369            nkv,
13370            t_kv,
13371            scale,
13372            kvl.k_tok_bytes,
13373            kvl.v_tok_bytes,
13374            crate::Engine::kv_fp8_on(),
13375        )?;
13376
13377        let mut ag = e.uninit(nh * hd)?;
13378        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
13379        Ok(e.matmul(&fa.wo, &ag, 1)?)
13380    }
13381}
13382
13383// ===================================================================================== //
13384//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
13385//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
13386//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
13387//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
13388//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
13389//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
13390// ===================================================================================== //
13391impl HybridModel {
13392    pub fn is_gemma4_e4b(&self) -> bool {
13393        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
13394    }
13395
13396    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
13397    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
13398    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
13399    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
13400        let g = self.cfg.gemma4.as_ref().unwrap();
13401        let swa = g.swa_pattern[il];
13402        let hd = if swa {
13403            g.key_length_swa
13404        } else {
13405            g.key_length_global
13406        } as usize;
13407        let Mixer::Full(fa) = &self.layers[il].mixer else {
13408            panic!("e4b layer {il} not full-attn")
13409        };
13410        let nh = fa.wq.out_features() / hd;
13411        let nkv = fa.wk.out_features() / hd;
13412        (
13413            hd,
13414            nkv,
13415            nh,
13416            if swa {
13417                g.rope_base_swa
13418            } else {
13419                g.rope_base_global
13420            },
13421            1.0,
13422            swa,
13423        )
13424    }
13425
13426    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
13427    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
13428        self.layers[il]
13429            .gemma4
13430            .as_ref()
13431            .and_then(|b| b.e4b.as_ref())
13432            .and_then(|e4| e4.kv_share.map(|t| t as usize))
13433    }
13434
13435    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
13436    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
13437    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
13438    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
13439    fn gemma4_e4b_inp_pl(
13440        &self,
13441        e: &Engine,
13442        tokens: &[u32],
13443        x_scaled: &CudaSlice<f32>,
13444        t: usize,
13445    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13446        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
13447        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
13448    }
13449
13450    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
13451    fn gemma4_e4b_inp_pl_dev(
13452        &self,
13453        e: &Engine,
13454        tok_d: &CudaSlice<u32>,
13455        x_scaled: &CudaSlice<f32>,
13456        t: usize,
13457    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13458        let aux = self.gemma4_aux.as_ref().unwrap();
13459        let m = aux.e4b.as_ref().unwrap();
13460        let n_embd = self.cfg.n_embd as usize;
13461        let n_layer = self.layers.len();
13462        let width = m.n_epl * n_layer;
13463        let tbl = m.tok_tbl_gpu.get_or_init(|| {
13464            e.upload_u8(&m.tok_embd_bytes)
13465                .expect("e4b per-layer token table upload")
13466        });
13467        let mut a =
13468            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
13469        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
13470        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
13471        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
13472        let mut pn = e.uninit(t * width)?;
13473        e.rms_norm(
13474            &p,
13475            m.proj_norm.float_data(),
13476            &mut pn,
13477            m.n_epl,
13478            t * n_layer,
13479            self.cfg.rms_eps,
13480        )?;
13481        let mut out = e.uninit(t * width)?;
13482        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
13483        Ok(out)
13484    }
13485
13486    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
13487    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
13488    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
13489    /// already holds this forward's rows — the target runs earlier in the stack).
13490    #[allow(clippy::too_many_arguments)]
13491    fn gemma4_e4b_attn(
13492        &self,
13493        e: &Engine,
13494        il: usize,
13495        hq: &CudaSlice<i8>,
13496        hdq: &CudaSlice<f32>,
13497        pos_d: &CudaSlice<i32>,
13498        t: usize,
13499        cache: &mut Cache,
13500        dc_bucket: Option<usize>,
13501    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
13502        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
13503        let eps = self.cfg.rms_eps;
13504        let aux = self.gemma4_aux.as_ref().unwrap();
13505        let ones = aux.ones(e);
13506        #[cfg(debug_assertions)]
13507        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
13508        let Mixer::Full(fa) = &self.layers[il].mixer else {
13509            unreachable!()
13510        };
13511        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
13512        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
13513        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
13514        let h0 = e.zeros(0)?;
13515        let h = &h0;
13516
13517        let ff = if swa {
13518            None
13519        } else {
13520            Some(
13521                aux.rope_freqs(e)
13522                    .expect("e4b global rope needs rope_freqs.weight"),
13523            )
13524        };
13525        #[cfg(debug_assertions)]
13526        if let Some(ff) = ff {
13527            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13528        }
13529        let share = self.gemma4_e4b_kv_target(il);
13530        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
13531        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13532        let mut q;
13533        if let Some(_tgt) = share {
13534            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13535            q = e.uninit(t * nh * hd)?;
13536            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
13537            // empty; q0 stands in for the unused k/v pointers).
13538            let mut kdummy = e.uninit(1)?;
13539            let mut vdummy = e.uninit(1)?;
13540            e.rms_norm_qkv_rope(
13541                &q0,
13542                &q0,
13543                &q0,
13544                fa.q_norm.float_data(),
13545                fa.q_norm.float_data(),
13546                ones,
13547                &mut q,
13548                &mut kdummy,
13549                &mut vdummy,
13550                hd,
13551                self.gemma4_rope_dims(il),
13552                nh * t,
13553                0,
13554                pos_d,
13555                nh,
13556                1,
13557                base,
13558                1.0,
13559                ff,
13560                eps,
13561            )?;
13562        } else {
13563            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
13564            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
13565            // q|k|v rows — the cat norm+rope twin consumes it directly.
13566            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13567            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13568            q = e.uninit(t * nh * hd)?;
13569            let mut k = e.uninit(t * nkv * hd)?;
13570            let mut v = e.uninit(t * nkv * hd)?;
13571            if t == 1 && cat.is_some() {
13572                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13573                e.rms_norm_qkv_rope_cat(
13574                    &qkv0,
13575                    fa.q_norm.float_data(),
13576                    fa.k_norm.float_data(),
13577                    ones,
13578                    &mut q,
13579                    &mut k,
13580                    &mut v,
13581                    hd,
13582                    self.gemma4_rope_dims(il),
13583                    nh,
13584                    nkv,
13585                    pos_d,
13586                    nh,
13587                    nkv,
13588                    base,
13589                    1.0,
13590                    ff,
13591                    eps,
13592                )?;
13593            } else {
13594                let (q0, k0, v0) = match if t == 1 {
13595                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13596                } else {
13597                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
13598                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
13599                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13600                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13601                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13602                    } else {
13603                        None
13604                    }
13605                } {
13606                    Some(triple) => triple,
13607                    None => (
13608                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13609                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13610                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13611                    ), // E4B: real v (K != V)
13612                };
13613                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
13614                // the normed rows; V ones-rms, never roped).
13615                e.rms_norm_qkv_rope(
13616                    &q0,
13617                    &k0,
13618                    &v0,
13619                    fa.q_norm.float_data(),
13620                    fa.k_norm.float_data(),
13621                    ones,
13622                    &mut q,
13623                    &mut k,
13624                    &mut v,
13625                    hd,
13626                    self.gemma4_rope_dims(il),
13627                    nh * t,
13628                    nkv * t,
13629                    pos_d,
13630                    nh,
13631                    nkv,
13632                    base,
13633                    1.0,
13634                    ff,
13635                    eps,
13636                )?;
13637            }
13638            let kvl = cache.kv[il].as_mut().unwrap();
13639            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
13640            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
13641            // degenerate tok-0 stream, 2026-07-12).
13642            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13643            if dc_bucket.is_some() {
13644                // DC arm (graph serving): append at the len_d slot, advance the counter
13645                // in-stream — replay-correct, no host len in the launch args. Host mirrors
13646                // are NOT touched here (the replay loop owns them; a bump at capture-record
13647                // time would double-count the capture iteration).
13648                debug_assert!(t == 1);
13649                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
13650                e.append_kv_quantized_row_dc_inc(
13651                    &k,
13652                    &v,
13653                    &mut kvl.k,
13654                    &mut kvl.v,
13655                    &mut kvl.len_d,
13656                    kvl.kv_dim_k,
13657                    kvl.kv_dim_v,
13658                    kvl.k_tok_bytes,
13659                    kvl.v_tok_bytes,
13660                    cls,
13661                )?;
13662            } else {
13663                e.append_kv_quantized_rows(
13664                    &k,
13665                    &v,
13666                    &mut kvl.k,
13667                    &mut kvl.v,
13668                    kvl.len,
13669                    t,
13670                    kvl.kv_dim_k,
13671                    kvl.kv_dim_v,
13672                    kvl.k_tok_bytes,
13673                    kvl.v_tok_bytes,
13674                    cls,
13675                )?;
13676                kvl.len += t;
13677            }
13678            kv_f32 = Some((k, v));
13679        }
13680        // attention: per-row causal fa over the (own or target) quantized cache. The cache
13681        // already contains this forward's rows in both arms; row i attends [.., base+i].
13682        let kvl_idx = share.unwrap_or(il);
13683        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13684        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
13685        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13686        let mut attn = e.uninit(t * nh * hd)?;
13687        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
13688        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
13689        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
13690        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
13691        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
13692        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
13693        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
13694        //     rows (the T=K verify kernel; the target appended this forward's rows already).
13695        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
13696        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
13697        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
13698        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13699            if let Some((kf, vf)) = &kv_f32 {
13700                if hd == 256 && t <= win {
13701                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13702                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13703                }
13704                if hd == 256 && swa && t > win {
13705                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13706                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13707                }
13708                if hd == 512 && !swa {
13709                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13710                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13711                }
13712            } else if share.is_some() {
13713                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13714                let k_view = e.view_u8(&kvl.k, kvl.k.len());
13715                let v_view = e.view_u8(&kvl.v, kvl.v.len());
13716                if hd == 256 && (!swa || t <= win) {
13717                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
13718                    e.fa_prefill_view(
13719                        &q,
13720                        &k_view,
13721                        &v_view,
13722                        &mut attn,
13723                        hd,
13724                        nh,
13725                        nkv,
13726                        t,
13727                        t,
13728                        scale,
13729                        true,
13730                        kvl.k_tok_bytes,
13731                        kvl.v_tok_bytes,
13732                        g,
13733                    )?;
13734                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13735                }
13736                // remaining shared classes (swa above the window; hd512 globals): dequant
13737                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
13738                let kv_dim = nkv * hd;
13739                let mut kf = e.uninit(t * kv_dim)?;
13740                let mut vf = e.uninit(t * kv_dim)?;
13741                e.fa_dequant_kv_view_f32(
13742                    &k_view,
13743                    &v_view,
13744                    &mut kf,
13745                    &mut vf,
13746                    kv_dim,
13747                    kv_dim,
13748                    t,
13749                    kvl.k_tok_bytes,
13750                    kvl.v_tok_bytes,
13751                    g,
13752                )?;
13753                if hd == 512 {
13754                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13755                } else {
13756                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13757                }
13758                return Ok(e.matmul(&fa.wo, &attn, t)?);
13759            }
13760        }
13761        if let Some(bucket) = dc_bucket {
13762            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
13763            // fa_decode_dc over the live counter. len_d already advanced past this token
13764            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
13765            // counter (advanced when the target ran earlier in the stack).
13766            assert!(t == 1);
13767            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
13768            // and under the window every live t_kv sits below it — cap the capture bucket
13769            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
13770            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
13771            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
13772            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13773                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13774            } else {
13775                bucket
13776            };
13777            let k_view = e.view_u8(&kvl.k, kvl.k.len());
13778            let v_view = e.view_u8(&kvl.v, kvl.v.len());
13779            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13780            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
13781            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
13782            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
13783            // captured into the dc graph like any other launch. Extending the cascade to
13784            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
13785            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
13786            // MEMRA_WPF=0 rollback seam.
13787            if crate::Engine::wpf_level() >= 1 {
13788                e.prefetch_weight_l2(&fa.wo)?;
13789            }
13790            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
13791            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
13792            if e.uses_q8_1_fast(&fa.wo) {
13793                let mut oq = e.alloc_i8_uninit(nh * hd)?;
13794                let mut od = e.zeros(nh * hd / 32)?;
13795                e.fa_decode_dc_q8(
13796                    &q,
13797                    &k_view,
13798                    &v_view,
13799                    &mut attn,
13800                    hd,
13801                    nh,
13802                    nkv,
13803                    &kvl.len_d,
13804                    bucket,
13805                    scale,
13806                    kvl.k_tok_bytes,
13807                    kvl.v_tok_bytes,
13808                    g,
13809                    Some((&mut oq, &mut od)),
13810                )?;
13811                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13812            }
13813            e.fa_decode_dc(
13814                &q,
13815                &k_view,
13816                &v_view,
13817                &mut attn,
13818                hd,
13819                nh,
13820                nkv,
13821                &kvl.len_d,
13822                bucket,
13823                scale,
13824                kvl.k_tok_bytes,
13825                kvl.v_tok_bytes,
13826                g,
13827            )?;
13828            return Ok(e.matmul(&fa.wo, &attn, t)?);
13829        }
13830        for i in 0..t {
13831            let avail = base_len + i + 1;
13832            let (off_tok, t_kv) = if swa && avail > win {
13833                (avail - win, win)
13834            } else {
13835                (0, avail)
13836            };
13837            let k_view = e.view_u8_range(
13838                &kvl.k,
13839                off_tok * kvl.k_tok_bytes,
13840                (off_tok + t_kv) * kvl.k_tok_bytes,
13841            );
13842            let v_view = e.view_u8_range(
13843                &kvl.v,
13844                off_tok * kvl.v_tok_bytes,
13845                (off_tok + t_kv) * kvl.v_tok_bytes,
13846            );
13847            let qv = e.view(&q, t * nh * hd);
13848            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13849            let mut q_one = e.uninit(nh * hd)?;
13850            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13851            let mut a_one = e.uninit(nh * hd)?;
13852            // read class MUST match the append class (globals are e4m3 under gkv): the
13853            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
13854            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
13855            e.fa_decode_kvmod(
13856                &q_one,
13857                &k_view,
13858                &v_view,
13859                &mut a_one,
13860                hd,
13861                nh,
13862                nkv,
13863                t_kv,
13864                scale,
13865                kvl.k_tok_bytes,
13866                kvl.v_tok_bytes,
13867                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13868            )?;
13869            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13870        }
13871        Ok(e.matmul(&fa.wo, &attn, t)?)
13872    }
13873
13874    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
13875    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
13876    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
13877    /// layer; does NOT advance cache.pos (caller owns pos).
13878    fn gemma4_e4b_trunk(
13879        &self,
13880        e: &Engine,
13881        tokens: &[u32],
13882        pos0: usize,
13883        cache: &mut Cache,
13884        head_last: bool,
13885    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13886        let n_embd = self.cfg.n_embd as usize;
13887        let t = tokens.len();
13888        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13889        let pos_d = e.htod_i32(&pos)?;
13890        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13891        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13892        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13893        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13894    }
13895
13896    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
13897    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
13898    /// eager chain by construction: SAME functions, not twins).
13899    fn gemma4_e4b_trunk_core(
13900        &self,
13901        e: &Engine,
13902        x_in: CudaSlice<f32>,
13903        inp_pl: CudaSlice<f32>,
13904        pos_d: &CudaSlice<i32>,
13905        t: usize,
13906        cache: &mut Cache,
13907        dc_bucket: Option<usize>,
13908        cap_logits: bool,
13909        head_last: bool,
13910    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13911        let n_embd = self.cfg.n_embd as usize;
13912        let eps = self.cfg.rms_eps;
13913        let n_layer = self.layers.len();
13914        let mut x = x_in;
13915        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13916        let n_epl = aux_e4b.n_epl;
13917
13918        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
13919        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
13920        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
13921        // head rides matmul_pre too. First layer's pair comes from a standalone fused
13922        // norm+quant.
13923        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13924        for il in 0..n_layer {
13925            let layer = &self.layers[il];
13926            let (hq, hdq) = match h_carry.take() {
13927                Some(p) => p,
13928                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13929            };
13930            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13931            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
13932            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
13933            let bits = layer.gemma4.as_ref().unwrap();
13934            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13935            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
13936            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
13937            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
13938            // the fused single-phase reduction is NOT FP-order-identical to the unfused
13939            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
13940            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
13941            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
13942            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
13943            // gate dropped, decode AND verify ride the same fused chain — parity by
13944            // construction, VERIFY-GATE 0.000e0.
13945            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13946            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13947                e,
13948                layer,
13949                &o,
13950                &x,
13951                t,
13952                Some(layer.post_attn_norm.float_data()),
13953                fuse_exit,
13954            )?;
13955            let mut resid = e.uninit(t * n_embd)?;
13956            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
13957            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
13958            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
13959            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
13960            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
13961            let g = if fuse_exit {
13962                // sn here = RAW f0 (post_ffw deferred).
13963                let (rq, rd) = e.rms_pre_add_q8_1(
13964                    &sn,
13965                    bits.post_ffw_norm.float_data(),
13966                    &attn_out,
13967                    &mut resid,
13968                    n_embd,
13969                    t,
13970                    self.cfg.rms_eps,
13971                )?;
13972                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13973            } else {
13974                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13975                e.matmul(&e4b.inp_gate, &resid, t)?
13976            };
13977            let mut act = e.uninit(t * n_epl)?;
13978            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13979                let ipv = e.view(&inp_pl, n_epl * n_layer);
13980                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13981                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13982                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13983            } else {
13984                let mut inp_this = e.uninit(t * n_epl)?;
13985                e.copy_rows_strided(
13986                    &inp_pl,
13987                    &mut inp_this,
13988                    n_epl,
13989                    t,
13990                    n_epl * n_layer,
13991                    il * n_epl,
13992                )?;
13993                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13994                e.matmul(&e4b.proj, &act, t)?
13995            };
13996            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
13997            // ONE launch (glue-fusion lane; last layer emits through output_norm).
13998            let next_norm = if il + 1 < n_layer {
13999                self.layers[il + 1].attn_norm.float_data()
14000            } else {
14001                self.output_norm.float_data()
14002            };
14003            let mut xn = e.uninit(t * n_embd)?;
14004            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
14005                &y,
14006                e4b.post_norm.float_data(),
14007                &resid,
14008                bits.layer_scale,
14009                next_norm,
14010                &mut xn,
14011                n_embd,
14012                t,
14013                eps,
14014            )?;
14015            h_carry = Some(pair);
14016            x = xn;
14017        }
14018        // the head consumes the last layer's fused (output_norm) emit. head_last callers
14019        // (prime, last_only forward) need only the final row's logits — the all-T head is
14020        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
14021        let (oq, odq) = h_carry.take().unwrap();
14022        let h0 = e.zeros(0)?;
14023        let hm = if head_last { 1 } else { t };
14024        let (hq, hd) = if head_last && t > 1 {
14025            let mut q1 = e.uninit_i8(n_embd)?;
14026            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
14027            let nb = n_embd / 32;
14028            let mut d1 = e.uninit(nb)?;
14029            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
14030            (q1, d1)
14031        } else {
14032            (oq, odq)
14033        };
14034        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
14035        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
14036        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
14037        // Logit-returning callers (host logits / spec prime) keep the capped emit.
14038        if cap_logits {
14039            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
14040            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
14041        }
14042        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
14043        Ok((ld, x))
14044    }
14045
14046    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
14047    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
14048    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
14049    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
14050    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
14051    /// covers exactly the layers that appended).
14052    pub fn gemma4_e4b_decode_step_t_am_dev(
14053        &self,
14054        e: &Engine,
14055        tok_d: &CudaSlice<u32>,
14056        t: usize,
14057        pos0: usize,
14058        cache: &mut Cache,
14059    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14060        let n_embd = self.cfg.n_embd as usize;
14061        let eps = self.cfg.rms_eps;
14062        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
14063        let pos_d = e.htod_i32(&pos)?;
14064        let embd_gpu = self
14065            .embd_gpu
14066            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
14067        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
14068        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
14069        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
14070        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
14071        let (ld, xp) =
14072            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
14073        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
14074        // emit is already capped, matching the eager chain bit-for-bit).
14075        let n_vocab = self.output.out_features();
14076        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
14077        for i in 0..t {
14078            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
14079        }
14080        let mut hn = e.uninit(t * n_embd)?;
14081        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
14082        cache.pos += t;
14083        Ok((vam, hn))
14084    }
14085
14086    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
14087    /// prime path — mirror of `gemma4_decode_step_t_h`).
14088    pub(crate) fn gemma4_e4b_decode_step_t_h(
14089        &self,
14090        e: &Engine,
14091        tokens: &[u32],
14092        pos0: usize,
14093        cache: &mut Cache,
14094    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14095        let n_embd = self.cfg.n_embd as usize;
14096        let eps = self.cfg.rms_eps;
14097        let t = tokens.len();
14098        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
14099        let mut hn = e.uninit(t * n_embd)?;
14100        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
14101        cache.pos += t;
14102        Ok((e.dtoh(&ld)?, hn))
14103    }
14104
14105    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
14106    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
14107    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
14108    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
14109    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
14110    pub fn gemma4_e4b_decode_step_dcg(
14111        &self,
14112        e: &Engine,
14113        token_d: &mut CudaSlice<u32>,
14114        pos_d: &mut CudaSlice<i32>,
14115        embd_gpu: &CudaSlice<u8>,
14116        embd_qt: i32,
14117        embd_rb: usize,
14118        cache: &mut Cache,
14119        n_vocab: usize,
14120        bucket: usize,
14121    ) -> Result<(), Box<dyn std::error::Error>> {
14122        let n_embd = self.cfg.n_embd as usize;
14123        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14124        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14125        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14126        let (ld, _x) =
14127            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
14128        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
14129        e.inc_seqlen(pos_d)?;
14130        Ok(())
14131    }
14132
14133    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
14134    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
14135    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
14136    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
14137    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
14138    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
14139    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
14140    #[allow(clippy::too_many_arguments)]
14141    pub fn gemma4_e4b_decode_step_dc(
14142        &self,
14143        e: &Engine,
14144        token_d: &CudaSlice<u32>,
14145        pos_d: &mut CudaSlice<i32>,
14146        embd_gpu: &CudaSlice<u8>,
14147        embd_qt: i32,
14148        embd_rb: usize,
14149        cache: &mut Cache,
14150        n_vocab: usize,
14151    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
14152        let n_embd = self.cfg.n_embd as usize;
14153        let eps = self.cfg.rms_eps;
14154        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
14155        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
14156        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
14157        let (ld, _x) =
14158            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
14159        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
14160        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
14161        e.inc_seqlen(pos_d)?;
14162        cache.pos += 1;
14163        let _ = eps;
14164        Ok(tok_out)
14165    }
14166
14167    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
14168    /// pre-output_norm hidden). Advances cache.pos.
14169    pub(crate) fn gemma4_e4b_decode_step_h(
14170        &self,
14171        e: &Engine,
14172        token: u32,
14173        cache: &mut Cache,
14174    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14175        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
14176        let logits = e.dtoh(&ld)?;
14177        cache.pos += 1;
14178        Ok((logits, x))
14179    }
14180
14181    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
14182    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
14183    /// fast; the prefill fa arms come later.
14184    pub(crate) fn gemma4_e4b_prime(
14185        &self,
14186        e: &Engine,
14187        tokens: &[u32],
14188        cache: &mut Cache,
14189    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14190        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
14191        // process-kill as gemma4_prime — refuse per-request.
14192        if cache.pos != 0 {
14193            return Err(
14194                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
14195                        call or decode tokenwise"
14196                    .into(),
14197            );
14198        }
14199        let n_embd = self.cfg.n_embd as usize;
14200        let t = tokens.len();
14201        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
14202        cache.pos += t;
14203        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
14204        let xv = e.view(&x, t * n_embd);
14205        let row = xv.slice((t - 1) * n_embd..t * n_embd);
14206        let mut h_seed = e.uninit(n_embd)?;
14207        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
14208        Ok((last, h_seed, x))
14209    }
14210
14211    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
14212    pub(crate) fn gemma4_e4b_forward(
14213        &self,
14214        e: &Engine,
14215        tokens: &[u32],
14216        last_only: bool,
14217    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
14218        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
14219        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
14220        Ok(e.dtoh(&ld)?) // head_last already reduced to the final row when last_only
14221    }
14222}
14223
14224#[cfg(test)]
14225mod prime_chunk_schedule_tests {
14226    use super::{
14227        PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
14228        fixed_prime_chunk_ranges_for_ring,
14229    };
14230
14231    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
14232        ranges.iter().map(|(start, end)| end - start).collect()
14233    }
14234
14235    fn auto_chunk(t: usize) -> usize {
14236        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
14237    }
14238
14239    #[test]
14240    fn fixed_schedule_retains_measured_geometry() {
14241        assert_eq!(
14242            sizes(&fixed_prime_chunk_ranges(461, 128)),
14243            vec![128, 128, 128, 77]
14244        );
14245        assert_eq!(
14246            sizes(&fixed_prime_chunk_ranges(1833, 230)),
14247            vec![230, 230, 230, 230, 230, 230, 230, 223]
14248        );
14249        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
14250        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
14251        assert_eq!(capped, vec![4096, 4088, 16]);
14252        assert!(capped.iter().all(|&rows| rows <= 4096));
14253        assert_eq!(
14254            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
14255            vec![4100],
14256            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
14257        );
14258    }
14259
14260    #[test]
14261    fn dynamic_schedule_matches_registered_shapes() {
14262        let cases = [
14263            (461, vec![64, 141, 132, 124]),
14264            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
14265            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
14266        ];
14267        for (t, expected) in cases {
14268            let chunk = auto_chunk(t);
14269            let fixed = fixed_prime_chunk_ranges(t, chunk);
14270            assert_eq!(
14271                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
14272                expected
14273            );
14274        }
14275    }
14276
14277    #[test]
14278    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
14279        for t in 256..=8192 {
14280            let chunk = auto_chunk(t);
14281            let fixed = fixed_prime_chunk_ranges(t, chunk);
14282            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
14283            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
14284            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
14285            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
14286            for pair in dynamic.windows(2) {
14287                assert_eq!(pair[0].1, pair[1].0, "T={t}");
14288            }
14289            assert!(
14290                dynamic
14291                    .iter()
14292                    .all(|(start, end)| end - start >= PRIME_MIN_T),
14293                "T={t} sizes={:?}",
14294                sizes(&dynamic)
14295            );
14296            if dynamic.len() >= 3 {
14297                let chunk_sizes = sizes(&dynamic);
14298                assert!(
14299                    chunk_sizes[0] < chunk_sizes[1],
14300                    "T={t} sizes={chunk_sizes:?}"
14301                );
14302                assert!(
14303                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
14304                    "T={t} sizes={chunk_sizes:?}"
14305                );
14306            }
14307        }
14308    }
14309}
14310
14311#[cfg(test)]
14312mod page_prefetch_tests {
14313    use super::{
14314        grouped_worker_prefetch_position, page_prefetch_positions,
14315        page_prefetch_window_from_values, worker_prefetch_positions,
14316    };
14317
14318    #[test]
14319    fn page_prefetch_window_keeps_existing_opt_in_default() {
14320        assert_eq!(page_prefetch_window_from_values(false, None), 0);
14321        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
14322        assert_eq!(page_prefetch_window_from_values(true, None), 1);
14323        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
14324        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
14325        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
14326    }
14327
14328    #[test]
14329    fn rolling_page_prefetch_advises_each_future_expert_once() {
14330        let advised: Vec<_> = (0..7)
14331            .flat_map(|position| page_prefetch_positions(position, 7, 3))
14332            .collect();
14333        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
14334
14335        let one_ahead: Vec<_> = (0..4)
14336            .flat_map(|position| page_prefetch_positions(position, 4, 1))
14337            .collect();
14338        assert_eq!(one_ahead, vec![1, 2, 3]);
14339        assert!(page_prefetch_positions(0, 4, 0).is_empty());
14340    }
14341
14342    #[test]
14343    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
14344        assert_eq!(grouped_worker_prefetch_position(0, None), None);
14345        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
14346            .chain(
14347                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
14348            )
14349            .collect();
14350        assert_eq!(positions, vec![0, 1, 2, 3]);
14351        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
14352    }
14353
14354    #[test]
14355    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
14356        let queued: Vec<_> = (0..8)
14357            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
14358            .collect();
14359        assert_eq!(queued, (0..8).collect::<Vec<_>>());
14360
14361        let one_at_a_time: Vec<_> = (0..4)
14362            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
14363            .collect();
14364        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
14365        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
14366    }
14367}
14368
14369pub struct G4DcSlots {
14370    x: CudaSlice<f32>,
14371    xn: CudaSlice<f32>,
14372    cur: CudaSlice<f32>,
14373    hq: CudaSlice<i8>,
14374    hd_: CudaSlice<f32>,
14375    q0: CudaSlice<f32>,
14376    k0: CudaSlice<f32>,
14377    v0: CudaSlice<f32>,
14378    q: CudaSlice<f32>,
14379    k: CudaSlice<f32>,
14380    v: CudaSlice<f32>,
14381    attn: CudaSlice<f32>,
14382    o: CudaSlice<f32>,
14383    attn_out: CudaSlice<f32>,
14384    zsh: CudaSlice<f32>,
14385    zq: CudaSlice<i8>,
14386    zd: CudaSlice<f32>,
14387    gate: CudaSlice<f32>,
14388    up: CudaSlice<f32>,
14389    act: CudaSlice<f32>,
14390    actq: CudaSlice<i8>,
14391    actd: CudaSlice<f32>,
14392    f0: CudaSlice<f32>,
14393    sn: CudaSlice<f32>,
14394    hn: CudaSlice<f32>,
14395    logits: CudaSlice<f32>,
14396}