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        _ => false,
374    }
375}
376
377/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
378/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
379fn moe_prewarm_enabled() -> bool {
380    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381    *E.get_or_init(|| {
382        std::env::var("MEMRA_MOE_PREWARM")
383            .map(|v| v != "0")
384            .unwrap_or(true)
385    })
386}
387
388/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
389/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
390/// can then vote for and exercise those experts on GPU before the cache is frozen.
391fn cpu_expert_profile_admit_enabled() -> bool {
392    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
393    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
394}
395
396/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
397/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
398/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
399pub const PRIME_MIN_T: usize = 16;
400const PRIME_PIPE_MICROBATCHES: usize = 8;
401const PRIME_PIPE_MIN_CHUNK: usize = 128;
402const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
403const PRIME_PIPE_LINEAR_WORK: usize = 8;
404
405fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
406    crate::pp::prime_pp_on()
407        && !crate::pp::pp2_streams_off()
408        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
409}
410
411/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative.
412/// Naked PP-2 primes use the measured pipeline geometry: up to eight microchunks, never
413/// below 128 tokens, while the legacy 4096-token cap remains the long-context bound.
414pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
415    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
416        let parsed = value
417            .parse::<usize>()
418            .unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
419        return if crate::cache::swa_ring_on() {
420            if parsed == 0 {
421                crate::cache::PRIME_CHUNK_MAX_TOKENS
422            } else {
423                parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
424            }
425        } else {
426            parsed
427        };
428    }
429    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
430    if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
431        chunk.min(
432            t.div_ceil(PRIME_PIPE_MICROBATCHES)
433                .max(PRIME_PIPE_MIN_CHUNK),
434        )
435    } else {
436        chunk
437    }
438}
439
440fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
441    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
442}
443
444fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
445    if chunk == 0 || t <= chunk {
446        return vec![(0, t)];
447    }
448    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
449    let mut start = 0usize;
450    while start < t {
451        let mut end = (start + chunk).min(t);
452        if t - end > 0 && t - end < PRIME_MIN_T {
453            if ring_on {
454                let shifted = t - PRIME_MIN_T;
455                end = if shifted > start { shifted } else { t };
456            } else {
457                end = t;
458            }
459        }
460        ranges.push((start, end));
461        start = end;
462    }
463    ranges
464}
465
466fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
467    let prefix = prefix as u128;
468    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
469}
470
471fn dynamic_prime_chunk_ranges(
472    t: usize,
473    fixed_chunk: usize,
474    fixed: &[(usize, usize)],
475) -> Vec<(usize, usize)> {
476    let n = fixed.len();
477    if n < 3 {
478        return fixed.to_vec();
479    }
480
481    let max_first = t - (n - 1) * PRIME_MIN_T;
482    let first = fixed_chunk
483        .div_ceil(2)
484        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
485        .min(max_first);
486    let mut ranges = Vec::with_capacity(n);
487    ranges.push((0, first));
488
489    let first_work = prime_chunk_work(first, t);
490    let work_span = prime_chunk_work(t, t) - first_work;
491    let denominator = (n - 1) as u128;
492    let mut previous = first;
493    for boundary in 1..n - 1 {
494        let target = first_work * denominator + work_span * (boundary as u128);
495        let remaining = n - 1 - boundary;
496        let mut low = previous + PRIME_MIN_T;
497        let mut high = t - remaining * PRIME_MIN_T;
498        while low < high {
499            let mid = low + (high - low) / 2;
500            if prime_chunk_work(mid, t) * denominator >= target {
501                high = mid;
502            } else {
503                low = mid + 1;
504            }
505        }
506        ranges.push((previous, low));
507        previous = low;
508    }
509    ranges.push((previous, t));
510    ranges
511}
512
513/// Internal prime ranges. The naked PP-2 pipeline defaults to a short-fill,
514/// equal-modeled-time schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
515/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
516pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
517    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
518    let chunk = prime_chunk_tokens(t, n_layers);
519    let fixed = fixed_prime_chunk_ranges(t, chunk);
520    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
521        Ok(value) => value == "dynamic",
522        Err(_) => true,
523    };
524    if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
525        fixed
526    } else {
527        dynamic_prime_chunk_ranges(t, chunk, &fixed)
528    }
529}
530
531impl HybridModel {
532    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
533    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
534    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
535    /// (it forces a dtoh + host hash per layer).
536    fn prime_trace_path() -> Option<&'static str> {
537        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
538        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
539            .as_deref()
540    }
541
542    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
543    pub fn forward(
544        &self,
545        e: &Engine,
546        tokens: &[u32],
547    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
548        if self.is_gemma4_e4b() {
549            return self.gemma4_e4b_forward(e, tokens, false);
550        }
551        if self.cfg.gemma4.is_some() {
552            return self.gemma4_forward(e, tokens, false);
553        }
554        let cfg = &self.cfg;
555        let n_embd = cfg.n_embd as usize;
556        let t = tokens.len();
557        let eps = cfg.rms_eps;
558        let pos: Vec<i32> = (0..t as i32).collect();
559        let pos_d = e.htod_i32(&pos)?;
560
561        let mut x = self.embed(e, tokens)?; // [T, n_embd]
562
563        for (il, layer) in self.layers.iter().enumerate() {
564            // attn_norm
565            let mut h = e.uninit(t * n_embd)?;
566            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
567
568            let mixed = match &layer.mixer {
569                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
570                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
571                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
572            };
573
574            // residual 1
575            let mut x1 = e.uninit(t * n_embd)?;
576            e.add(&x, &mixed, &mut x1, t * n_embd)?;
577
578            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
579            let mut z = e.uninit(t * n_embd)?;
580            e.rms_norm(
581                &x1,
582                layer.post_attn_norm.float_data(),
583                &mut z,
584                n_embd,
585                t,
586                eps,
587            )?;
588            let ffn_out = match &layer.ffn {
589                crate::hybrid::Ffn::Dense {
590                    ffn_gate,
591                    ffn_up,
592                    ffn_down,
593                } => {
594                    let n_ff = ffn_gate.out_features();
595                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
596                    let up = g2.pop().unwrap();
597                    let gate = g2.pop().unwrap();
598                    let mut act = e.uninit(t * n_ff)?;
599                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
600                    // both the dense MLP and the shared expert, and its limit is
601                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
602                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
603                    Self::ffn_act_lim(
604                        e,
605                        &self.cfg,
606                        &gate,
607                        &up,
608                        1.0,
609                        1.0,
610                        self.cfg.clamp_shexp_at(il as u32),
611                        &mut act,
612                        t * n_ff,
613                    )?;
614                    e.matmul(ffn_down, &act, t)?
615                }
616                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
617            };
618            let mut x2 = e.uninit(t * n_embd)?;
619            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
620            x = x2;
621        }
622
623        let mut hn = e.uninit(t * n_embd)?;
624        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
625        let logits = e.matmul(&self.output, &hn, t)?;
626        Ok(e.dtoh(&logits)?)
627    }
628
629    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
630    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
631    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
632    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
633    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
634    pub fn forward_last(
635        &self,
636        e: &Engine,
637        tokens: &[u32],
638    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
639        if self.cfg.gemma4.is_some() {
640            return self.gemma4_forward(e, tokens, true);
641        }
642        let cfg = &self.cfg;
643        let n_embd = cfg.n_embd as usize;
644        let t = tokens.len();
645        let eps = cfg.rms_eps;
646        let pos: Vec<i32> = (0..t as i32).collect();
647        let pos_d = e.htod_i32(&pos)?;
648
649        let mut x = self.embed(e, tokens)?; // [T, n_embd]
650        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
651        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
652        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
653        for (il, layer) in self.layers.iter().enumerate() {
654            let mut h = e.uninit(t * n_embd)?;
655            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
656            if probe {
657                e.stream().synchronize()?;
658                eprintln!("[probe] L{il} norm ok");
659            }
660            let mixed = match &layer.mixer {
661                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
662                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
663                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
664            };
665            if probe {
666                e.stream().synchronize()?;
667                eprintln!("[probe] L{il} mixer ok");
668            }
669            let mut x1 = e.uninit(t * n_embd)?;
670            e.add(&x, &mixed, &mut x1, t * n_embd)?;
671            let mut z = e.uninit(t * n_embd)?;
672            e.rms_norm(
673                &x1,
674                layer.post_attn_norm.float_data(),
675                &mut z,
676                n_embd,
677                t,
678                eps,
679            )?;
680            let ffn_out = match &layer.ffn {
681                crate::hybrid::Ffn::Dense {
682                    ffn_gate,
683                    ffn_up,
684                    ffn_down,
685                } => {
686                    let n_ff = ffn_gate.out_features();
687                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
688                    let up = g2.pop().unwrap();
689                    let gate = g2.pop().unwrap();
690                    let mut act = e.uninit(t * n_ff)?;
691                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
692                    Self::ffn_act_lim(
693                        e,
694                        &self.cfg,
695                        &gate,
696                        &up,
697                        1.0,
698                        1.0,
699                        self.cfg.clamp_shexp_at(il as u32),
700                        &mut act,
701                        t * n_ff,
702                    )?;
703                    e.matmul(ffn_down, &act, t)?
704                }
705                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
706            };
707            if probe {
708                e.stream().synchronize()?;
709                eprintln!("[probe] L{il} ffn ok");
710            }
711            let mut x2 = e.uninit(t * n_embd)?;
712            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
713            x = x2;
714        }
715        // norm over all T, then slice the LAST row and run lm_head on that single row.
716        let mut hn = e.uninit(t * n_embd)?;
717        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
718        let last = e.view(&hn, t * n_embd); // [T, n_embd]
719        let last_row = last.slice((t - 1) * n_embd..t * n_embd); // [1, n_embd]
720        let mut hlast = e.uninit(n_embd)?;
721        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
722        let logits = e.matmul(&self.output, &hlast, 1)?; // [1, n_vocab] — lm_head on ONE row
723        Ok(e.dtoh(&logits)?)
724    }
725
726    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
727    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
728    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
729    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
730    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
731    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
732    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
733    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
734    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
735    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
736    ///       argmax gate is the accuracy authority, exactly as for forward_last);
737    ///   (c) `cache.pos`/KV len/len_d advance by T.
738    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
739    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
740    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
741    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
742    ///
743    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
744    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
745    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
746    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
747    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
748    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
749    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
750    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
751    /// differently under load — research/tick-seg-20260807, receipt in
752    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
753    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
754    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
755    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
756    /// caller that SPLITS one request across calls passes the remainder.
757    pub fn prime_cache(
758        &self,
759        e: &Engine,
760        tokens: &[u32],
761        cache: &mut Cache,
762        queued_after: usize,
763    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
764        let n_embd = self.cfg.n_embd as usize;
765        let t = tokens.len();
766        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
767        // session cache — every chunk (including the first) takes the continuation arm
768        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
769        assert!(
770            t >= PRIME_MIN_T,
771            "prime_cache needs T >= {PRIME_MIN_T} (caller gates)"
772        );
773        assert!(
774            cache.pos + t <= cache.max_ctx,
775            "prime_cache: prompt exceeds cache max_ctx"
776        );
777
778        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
779        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
780        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
781        // each chunk runs the full layer stack with transients sized to the chunk, appending its
782        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
783        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
784        // exactly the state carry it was built for). Full-attn chunks after the first attend to
785        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
786        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
787        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
788        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
789        if self.is_gemma4_e4b() {
790            return self.gemma4_e4b_prime(e, tokens, cache);
791        }
792        if self.cfg.gemma4.is_some() {
793            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
794            return self.gemma4_prime(e, tokens, cache);
795        }
796        let ranges = prime_chunk_ranges(t, self.layers.len());
797        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
798        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
799        // the prefill's ARITHMETIC, so two rigs with different values produced different
800        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
801        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
802        // (VERDICT.md) — and it is NOT what docs originally said:
803        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
804        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
805        //     output head), so growing a chunk cannot move an existing row's value.
806        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
807        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
808        //     not describe our leak.
809        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
810        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
811        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
812        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
813        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
814        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
815        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
816        // the source — every row is in one numeric class, so the chunk size no longer steers
817        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
818        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
819        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
820        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
821        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
822        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
823        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
824        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
825        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
826        // across calls, the request still ends at the same absolute position, whatever the tick
827        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
828        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
829        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
830        // default. Read per call, not cached (the probe flips it in-process between arms). Never
831        // on in a measured default run.
832        let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
833        let seq_end = if legacy_calllocal {
834            cache.pos + t
835        } else {
836            cache.pos + t + queued_after
837        };
838        if ranges.len() == 1 {
839            return self.prime_chunk(e, tokens, cache, seq_end);
840        }
841        // PIPELINED PP-2 PRIME (lane/cx-pipeline-prime, 2026-08-08): overlap stage 0 of
842        // chunk N+1 with stage 1 of chunk N. The serial split stays reachable through
843        // MEMRA_PRIME_PIPE=0 and is the exactness oracle. N>2 keeps the serial walker;
844        // this lane owns the balanced two-stage schedule only.
845        if crate::pp::prime_pipe_on() && crate::pp::prime_pp_on() && !crate::pp::pp2_streams_off() {
846            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
847                if crate::pp::pp_multi_stream_same_device() {
848                    return Err(
849                        "prime chunk pipeline refused with 2 stage streams on one device — \
850                         that concurrent-stream placement remains quarantined by the deferred \
851                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
852                         the serial split."
853                            .into(),
854                    );
855                }
856                return self.prime_cache_pp2_pipelined(e, tokens, cache, seq_end, &ranges, &fence);
857            }
858        }
859        let mut hiddens = e.uninit(t * n_embd)?;
860        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
861        for &(start, end) in &ranges {
862            let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
863            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
864            last = Some((l, hs));
865        }
866        let (logits, h_seed) = last.unwrap();
867        Ok((logits, h_seed, hiddens))
868    }
869
870    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
871    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
872    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
873    /// norm, lm head, and caller hidden-stack copy as the serial split.
874    fn prime_cache_pp2_pipelined(
875        &self,
876        e: &Engine,
877        tokens: &[u32],
878        cache: &mut Cache,
879        seq_end: usize,
880        ranges: &[(usize, usize)],
881        fence: &[usize],
882    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
883        debug_assert_eq!(fence.len(), 3);
884        debug_assert!(ranges.len() >= 2);
885        let rt = crate::pp::PpNRt::get(e)?;
886        assert_eq!(
887            rt.n_stages(),
888            2,
889            "prime pipeline requires exactly two PP stages"
890        );
891        let n_embd = self.cfg.n_embd as usize;
892        let t = tokens.len();
893        let initial_base = cache.pos;
894        let caller_stream = e.stream();
895
896        // #87 reverse publication before any new stage allocation, then prewarm both
897        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
898        // after stage 1(N) is queued would synchronize that stream and erase the first
899        // overlap on a two-chunk prompt.
900        rt.fence_stages_behind(&caller_stream)?;
901        let max_payload = ranges.iter().map(|(s, e)| (e - s) * n_embd).max().unwrap();
902        rt.prepare_overlap_slots(0, max_payload)?;
903
904        let mut hiddens = e.uninit(t * n_embd)?;
905        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
906        let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
907        let (cache0, cache1) = stage_caches.parts();
908        let (first_start, first_end) = ranges[0];
909        let mut slot = self.prime_pp2_stage0_enqueue(
910            e,
911            rt,
912            &tokens[first_start..first_end],
913            cache0,
914            seq_end,
915            fence,
916            initial_base + first_start,
917            true,
918        )?;
919        cache0.pos = initial_base + first_end;
920
921        for (i, &(start, end)) in ranges.iter().enumerate() {
922            let base = initial_base + start;
923            debug_assert_eq!(
924                cache1.pos, base,
925                "stage 1 must drain chunks in original position order"
926            );
927            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
928                let next_base = initial_base + next_start;
929                debug_assert_eq!(
930                    cache0.pos, next_base,
931                    "stage 0 must issue chunks in original position order"
932                );
933                let cache0_stage = &mut *cache0;
934                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
935                // on one host thread therefore serialize even if the calls are ordered as
936                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
937                // stage 1 consumes slot N while stage 0 produces slot N+1.
938                std::thread::scope(|scope| -> Result<_, Box<dyn std::error::Error>> {
939                    let stage0 = scope.spawn(move || -> Result<usize, String> {
940                        let next = self
941                            .prime_pp2_stage0_enqueue(
942                                e,
943                                rt,
944                                &tokens[next_start..next_end],
945                                cache0_stage,
946                                seq_end,
947                                fence,
948                                next_base,
949                                true,
950                            )
951                            .map_err(|err| err.to_string())?;
952                        cache0_stage.pos = initial_base + next_end;
953                        Ok(next)
954                    });
955                    let x = self.prime_pp2_stage1_enqueue(
956                        e,
957                        rt,
958                        slot,
959                        end - start,
960                        cache1,
961                        seq_end,
962                        fence,
963                        base,
964                        true,
965                    )?;
966                    let out = {
967                        rt.bind_stage(1)?;
968                        let _st1 = rt.enter(1);
969                        let e1 = rt.engine(1, e);
970                        self.prime_chunk_epilogue(e1, x, end - start, cache1)?
971                    };
972                    let next = stage0
973                        .join()
974                        .map_err(|_| "pipeprime stage-0 host walker panicked")?
975                        .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
976                    Ok((out, Some(next)))
977                })?
978            } else {
979                let x = self.prime_pp2_stage1_enqueue(
980                    e,
981                    rt,
982                    slot,
983                    end - start,
984                    cache1,
985                    seq_end,
986                    fence,
987                    base,
988                    true,
989                )?;
990                let out = {
991                    rt.bind_stage(1)?;
992                    let _st1 = rt.enter(1);
993                    let e1 = rt.engine(1, e);
994                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
995                };
996                (out, None)
997            };
998
999            rt.publish_to(1, &caller_stream)?;
1000            e.copy_into(&mut hiddens, start * n_embd, &out.2, (end - start) * n_embd)?;
1001            last = Some((out.0, out.1));
1002            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1003
1004            if let Some(next) = next_slot {
1005                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
1006                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
1007                // Stage 0(N+1) is already queued before this wait is appended, so its
1008                // overlap with stage 1(N) is preserved.
1009                rt.fence_stages_behind(&caller_stream)?;
1010                slot = next;
1011            }
1012        }
1013
1014        debug_assert_eq!(cache0.pos, initial_base + t);
1015        debug_assert_eq!(cache1.pos, initial_base + t);
1016        let (logits, h_seed) = last.unwrap();
1017        Ok((logits, h_seed, hiddens))
1018    }
1019
1020    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
1021    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
1022    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
1023    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
1024    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
1025    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
1026    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
1027        if Engine::gdn_db_on()
1028            && Engine::gdn_chunked_enabled()
1029            && t >= 16
1030            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
1031            && num_k * 2 == num_v
1032        {
1033            num_k
1034        } else {
1035            num_v
1036        }
1037    }
1038
1039    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
1040    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
1041    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
1042    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
1043    fn f16out_on(e: &Engine, t: usize) -> bool {
1044        crate::f16_ffi::pp_f16_enabled()
1045            && t >= 16
1046            && !e.verify_exact_on()
1047            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
1048    }
1049
1050    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
1051    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
1052    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
1053    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
1054    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
1055    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
1056    /// see one entry, byte-identical behavior.
1057    pub fn prime_slabs_get(
1058        &self,
1059        e: &Engine,
1060        t: usize,
1061        n_embd: usize,
1062        n_ff_max: usize,
1063    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
1064        let mut slabs = self.prime_slabs.lock().unwrap();
1065        let dev = e.ctx().ordinal();
1066        let need_new = match slabs.get(&dev) {
1067            None => true,
1068            Some(sl) => sl.lock().unwrap().t_cap < t,
1069        };
1070        if need_new {
1071            slabs.insert(
1072                dev,
1073                std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
1074                    t_cap: t,
1075                    h: e.uninit(t * n_embd)?,
1076                    x1: e.uninit(t * n_embd)?,
1077                    z: e.uninit(t * n_embd)?,
1078                    act: e.uninit(t * n_ff_max)?,
1079                    xa: e.uninit(t * n_embd)?,
1080                    xb: e.uninit(t * n_embd)?,
1081                    h16: e.alloc_u8_uninit(t * n_embd * 2)?,
1082                    z16: e.alloc_u8_uninit(t * n_embd * 2)?,
1083                    gate: e.uninit(t * n_ff_max)?,
1084                    up: e.uninit(t * n_ff_max)?,
1085                    ffn_out: e.uninit(t * n_embd)?,
1086                    seg_glue: Vec::new(),
1087                    mixed: e.uninit(t * n_embd)?,
1088                    seg_mid: Vec::new(),
1089                    seg_t: 0,
1090                })),
1091            );
1092        }
1093        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
1094    }
1095
1096    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
1097    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
1098    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
1099    fn prime_chunk(
1100        &self,
1101        e: &Engine,
1102        tokens: &[u32],
1103        cache: &mut Cache,
1104        seq_end: usize,
1105    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1106        if crate::pp::pp_host_bounce_active()
1107            && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1108        {
1109            return Err(
1110                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1111                 has no active prime stage split and would peer-read remote weights; keep \
1112                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1113                    .into(),
1114            );
1115        }
1116        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
1117        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
1118        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
1119        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
1120        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
1121        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
1122        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
1123        // loader is off and there is nothing remote to split for.
1124        if self.cfg.gemma4.is_none() && !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1125            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1126                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1127            }
1128        }
1129        if crate::pp::pp_host_bounce_active() {
1130            return Err(
1131                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1132                 refusing an unsplit remote-weight walk"
1133                    .into(),
1134            );
1135        }
1136        let t = tokens.len();
1137        let base = cache.pos;
1138        debug_assert!(
1139            seq_end >= base + t,
1140            "prime_chunk: seq_end must cover this chunk"
1141        );
1142        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1143        let pos_d = e.htod_i32(&pos)?;
1144
1145        let x_embed = self.embed(e, tokens)?; // [T, n_embd]
1146        let x = self.prime_layers(
1147            e,
1148            x_embed,
1149            0,
1150            self.layers.len(),
1151            &pos_d,
1152            t,
1153            base,
1154            cache,
1155            seq_end,
1156        )?;
1157        self.prime_chunk_epilogue(e, x, t, cache)
1158    }
1159
1160    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
1161    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
1162    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
1163    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
1164    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
1165    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
1166    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
1167    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
1168    ///     the plain add (materialize) and the next stage hoists its own first norm — the
1169    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
1170    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
1171    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
1172    ///     each stage walks through its own resident transients;
1173    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
1174    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
1175    #[allow(clippy::too_many_arguments)]
1176    fn prime_layers(
1177        &self,
1178        e: &Engine,
1179        x_in: CudaSlice<f32>,
1180        lo: usize,
1181        hi: usize,
1182        pos_d: &CudaSlice<i32>,
1183        t: usize,
1184        base: usize,
1185        cache: &mut Cache,
1186        seq_end: usize,
1187    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1188        let cfg = &self.cfg;
1189        let n_embd = cfg.n_embd as usize;
1190        let eps = cfg.rms_eps;
1191        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
1192        // standalone convert launches). Only when the f16 lane serves and T reaches the
1193        // GEMM tier; bit-identical either way.
1194        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1195        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
1196        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
1197        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
1198        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
1199        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
1200        let n_ff_max = self
1201            .layers
1202            .iter()
1203            .map(|l| match &l.ffn {
1204                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1205                _ => n_embd,
1206            })
1207            .max()
1208            .unwrap_or(n_embd)
1209            .max(n_embd);
1210        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1211        let slab = if use_slabs {
1212            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1213        } else {
1214            None
1215        };
1216        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1217        let mut x_own; // fallback storage when slabs are off
1218        type SlabRefs<'a> = (
1219            &'a mut CudaSlice<f32>,
1220            &'a mut CudaSlice<f32>,
1221            &'a mut CudaSlice<f32>,
1222            &'a mut CudaSlice<f32>,
1223            &'a mut CudaSlice<u8>,
1224            &'a mut CudaSlice<u8>,
1225            &'a mut CudaSlice<f32>,
1226            &'a mut CudaSlice<f32>,
1227            &'a mut CudaSlice<f32>,
1228        );
1229        let (mut x_cur, mut x_nxt, sl): (
1230            &mut CudaSlice<f32>,
1231            &mut CudaSlice<f32>,
1232            Option<SlabRefs>,
1233        );
1234        let mut seg: Option<(
1235            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1236            &mut Vec<Option<cudarc::driver::CudaGraph>>,
1237            &mut CudaSlice<f32>,
1238            &mut usize,
1239        )> = None;
1240        let mut x_own2;
1241        match slab_guard.as_mut() {
1242            Some(g) => {
1243                let slabs = &mut **g;
1244                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1245                let PrimeSlabs {
1246                    xa,
1247                    xb,
1248                    h,
1249                    x1,
1250                    z,
1251                    act,
1252                    h16,
1253                    z16,
1254                    gate,
1255                    up,
1256                    ffn_out,
1257                    seg_glue,
1258                    mixed,
1259                    seg_mid,
1260                    seg_t,
1261                    ..
1262                } = slabs;
1263                x_cur = xa;
1264                x_nxt = xb;
1265                seg = Some((seg_glue, seg_mid, mixed, seg_t));
1266                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1267            }
1268            None => {
1269                x_own = x_in;
1270                x_own2 = e.uninit(t * n_embd)?;
1271                x_cur = &mut x_own;
1272                x_nxt = &mut x_own2;
1273                sl = None;
1274            }
1275        }
1276        let mut alloc_h;
1277        let mut alloc_x1;
1278        let mut alloc_z;
1279        let mut alloc_act;
1280        let mut alloc_h16;
1281        let mut alloc_z16;
1282        let mut alloc_gate;
1283        let mut alloc_up;
1284        let mut alloc_fo;
1285        let (h, x1, z, act): (
1286            &mut CudaSlice<f32>,
1287            &mut CudaSlice<f32>,
1288            &mut CudaSlice<f32>,
1289            &mut CudaSlice<f32>,
1290        );
1291        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1292        let (sl_gate, sl_up, sl_fo): (
1293            &mut CudaSlice<f32>,
1294            &mut CudaSlice<f32>,
1295            &mut CudaSlice<f32>,
1296        );
1297        match sl {
1298            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1299                h = a;
1300                x1 = b;
1301                z = c;
1302                act = d;
1303                h16 = e16;
1304                z16 = f16b;
1305                sl_gate = g;
1306                sl_up = u;
1307                sl_fo = fo;
1308            }
1309            None => {
1310                alloc_h = e.uninit(t * n_embd)?;
1311                alloc_x1 = e.uninit(t * n_embd)?;
1312                alloc_z = e.uninit(t * n_embd)?;
1313                alloc_act = e.uninit(t * n_ff_max)?;
1314                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1315                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1316                alloc_gate = e.uninit(t * n_ff_max)?;
1317                alloc_up = e.uninit(t * n_ff_max)?;
1318                alloc_fo = e.uninit(t * n_embd)?;
1319                h = &mut alloc_h;
1320                x1 = &mut alloc_x1;
1321                z = &mut alloc_z;
1322                act = &mut alloc_act;
1323                h16 = &mut alloc_h16;
1324                z16 = &mut alloc_z16;
1325                sl_gate = &mut alloc_gate;
1326                sl_up = &mut alloc_up;
1327                sl_fo = &mut alloc_fo;
1328            }
1329        }
1330        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
1331        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
1332        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
1333        // first prime at this t (capture does not execute -> launch right after).
1334        let n_layers = self.layers.len();
1335        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
1336        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
1337        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
1338        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
1339        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
1340        // machinery stays (byte-identical) as their foundation.
1341        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
1342        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
1343        // step35 rides its own mixer through the normal per-layer arm below.
1344        let use_seg = f16fuse
1345            && seg.is_some()
1346            && self.cfg.step35.is_none()
1347            && lo == 0
1348            && hi == n_layers
1349            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1350        if let Some((sg, sm, _, st)) = seg.as_mut() {
1351            if **st != t {
1352                sg.clear();
1353                sg.extend((0..n_layers).map(|_| None));
1354                sm.clear();
1355                sm.extend((0..n_layers).map(|_| None));
1356                **st = t;
1357            }
1358        }
1359        {
1360            let layer_lo = &self.layers[lo];
1361            if f16fuse {
1362                e.rms_norm_f16out(
1363                    x_cur,
1364                    layer_lo.attn_norm.float_data(),
1365                    h,
1366                    h16,
1367                    n_embd,
1368                    t,
1369                    eps,
1370                )?;
1371            } else {
1372                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1373            }
1374        }
1375        for il in lo..hi {
1376            let layer = &self.layers[il];
1377            let hx16 = if f16fuse { Some(&*h16) } else { None };
1378            if use_seg {
1379                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
1380                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
1381                let (pre, pre16, w_out) = match &layer.mixer {
1382                    Mixer::Full(fa) => {
1383                        let g3 = match hx16 {
1384                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1385                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1386                        };
1387                        let (pre, pre16) =
1388                            self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1389                        (pre, pre16, &fa.wo)
1390                    }
1391                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1392                    Mixer::Linear(la) => {
1393                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1394                        let g4 = match hx16 {
1395                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1396                            None => e.matmul_group(&ws, h, t)?,
1397                        };
1398                        let (pre, pre16) =
1399                            self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1400                        (pre, pre16, &la.ssm_out)
1401                    }
1402                };
1403                {
1404                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
1405                    let pre_n = pre.len() / t;
1406                    let xh_pre = match pre16 {
1407                        Some(x) => x,
1408                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
1409                    };
1410                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1411                        let y = e.matmul(w_out, &pre, t)?;
1412                        e.copy_into(mslab, 0, &y, t * n_embd)?;
1413                    }
1414                    if sm[il].is_none() {
1415                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1416                        let w_post = layer.post_attn_norm.float_data();
1417                        e.stream().synchronize()?;
1418                        e.stream()
1419                            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1420                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1421                            e.add(x_cur, mslab, x1, t * n_embd)?;
1422                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1423                            Ok(())
1424                        })();
1425                        let g = e.stream().end_capture(
1426                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1427                        r?;
1428                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1429                    }
1430                    sm[il].as_ref().unwrap().launch()?;
1431                }
1432            } else {
1433                let mixed = match &layer.mixer {
1434                    Mixer::Full(fa) => {
1435                        self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il, seq_end)?
1436                    }
1437                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1438                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1439                };
1440                if f16fuse {
1441                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
1442                    // bit-identical) — the standalone add pass disappears.
1443                    e.add_rms_norm_f16out(
1444                        x_cur,
1445                        &mixed,
1446                        layer.post_attn_norm.float_data(),
1447                        x1,
1448                        z,
1449                        z16,
1450                        n_embd,
1451                        t,
1452                        eps,
1453                    )?;
1454                } else {
1455                    e.add(x_cur, &mixed, x1, t * n_embd)?;
1456                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1457                }
1458            }
1459            let zx16 = if f16fuse { Some(&*z16) } else { None };
1460            match &layer.ffn {
1461                crate::hybrid::Ffn::Dense {
1462                    ffn_gate,
1463                    ffn_up,
1464                    ffn_down,
1465                } => {
1466                    let n_ff = ffn_gate.out_features();
1467                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
1468                    // the allocating group + copy when a mirror is missing.
1469                    let mut into_ok = false;
1470                    if let Some(xh) = zx16 {
1471                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1472                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1473                    }
1474                    if !into_ok {
1475                        let mut g2 = match zx16 {
1476                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1477                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1478                        };
1479                        let up_y = g2.pop().unwrap();
1480                        let gate_y = g2.pop().unwrap();
1481                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1482                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1483                    }
1484                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
1485                    // operand in-epilogue; non-silu activations keep the standalone convert.
1486                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
1487                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
1488                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1489                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() && d_lim.is_none()
1490                    {
1491                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1492                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1493                        Some(a16)
1494                    } else {
1495                        Self::ffn_act_lim(
1496                            e,
1497                            &self.cfg,
1498                            sl_gate,
1499                            sl_up,
1500                            1.0,
1501                            1.0,
1502                            d_lim,
1503                            act,
1504                            t * n_ff,
1505                        )?;
1506                        None
1507                    };
1508                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
1509                    let xh_act = match act16 {
1510                        Some(x) => x,
1511                        None => e.f16_act(act, t * n_ff, n_ff)?,
1512                    };
1513                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1514                        let y = e.matmul(ffn_down, &*act, t)?;
1515                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1516                    }
1517                }
1518                crate::hybrid::Ffn::Moe(m) => {
1519                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1520                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1521                }
1522            }
1523            if use_seg && il + 1 < hi {
1524                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
1525                let w_next = self.layers[il + 1].attn_norm.float_data();
1526                let (sg, _, _, _) = seg.as_mut().unwrap();
1527                if sg[il].is_none() {
1528                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1529                    e.stream().synchronize()?;
1530                    e.stream()
1531                        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1532                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1533                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1534                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1535                        Ok(())
1536                    })();
1537                    let g = e.stream().end_capture(
1538                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1539                    );
1540                    r?;
1541                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1542                }
1543                sg[il].as_ref().unwrap().launch()?;
1544            } else {
1545                if il + 1 < hi {
1546                    let w_next = self.layers[il + 1].attn_norm.float_data();
1547                    if f16fuse {
1548                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1549                    } else {
1550                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1551                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1552                    }
1553                } else {
1554                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1555                }
1556            }
1557            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
1558            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
1559            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
1560            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
1561            // unset (the default) costs one OnceLock read per layer.
1562            if let Some(path) = Self::prime_trace_path() {
1563                let row = (base + t - 1) as usize;
1564                let host = e.dtoh(x_nxt)?;
1565                let last = &host[(t - 1) * n_embd..t * n_embd];
1566                use std::io::Write as _;
1567                let mut f = std::fs::OpenOptions::new()
1568                    .create(true)
1569                    .append(true)
1570                    .open(path)?;
1571                let mut h64: u64 = 0xcbf29ce484222325;
1572                for v in last {
1573                    h64 ^= v.to_bits() as u64;
1574                    h64 = h64.wrapping_mul(0x100000001b3);
1575                }
1576                writeln!(
1577                    f,
1578                    "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1579                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1580                    last[0], last[1], last[2]
1581                )?;
1582            }
1583            std::mem::swap(&mut x_cur, &mut x_nxt);
1584        }
1585        // hidden-stack return: clone the final x out of the slab
1586        let mut x = e.uninit(t * n_embd)?;
1587        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1588        drop(slab_guard);
1589        Ok(x)
1590    }
1591
1592    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
1593    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
1594    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
1595    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
1596    fn prime_chunk_epilogue(
1597        &self,
1598        e: &Engine,
1599        x: CudaSlice<f32>,
1600        t: usize,
1601        cache: &mut Cache,
1602    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1603        let n_embd = self.cfg.n_embd as usize;
1604        let eps = self.cfg.rms_eps;
1605        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
1606        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
1607        // the post-norm copy happens after hn exists).
1608        let mut h_seed = e.uninit(n_embd)?;
1609        if !crate::spec::spec_hpost() {
1610            e.copy_view_into(
1611                &mut h_seed,
1612                0,
1613                &x.slice((t - 1) * n_embd..t * n_embd),
1614                n_embd,
1615            )?;
1616        }
1617        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
1618        let mut hn = e.uninit(t * n_embd)?;
1619        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1620        if crate::spec::spec_hpost() {
1621            e.copy_view_into(
1622                &mut h_seed,
1623                0,
1624                &hn.slice((t - 1) * n_embd..t * n_embd),
1625                n_embd,
1626            )?;
1627        }
1628        let last = e.view(&hn, t * n_embd);
1629        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1630        let mut hlast = e.uninit(n_embd)?;
1631        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1632        let logits = e.matmul(&self.output, &hlast, 1)?;
1633        cache.pos += t;
1634        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
1635        // post-norm stack hn (MEMRA_SPEC_HPOST).
1636        Ok((
1637            e.dtoh(&logits)?,
1638            h_seed,
1639            if crate::spec::spec_hpost() { hn } else { x },
1640        ))
1641    }
1642
1643    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
1644    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
1645    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
1646    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
1647    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
1648    /// prefill kernels. Structure mirrors the verify split exactly:
1649    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
1650    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
1651    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
1652    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
1653    ///                  there via the sharded loader) → `publish_to`
1654    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
1655    /// round's stage-freed buffers must not be reused under the caller's queued reads);
1656    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
1657    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
1658    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
1659    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
1660    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
1661    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
1662    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
1663    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
1664    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
1665    /// and its liveness counter is bumped here — the gate goes green with this function.
1666    fn prime_chunk_ppn(
1667        &self,
1668        e: &Engine,
1669        tokens: &[u32],
1670        cache: &mut Cache,
1671        seq_end: usize,
1672        fence: &[usize],
1673    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1674        let rt = crate::pp::PpNRt::get(e)?;
1675        let n_st = fence.len() - 1;
1676        assert_eq!(
1677            rt.n_stages(),
1678            n_st,
1679            "PpNRt stage count {} != fence stages {n_st}",
1680            rt.n_stages()
1681        );
1682        let n_embd = self.cfg.n_embd as usize;
1683        let t = tokens.len();
1684        let base = cache.pos;
1685        debug_assert!(
1686            seq_end >= base + t,
1687            "prime_chunk_ppn: seq_end must cover this chunk"
1688        );
1689        let payload = t * n_embd;
1690        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1691        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1692        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1693        let caller_stream = e.stream();
1694        rt.fence_stages_behind(&caller_stream)?;
1695
1696        if n_st == 2 {
1697            let slot =
1698                self.prime_pp2_stage0_enqueue(e, rt, tokens, cache, seq_end, fence, base, false)?;
1699            let x =
1700                self.prime_pp2_stage1_enqueue(e, rt, slot, t, cache, seq_end, fence, base, false)?;
1701            let out = {
1702                rt.bind_stage(1)?;
1703                let _st1 = rt.enter(1);
1704                let e1 = rt.engine(1, e);
1705                self.prime_chunk_epilogue(e1, x, t, cache)?
1706            };
1707            rt.publish_to(1, &caller_stream)?;
1708            crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1709            return Ok(out);
1710        }
1711
1712        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1713
1714        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1715        let mut slot = {
1716            let _st0 = rt.enter(0);
1717            let e0 = rt.engine(0, e);
1718            let pos_d = e0.htod_i32(&pos)?;
1719            let x = self.embed(e0, tokens)?;
1720            let x =
1721                self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1722            rt.tx(0, &x, payload)?
1723            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1724        };
1725
1726        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1727        for s in 1..n_st - 1 {
1728            let _st = rt.enter(s);
1729            let es = rt.engine(s, e);
1730            let pos_d = es.htod_i32(&pos)?;
1731            let x = rt.rx(s - 1, slot, payload)?;
1732            let x = self.prime_layers(
1733                es,
1734                x,
1735                fence[s],
1736                fence[s + 1],
1737                &pos_d,
1738                t,
1739                base,
1740                cache,
1741                seq_end,
1742            )?;
1743            slot = rt.tx(s, &x, payload)?;
1744        }
1745
1746        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1747        let _stl = rt.enter(n_st - 1);
1748        let el = rt.engine(n_st - 1, e);
1749        let pos_d = el.htod_i32(&pos)?;
1750        let x = rt.rx(n_st - 2, slot, payload)?;
1751        let x = self.prime_layers(
1752            el,
1753            x,
1754            fence[n_st - 1],
1755            fence[n_st],
1756            &pos_d,
1757            t,
1758            base,
1759            cache,
1760            seq_end,
1761        )?;
1762        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1763        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1764        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1765        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1766        // stage stream host-side, but the law is stated in events, not in a dtoh side
1767        // effect a later deferred form would remove.
1768        rt.publish_to(n_st - 1, &caller_stream)?;
1769        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1770        Ok(out)
1771    }
1772
1773    fn prime_pp2_stage0_enqueue(
1774        &self,
1775        e: &Engine,
1776        rt: &crate::pp::PpNRt,
1777        tokens: &[u32],
1778        cache: &mut Cache,
1779        seq_end: usize,
1780        fence: &[usize],
1781        base: usize,
1782        pipelined: bool,
1783    ) -> Result<usize, Box<dyn std::error::Error>> {
1784        let t = tokens.len();
1785        let n_embd = self.cfg.n_embd as usize;
1786        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1787        rt.bind_stage(0)?;
1788        let _st0 = rt.enter(0);
1789        let e0 = rt.engine(0, e);
1790        let pos_d = e0.htod_i32(&pos)?;
1791        let x = self.embed(e0, tokens)?;
1792        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1793        let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end)?;
1794        if pipelined {
1795            rt.tx_pipelined(0, &x, t * n_embd)
1796        } else {
1797            rt.tx(0, &x, t * n_embd)
1798        }
1799    }
1800
1801    fn prime_pp2_stage1_enqueue(
1802        &self,
1803        e: &Engine,
1804        rt: &crate::pp::PpNRt,
1805        slot: usize,
1806        t: usize,
1807        cache: &mut Cache,
1808        seq_end: usize,
1809        fence: &[usize],
1810        base: usize,
1811        pipelined: bool,
1812    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1813        let n_embd = self.cfg.n_embd as usize;
1814        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1815        rt.bind_stage(1)?;
1816        let _st1 = rt.enter(1);
1817        let e1 = rt.engine(1, e);
1818        let pos_d = e1.htod_i32(&pos)?;
1819        let x = rt.rx(0, slot, t * n_embd)?;
1820        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1821        self.prime_layers(e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end)
1822    }
1823
1824    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
1825    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
1826    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
1827    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
1828    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
1829    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
1830    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
1831    /// bookkeeping still runs on the host per call — the real replay path moves the write
1832    /// slot to the len_d device counter (increment 3).
1833    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
1834    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
1835    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
1836    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
1837    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
1838    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
1839    pub fn prime_chunk_captured(
1840        &self,
1841        e: &Engine,
1842        x_in: &CudaSlice<f32>,
1843        pos_d: &CudaSlice<i32>,
1844        t: usize,
1845        cache: &mut Cache,
1846        len_d: &CudaSlice<i32>,
1847        logits_out: &mut CudaSlice<f32>,
1848        h_seed_out: &mut CudaSlice<f32>,
1849    ) -> Result<(), Box<dyn std::error::Error>> {
1850        let cfg = &self.cfg;
1851        let n_embd = cfg.n_embd as usize;
1852        let eps = cfg.rms_eps;
1853        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1854        let mut x = e.uninit(t * n_embd)?;
1855        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1856        for (il, layer) in self.layers.iter().enumerate() {
1857            let mut h = e.uninit(t * n_embd)?;
1858            let mut hx16: Option<CudaSlice<u8>> = None;
1859            if f16fuse {
1860                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1861                e.rms_norm_f16out(
1862                    &x,
1863                    layer.attn_norm.float_data(),
1864                    &mut h,
1865                    &mut b16,
1866                    n_embd,
1867                    t,
1868                    eps,
1869                )?;
1870                hx16 = Some(b16);
1871            } else {
1872                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1873            }
1874            let mixed = match &layer.mixer {
1875                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
1876                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
1877                // come from the caller (see step35_attn_pre_wo's doc note).
1878                Mixer::Full(fa) => {
1879                    self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il, t)?
1880                }
1881                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1882                Mixer::Linear(la) => {
1883                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1884                    let g4 = match hx16.as_ref() {
1885                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1886                        None => e.matmul_group(&ws, &h, t)?,
1887                    };
1888                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1889                }
1890            };
1891            let mut x1 = e.uninit(t * n_embd)?;
1892            e.add(&x, &mixed, &mut x1, t * n_embd)?;
1893            let mut z = e.uninit(t * n_embd)?;
1894            let mut zx16: Option<CudaSlice<u8>> = None;
1895            if f16fuse {
1896                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1897                e.rms_norm_f16out(
1898                    &x1,
1899                    layer.post_attn_norm.float_data(),
1900                    &mut z,
1901                    &mut b16,
1902                    n_embd,
1903                    t,
1904                    eps,
1905                )?;
1906                zx16 = Some(b16);
1907            } else {
1908                e.rms_norm(
1909                    &x1,
1910                    layer.post_attn_norm.float_data(),
1911                    &mut z,
1912                    n_embd,
1913                    t,
1914                    eps,
1915                )?;
1916            }
1917            let ffn_out = match &layer.ffn {
1918                crate::hybrid::Ffn::Dense {
1919                    ffn_gate,
1920                    ffn_up,
1921                    ffn_down,
1922                } => {
1923                    let n_ff = ffn_gate.out_features();
1924                    let mut g2 = match &zx16 {
1925                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1926                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1927                    };
1928                    let up = g2.pop().unwrap();
1929                    let gate = g2.pop().unwrap();
1930                    let mut act = e.uninit(t * n_ff)?;
1931                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
1932                    Self::ffn_act_lim(
1933                        e,
1934                        &self.cfg,
1935                        &gate,
1936                        &up,
1937                        1.0,
1938                        1.0,
1939                        self.cfg.clamp_shexp_at(il as u32),
1940                        &mut act,
1941                        t * n_ff,
1942                    )?;
1943                    e.matmul(ffn_down, &act, t)?
1944                }
1945                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
1946            };
1947            let mut x2 = e.uninit(t * n_embd)?;
1948            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1949            x = x2;
1950        }
1951        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
1952        if !crate::spec::spec_hpost() {
1953            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1954        }
1955        let mut hn = e.uninit(t * n_embd)?;
1956        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1957        if crate::spec::spec_hpost() {
1958            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1959        }
1960        let mut hlast = e.uninit(n_embd)?;
1961        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1962        let logits = e.matmul(&self.output, &hlast, 1)?;
1963        let nv = logits.len();
1964        e.copy_into(logits_out, 0, &logits, nv)?;
1965        Ok(())
1966    }
1967
1968    fn step35_prime_batch_on() -> bool {
1969        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
1970    }
1971
1972    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
1973    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
1974    #[allow(clippy::too_many_arguments)]
1975    fn step35_prime_batch_layers(
1976        &self,
1977        e: &Engine,
1978        mut x: CudaSlice<f32>,
1979        lo: usize,
1980        hi: usize,
1981        ts: &[usize],
1982        offs: &[usize],
1983        pos_ds: &[CudaSlice<i32>],
1984        caches: &mut [&mut Cache],
1985    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1986        let cfg = &self.cfg;
1987        let n_embd = cfg.n_embd as usize;
1988        let eps = cfg.rms_eps;
1989        let b = ts.len();
1990        let total: usize = ts.iter().sum();
1991        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
1992
1993        let split = |e: &Engine,
1994                     y: &CudaSlice<f32>,
1995                     dim: usize|
1996         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1997            let mut out = Vec::with_capacity(b);
1998            for s in 0..b {
1999                let mut ys = e.uninit(ts[s] * dim)?;
2000                e.copy_view_into(
2001                    &mut ys,
2002                    0,
2003                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2004                    ts[s] * dim,
2005                )?;
2006                out.push(ys);
2007            }
2008            Ok(out)
2009        };
2010
2011        for il in lo..hi {
2012            let layer = &self.layers[il];
2013            let Mixer::Full(fa) = &layer.mixer else {
2014                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2015            };
2016
2017            let mut h = e.uninit(total * n_embd)?;
2018            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2019            if f16fuse {
2020                e.rms_norm_f16out(
2021                    &x,
2022                    layer.attn_norm.float_data(),
2023                    &mut h,
2024                    &mut hx16,
2025                    n_embd,
2026                    total,
2027                    eps,
2028                )?;
2029            } else {
2030                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, total, eps)?;
2031            }
2032
2033            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
2034            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
2035            // application stay verbatim.
2036            let gate_w = fa
2037                .attn_gate
2038                .as_ref()
2039                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2040            let mut g4 = if f16fuse {
2041                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
2042            } else {
2043                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
2044            };
2045            let gate = g4.pop().unwrap();
2046            let mut parts: Vec<Vec<CudaSlice<f32>>> =
2047                (0..b).map(|_| Vec::with_capacity(3)).collect();
2048            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
2049                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2050                    parts[s].push(ys);
2051                }
2052            }
2053            let gates = split(e, &gate, gate_w.out_features())?;
2054            let geometry = self.step35_geom(il);
2055            let hd = geometry.head_dim_k as usize;
2056            let nh = geometry.n_head as usize;
2057            let mut ag_cat = e.uninit(total * nh * hd)?;
2058            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
2059                let ag = self.step35_attn_pre_wo(
2060                    e,
2061                    fa,
2062                    g3s,
2063                    None,
2064                    Some(&gate),
2065                    &pos_ds[s],
2066                    ts[s],
2067                    Some(&mut *caches[s]),
2068                    il,
2069                    ts[s],
2070                )?;
2071                e.copy_into(&mut ag_cat, offs[s] * nh * hd, &ag, ts[s] * nh * hd)?;
2072            }
2073            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
2074
2075            let mut x1 = e.uninit(total * n_embd)?;
2076            let mut z = e.uninit(total * n_embd)?;
2077            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2078            if f16fuse {
2079                e.add_rms_norm_f16out(
2080                    &x,
2081                    &mixed,
2082                    layer.post_attn_norm.float_data(),
2083                    &mut x1,
2084                    &mut z,
2085                    &mut zx16,
2086                    n_embd,
2087                    total,
2088                    eps,
2089                )?;
2090            } else {
2091                e.add(&x, &mixed, &mut x1, total * n_embd)?;
2092                e.rms_norm(
2093                    &x1,
2094                    layer.post_attn_norm.float_data(),
2095                    &mut z,
2096                    n_embd,
2097                    total,
2098                    eps,
2099                )?;
2100            }
2101
2102            let ffn_out = match &layer.ffn {
2103                crate::hybrid::Ffn::Dense {
2104                    ffn_gate,
2105                    ffn_up,
2106                    ffn_down,
2107                } => {
2108                    let n_ff = ffn_gate.out_features();
2109                    let mut g2 = if f16fuse {
2110                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
2111                    } else {
2112                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
2113                    };
2114                    let up = g2.pop().unwrap();
2115                    let gate = g2.pop().unwrap();
2116                    let mut act = e.uninit(total * n_ff)?;
2117                    let d_lim = cfg.clamp_shexp_at(il as u32);
2118                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
2119                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2120                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2121                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2122                            Some(y) => y,
2123                            None => e.matmul(ffn_down, &act, total)?,
2124                        }
2125                    } else {
2126                        Self::ffn_act_lim(
2127                            e,
2128                            cfg,
2129                            &gate,
2130                            &up,
2131                            1.0,
2132                            1.0,
2133                            d_lim,
2134                            &mut act,
2135                            total * n_ff,
2136                        )?;
2137                        e.matmul(ffn_down, &act, total)?
2138                    }
2139                }
2140                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
2141            };
2142            let mut x2 = e.uninit(total * n_embd)?;
2143            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2144            x = x2;
2145        }
2146        Ok(x)
2147    }
2148
2149    fn step35_prime_batch_epilogue(
2150        &self,
2151        e: &Engine,
2152        x: CudaSlice<f32>,
2153        ts: &[usize],
2154        offs: &[usize],
2155        caches: &mut [&mut Cache],
2156    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2157        let n_embd = self.cfg.n_embd as usize;
2158        let total: usize = ts.iter().sum();
2159        let mut hn = e.uninit(total * n_embd)?;
2160        e.rms_norm(
2161            &x,
2162            self.output_norm.float_data(),
2163            &mut hn,
2164            n_embd,
2165            total,
2166            self.cfg.rms_eps,
2167        )?;
2168
2169        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
2170        let mut out = Vec::with_capacity(ts.len());
2171        for s in 0..ts.len() {
2172            let mut hidden = e.uninit(ts[s] * n_embd)?;
2173            e.copy_view_into(
2174                &mut hidden,
2175                0,
2176                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
2177                ts[s] * n_embd,
2178            )?;
2179            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2180            let mut h_seed = e.uninit(n_embd)?;
2181            e.copy_view_into(
2182                &mut h_seed,
2183                0,
2184                &hidden_src.slice(last0..last0 + n_embd),
2185                n_embd,
2186            )?;
2187            // Exactness-first: the serial reference runs the output head at m=1.
2188            let mut hlast = e.uninit(n_embd)?;
2189            e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2190            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
2191            caches[s].pos += ts[s];
2192            out.push((logits, h_seed, hidden));
2193        }
2194        Ok(out)
2195    }
2196
2197    fn step35_prime_cache_batch(
2198        &self,
2199        e: &Engine,
2200        prompts: &[&[u32]],
2201        caches: &mut [&mut Cache],
2202    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2203        if crate::pp::pp_host_bounce_active()
2204            && (!crate::pp::prime_pp_on() || crate::pp::pp_cuts(self.layers.len()).is_none())
2205        {
2206            return Err(
2207                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
2208                 stage split; refusing an unsplit remote-weight walk"
2209                    .into(),
2210            );
2211        }
2212        if !Self::step35_prime_batch_on() {
2213            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
2214        }
2215        if caches.iter().any(|c| c.pos != 0) {
2216            return Err(
2217                "step35 batched prime currently supports complete fresh prompts only; \
2218                 continuation/tick chunks require per-request queued_after"
2219                    .into(),
2220            );
2221        }
2222
2223        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2224        for &t in &ts {
2225            assert!(
2226                t >= PRIME_MIN_T,
2227                "step35 batched prime needs T >= {PRIME_MIN_T}"
2228            );
2229        }
2230        for (s, c) in caches.iter().enumerate() {
2231            assert!(
2232                ts[s] <= c.max_ctx,
2233                "step35 batched prime exceeds cache max_ctx"
2234            );
2235        }
2236        let offs: Vec<usize> = ts
2237            .iter()
2238            .scan(0usize, |a, &t| {
2239                let o = *a;
2240                *a += t;
2241                Some(o)
2242            })
2243            .collect();
2244        let total: usize = ts.iter().sum();
2245        let payload = total * self.cfg.n_embd as usize;
2246        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2247        let positions: Vec<Vec<i32>> = ts.iter().map(|&t| (0..t as i32).collect()).collect();
2248        let upload_positions =
2249            |e: &Engine| -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2250                positions
2251                    .iter()
2252                    .map(|p| e.htod_i32(p))
2253                    .collect::<Result<_, _>>()
2254            };
2255
2256        static ONCE: std::sync::Once = std::sync::Once::new();
2257        ONCE.call_once(|| {
2258            eprintln!(
2259                "[step35-prime-batch] first concat prime: B={} tokens={total}",
2260                prompts.len()
2261            );
2262        });
2263
2264        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
2265            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2266                let rt = crate::pp::PpNRt::get(e)?;
2267                let n_st = fence.len() - 1;
2268                assert_eq!(
2269                    rt.n_stages(),
2270                    n_st,
2271                    "step35 prime batch stage count mismatch"
2272                );
2273                let caller_stream = e.stream();
2274                rt.fence_stages_behind(&caller_stream)?;
2275
2276                let mut slot = {
2277                    let _st0 = rt.enter(0);
2278                    let e0 = rt.engine(0, e);
2279                    let pos_ds = upload_positions(e0)?;
2280                    let x = self.embed(e0, &cat_tokens)?;
2281                    let x = self.step35_prime_batch_layers(
2282                        e0, x, fence[0], fence[1], &ts, &offs, &pos_ds, caches,
2283                    )?;
2284                    rt.tx(0, &x, payload)?
2285                };
2286                for s in 1..n_st - 1 {
2287                    let _st = rt.enter(s);
2288                    let es = rt.engine(s, e);
2289                    let pos_ds = upload_positions(es)?;
2290                    let x = rt.rx(s - 1, slot, payload)?;
2291                    let x = self.step35_prime_batch_layers(
2292                        es,
2293                        x,
2294                        fence[s],
2295                        fence[s + 1],
2296                        &ts,
2297                        &offs,
2298                        &pos_ds,
2299                        caches,
2300                    )?;
2301                    slot = rt.tx(s, &x, payload)?;
2302                }
2303
2304                let _stl = rt.enter(n_st - 1);
2305                let el = rt.engine(n_st - 1, e);
2306                let pos_ds = upload_positions(el)?;
2307                let x = rt.rx(n_st - 2, slot, payload)?;
2308                let x = self.step35_prime_batch_layers(
2309                    el,
2310                    x,
2311                    fence[n_st - 1],
2312                    fence[n_st],
2313                    &ts,
2314                    &offs,
2315                    &pos_ds,
2316                    caches,
2317                )?;
2318                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2319                rt.publish_to(n_st - 1, &caller_stream)?;
2320                crate::pp::STEP35_PRIME_BATCH_SPLITS
2321                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2322                out
2323            } else {
2324                let pos_ds = upload_positions(e)?;
2325                let x = self.embed(e, &cat_tokens)?;
2326                let x = self.step35_prime_batch_layers(
2327                    e,
2328                    x,
2329                    0,
2330                    self.layers.len(),
2331                    &ts,
2332                    &offs,
2333                    &pos_ds,
2334                    caches,
2335                )?;
2336                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2337            }
2338        } else {
2339            let pos_ds = upload_positions(e)?;
2340            let x = self.embed(e, &cat_tokens)?;
2341            let x = self.step35_prime_batch_layers(
2342                e,
2343                x,
2344                0,
2345                self.layers.len(),
2346                &ts,
2347                &offs,
2348                &pos_ds,
2349                caches,
2350            )?;
2351            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2352        };
2353        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2354        Ok(out)
2355    }
2356
2357    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
2358    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
2359    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
2360    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
2361    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
2362    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
2363    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
2364    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
2365    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
2366    /// over the quantized past; Linear: the stateful pad_view twin — the same state
2367    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
2368    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
2369    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
2370    /// back to single-chunk serving).
2371    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
2372    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
2373    pub fn prime_cache_batch(
2374        &self,
2375        e: &Engine,
2376        prompts: &[&[u32]],
2377        caches: &mut [&mut Cache],
2378    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2379        let cfg = &self.cfg;
2380        let n_embd = cfg.n_embd as usize;
2381        let eps = cfg.rms_eps;
2382        let b = prompts.len();
2383        assert!(b >= 1 && b == caches.len());
2384        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2385        let carried = pos0s.iter().any(|&p| p > 0);
2386        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
2387        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
2388        // generic concat attn core below (uniform geometry, no per-layer swa window, no
2389        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
2390        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
2391        if cfg.gemma4.is_some() {
2392            return Err(
2393                "prime_cache_batch: gemma4 has no batched prime core (per-layer \
2394                        swa/global geometry, softcapped head) — use gemma4_prime per sequence"
2395                    .into(),
2396            );
2397        }
2398        // Step35 has a dedicated concat walk: the generic core below cannot express its
2399        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
2400        if cfg.step35.is_some() {
2401            return self.step35_prime_cache_batch(e, prompts, caches);
2402        }
2403        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2404        for &t in &ts {
2405            assert!(
2406                t >= PRIME_MIN_T,
2407                "prime_cache_batch needs T >= {PRIME_MIN_T}"
2408            );
2409        }
2410        for (s, c) in caches.iter().enumerate() {
2411            assert!(
2412                c.pos + ts[s] <= c.max_ctx,
2413                "prime_cache_batch: prompt exceeds cache max_ctx"
2414            );
2415        }
2416        let total: usize = ts.iter().sum();
2417        let offs: Vec<usize> = ts
2418            .iter()
2419            .scan(0usize, |a, &t| {
2420                let o = *a;
2421                *a += t;
2422                Some(o)
2423            })
2424            .collect();
2425        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
2426        let pos_ds: Vec<CudaSlice<i32>> = ts
2427            .iter()
2428            .zip(&pos0s)
2429            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2430            .collect::<Result<_, _>>()?;
2431        // split a concat [total, dim] buffer into per-seq copies
2432        let split = |e: &Engine,
2433                     y: &CudaSlice<f32>,
2434                     dim: usize|
2435         -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2436            let mut out = Vec::with_capacity(b);
2437            for s in 0..b {
2438                let mut ys = e.uninit(ts[s] * dim)?;
2439                e.copy_view_into(
2440                    &mut ys,
2441                    0,
2442                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
2443                    ts[s] * dim,
2444                )?;
2445                out.push(ys);
2446            }
2447            Ok(out)
2448        };
2449
2450        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2451        let mut x = self.embed(e, &cat_tokens)?; // [total, n_embd]
2452        for (il, layer) in self.layers.iter().enumerate() {
2453            let mut h = e.uninit(total * n_embd)?;
2454            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2455            e.rms_norm_f16out(
2456                &x,
2457                layer.attn_norm.float_data(),
2458                &mut h,
2459                &mut hx16,
2460                n_embd,
2461                total,
2462                eps,
2463            )?;
2464            // mixer: projection GROUP on the concat (m = total), stateful core per seq
2465            let mut mixed = e.uninit(total * n_embd)?;
2466            match &layer.mixer {
2467                Mixer::Full(fa) => {
2468                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2469                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
2470                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
2471                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
2472                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
2473                    // back to the per-seq dispatch.
2474                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
2475                    let (n_head, n_head_kv, head_dim) = (
2476                        geometry.n_head as usize,
2477                        geometry.n_head_kv as usize,
2478                        geometry.head_dim_k as usize,
2479                    );
2480                    let fa_scale = geometry.attention_scale();
2481                    let use_favl = !carried
2482                        && (2..=8).contains(&b)
2483                        && (head_dim == 256 || head_dim == 128)
2484                        && geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ
2485                        && std::env::var("MEMRA_NOFA").is_err()
2486                        && std::env::var("MEMRA_FA_FLOOR").is_err()
2487                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2488                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2489                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2490                    if use_favl {
2491                        let (qf_w, kf_w, vf_w) = (
2492                            fa.wq.out_features(),
2493                            fa.wk.out_features(),
2494                            fa.wv.out_features(),
2495                        );
2496                        struct APre {
2497                            q: CudaSlice<f32>,
2498                            gate: Option<CudaSlice<f32>>,
2499                            qn: CudaSlice<f32>,
2500                            kn: CudaSlice<f32>,
2501                        }
2502                        let mut aps = Vec::with_capacity(b);
2503                        for &t in ts.iter().take(b) {
2504                            aps.push(APre {
2505                                q: e.uninit(t * n_head * head_dim)?,
2506                                gate: Some(e.uninit(t * n_head * head_dim)?),
2507                                qn: e.uninit(t * n_head * head_dim)?,
2508                                kn: e.uninit(t * n_head_kv * head_dim)?,
2509                            });
2510                        }
2511                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2512                            let kvl = caches[0].kv[il].as_ref().unwrap();
2513                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2514                        };
2515                        let pargs: Vec<crate::AttnPreVl> = (0..b)
2516                            .map(|s| {
2517                                let (o, t) = (offs[s], ts[s]);
2518                                let kvl = caches[s].kv[il].as_ref().unwrap();
2519                                assert!(
2520                                    kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2521                                    "prime_cache_batch attn vl: fresh + capacity"
2522                                );
2523                                crate::AttnPreVl {
2524                                    qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2525                                    kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2526                                    vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2527                                    q: e.addr_f32(&aps[s].q),
2528                                    gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2529                                    qn: e.addr_f32(&aps[s].qn),
2530                                    kn: e.addr_f32(&aps[s].kn),
2531                                    kc: e.addr_u8(&kvl.k),
2532                                    vc: e.addr_u8(&kvl.v),
2533                                    t: t as i32,
2534                                    pad: 0,
2535                                }
2536                            })
2537                            .collect();
2538                        e.attn_pre_vl8(
2539                            &pargs,
2540                            fa.q_norm.float_data(),
2541                            fa.k_norm.float_data(),
2542                            head_dim,
2543                            geometry.n_rot as usize,
2544                            n_head,
2545                            n_head_kv,
2546                            self.cfg.rms_eps,
2547                            geometry.rope_base,
2548                            1.0,
2549                            kv_dim_k,
2550                            kv_dim_v,
2551                            ktb,
2552                            vtb,
2553                        )?;
2554                        for s in 0..b {
2555                            let kvl = caches[s].kv[il].as_mut().unwrap();
2556                            kvl.len += ts[s];
2557                            let new_len = kvl.len as i32;
2558                            e.set_i32_one(&mut kvl.len_d, new_len)?;
2559                        }
2560                        let mut attns = Vec::with_capacity(b);
2561                        let mut mirrors = Vec::with_capacity(b);
2562                        for &t in ts.iter().take(b) {
2563                            attns.push(e.uninit(t * n_head * head_dim)?);
2564                            let n = t * n_head_kv * head_dim;
2565                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2566                        }
2567                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
2568                        // promoted single-seq config is on; else the mma favl.
2569                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2570                            Ok("0") => false,
2571                            Ok("1") => true,
2572                            _ => cfg!(memra_hopper_mma),
2573                        };
2574                        if fa3_on {
2575                            let mut q16s = Vec::with_capacity(b);
2576                            let mut v16s = Vec::with_capacity(b);
2577                            for s in 0..b {
2578                                let t = ts[s];
2579                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2580                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2581                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2582                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2583                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2584                                e.f32_to_bf16_v(
2585                                    &g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2586                                    &mut v16,
2587                                    t * n_head_kv * head_dim,
2588                                )?;
2589                                q16s.push(q16);
2590                                v16s.push((k16, v16));
2591                            }
2592                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2593                            let mut kp = qp;
2594                            let mut vp = qp;
2595                            let mut op = [core::ptr::null_mut::<f32>(); 8];
2596                            let mut tsv = [0i32; 8];
2597                            for s in 0..b {
2598                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2599                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2600                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2601                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
2602                                tsv[s] = ts[s] as i32;
2603                            }
2604                            let rc = unsafe {
2605                                crate::fa3_vl_raw(
2606                                    qp.as_ptr(),
2607                                    kp.as_ptr(),
2608                                    vp.as_ptr(),
2609                                    op.as_ptr(),
2610                                    tsv.as_ptr(),
2611                                    b as i32,
2612                                    n_head as i32,
2613                                    n_head_kv as i32,
2614                                    head_dim as i32,
2615                                    fa_scale,
2616                                    e.stream().cu_stream() as *mut core::ffi::c_void,
2617                                )
2618                            };
2619                            if rc != 0 {
2620                                return Err(format!("memra_fa3_vl rc={rc}").into());
2621                            }
2622                        } else {
2623                            let fargs: Vec<crate::FaSeqVl> = (0..b)
2624                                .map(|s| crate::FaSeqVl {
2625                                    q: e.addr_f32(&aps[s].qn),
2626                                    k16: e.addr_u8(&mirrors[s].0),
2627                                    v16: e.addr_u8(&mirrors[s].1),
2628                                    o: e.addr_f32(&attns[s]),
2629                                    kf: e.addr_f32(&aps[s].kn),
2630                                    vf: e.addr_f32v(
2631                                        &g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w),
2632                                    ),
2633                                    t: ts[s] as i32,
2634                                    pad: 0,
2635                                })
2636                                .collect();
2637                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2638                        }
2639                        for (s, attn) in attns.into_iter().enumerate() {
2640                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
2641                                e,
2642                                attn,
2643                                &aps[s].gate,
2644                                ts[s],
2645                                n_head,
2646                                head_dim,
2647                            )?;
2648                            let mut done = false;
2649                            if let Some(xh) = &ag16 {
2650                                done = e.try_f16_gemm_pre_into_off(
2651                                    &fa.wo,
2652                                    xh,
2653                                    ts[s],
2654                                    &mut mixed,
2655                                    offs[s] * n_embd,
2656                                )?;
2657                            }
2658                            if !done {
2659                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2660                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2661                            }
2662                        }
2663                    } else {
2664                        let mut parts: Vec<Vec<CudaSlice<f32>>> =
2665                            (0..b).map(|_| Vec::new()).collect();
2666                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2667                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2668                                parts[s].push(ys);
2669                            }
2670                        }
2671                        for (s, g3s) in parts.into_iter().enumerate() {
2672                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
2673                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
2674                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il,
2675                            )?;
2676                            let mut done = false;
2677                            if let Some(xh) = &ag16 {
2678                                done = e.try_f16_gemm_pre_into_off(
2679                                    &fa.wo,
2680                                    xh,
2681                                    ts[s],
2682                                    &mut mixed,
2683                                    offs[s] * n_embd,
2684                                )?;
2685                            }
2686                            if !done {
2687                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2688                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2689                            }
2690                        }
2691                    }
2692                }
2693                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2694                Mixer::Linear(la) => {
2695                    // task #16: NO split copies (cores read row-offset views of the concat
2696                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
2697                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
2698                    // varlen K5 launch for all sequences.
2699                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2700                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2701                    let outs =
2702                        self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2703                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2704                        let (o, t) = (offs[s], ts[s]);
2705                        let mut done = false;
2706                        if let Some(xh) = &gn16 {
2707                            done = e.try_f16_gemm_pre_into_off(
2708                                &la.ssm_out,
2709                                xh,
2710                                t,
2711                                &mut mixed,
2712                                o * n_embd,
2713                            )?;
2714                        }
2715                        if !done {
2716                            let m = e.matmul(&la.ssm_out, &gn, t)?;
2717                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2718                        }
2719                    }
2720                }
2721            }
2722            let mut x1 = e.uninit(total * n_embd)?;
2723            let mut z = e.uninit(total * n_embd)?;
2724            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2725            e.add_rms_norm_f16out(
2726                &x,
2727                &mixed,
2728                layer.post_attn_norm.float_data(),
2729                &mut x1,
2730                &mut z,
2731                &mut zx16,
2732                n_embd,
2733                total,
2734                eps,
2735            )?;
2736            let ffn_out = match &layer.ffn {
2737                crate::hybrid::Ffn::Dense {
2738                    ffn_gate,
2739                    ffn_up,
2740                    ffn_down,
2741                } => {
2742                    let n_ff = ffn_gate.out_features();
2743                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2744                    let up = g2.pop().unwrap();
2745                    let gate = g2.pop().unwrap();
2746                    let mut act = e.uninit(total * n_ff)?;
2747                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
2748                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
2749                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
2750                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
2751                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2752                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2753                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2754                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2755                            Some(y) => y,
2756                            None => e.matmul(ffn_down, &act, total)?,
2757                        }
2758                    } else {
2759                        Self::ffn_act_lim(
2760                            e,
2761                            &self.cfg,
2762                            &gate,
2763                            &up,
2764                            1.0,
2765                            1.0,
2766                            d_lim,
2767                            &mut act,
2768                            total * n_ff,
2769                        )?;
2770                        e.matmul(ffn_down, &act, total)?
2771                    }
2772                }
2773                crate::hybrid::Ffn::Moe(m) => {
2774                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2775                }
2776            };
2777            let mut x2 = e.uninit(total * n_embd)?;
2778            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2779            x = x2;
2780        }
2781        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
2782        let mut hn = e.uninit(total * n_embd)?;
2783        e.rms_norm(
2784            &x,
2785            self.output_norm.float_data(),
2786            &mut hn,
2787            n_embd,
2788            total,
2789            eps,
2790        )?;
2791        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
2792        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
2793        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
2794        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
2795        // argmax battery arbitrates, same as every other prefill GEMM change.
2796        let mut hcat = e.uninit(b * n_embd)?;
2797        for s in 0..b {
2798            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2799            e.copy_view_into(
2800                &mut hcat,
2801                s * n_embd,
2802                &hn.slice(last0..last0 + n_embd),
2803                n_embd,
2804            )?;
2805        }
2806        let logits_cat = if b >= 2 {
2807            e.try_f16_gemm(&self.output, &hcat, b)?
2808        } else {
2809            None
2810        };
2811        let logits_host: Option<Vec<f32>> = match &logits_cat {
2812            Some(lc) => Some(e.dtoh(lc)?),
2813            None => None,
2814        };
2815        let n_vocab = self.output.out_features();
2816        let mut hidden_all = if crate::spec::spec_hpost() {
2817            split(e, &hn, n_embd)?
2818        } else {
2819            split(e, &x, n_embd)?
2820        };
2821        let mut out = Vec::with_capacity(b);
2822        for s in 0..b {
2823            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2824            let mut h_seed = e.uninit(n_embd)?;
2825            if !crate::spec::spec_hpost() {
2826                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2827            } else {
2828                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2829            }
2830            let logits = match &logits_host {
2831                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2832                None => {
2833                    let mut hlast = e.uninit(n_embd)?;
2834                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2835                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2836                }
2837            };
2838            caches[s].pos += ts[s];
2839            out.push((logits, h_seed, hidden_all.remove(0)));
2840        }
2841        Ok(out)
2842    }
2843
2844    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
2845    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
2846    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
2847    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
2848    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
2849    ///
2850    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
2851    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
2852    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
2853    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
2854    #[allow(clippy::too_many_arguments)]
2855    fn full_attn_prime(
2856        &self,
2857        e: &Engine,
2858        fa: &FullAttnLayer,
2859        h: &CudaSlice<f32>,
2860        hx: Option<&CudaSlice<u8>>,
2861        pos_d: &CudaSlice<i32>,
2862        t: usize,
2863        cache: &mut Cache,
2864        il: usize,
2865        seq_end: usize,
2866    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2867        if self.cfg.step35.is_some() {
2868            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2869        }
2870        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
2871        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
2872        // this single-seq path composes proj+core identically (byte-for-byte the old body).
2873        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
2874        let g3 = match hx {
2875            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2876            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2877        };
2878        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2879    }
2880
2881    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
2882    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
2883    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
2884    fn full_attn_prime_core(
2885        &self,
2886        e: &Engine,
2887        fa: &FullAttnLayer,
2888        g3: Vec<CudaSlice<f32>>,
2889        pos_d: &CudaSlice<i32>,
2890        t: usize,
2891        cache: &mut Cache,
2892        il: usize,
2893    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2894        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2895        if let Some(xh) = &ag16 {
2896            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2897                return Ok(y);
2898            }
2899        }
2900        Ok(e.matmul(&fa.wo, &attn_g, t)?)
2901    }
2902
2903    fn full_attn_prime_core_inner(
2904        &self,
2905        e: &Engine,
2906        fa: &FullAttnLayer,
2907        g3: Vec<CudaSlice<f32>>,
2908        pos_d: &CudaSlice<i32>,
2909        t: usize,
2910        cache: &mut Cache,
2911        il: usize,
2912    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2913        let cfg = &self.cfg;
2914        let geometry = cfg.full_attention_geometry_at(il as u32);
2915        let n_head = geometry.n_head as usize;
2916        let n_head_kv = geometry.n_head_kv as usize;
2917        let head_dim = geometry.head_dim_k as usize;
2918        let scale = geometry.attention_scale();
2919        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2920        let AttnPre { q, k, v, gate } = pre;
2921        let mut attn = e.uninit(t * n_head * head_dim)?;
2922        self.full_attn_prime_fa_dispatch(
2923            e, &q, &k, &v, &mut attn, base_len, t, cache, il, head_dim, n_head, n_head_kv, scale,
2924        )?;
2925        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2926    }
2927
2928    /// task #18 (attn side): projections tail through KV append — everything before the
2929    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
2930    /// present BEFORE this chunk's append (base_len; 0 == fresh).
2931    #[allow(clippy::type_complexity)]
2932    fn full_attn_prime_pre_fa(
2933        &self,
2934        e: &Engine,
2935        fa: &FullAttnLayer,
2936        mut g3: Vec<CudaSlice<f32>>,
2937        pos_d: &CudaSlice<i32>,
2938        t: usize,
2939        cache: &mut Cache,
2940        il: usize,
2941    ) -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
2942        let cfg = &self.cfg;
2943        let geometry = cfg.full_attention_geometry_at(il as u32);
2944        let n_head = geometry.n_head as usize;
2945        let n_head_kv = geometry.n_head_kv as usize;
2946        let head_dim = geometry.head_dim_k as usize;
2947        let eps = cfg.rms_eps;
2948
2949        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
2950        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
2951        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
2952        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2953        let v = g3.pop().unwrap();
2954        let mut k = g3.pop().unwrap();
2955        let qf = g3.pop().unwrap();
2956        let (mut q, gate) = if gated {
2957            let mut q = e.uninit(t * n_head * head_dim)?;
2958            let mut gate = e.uninit(t * n_head * head_dim)?;
2959            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2960            (q, Some(gate))
2961        } else {
2962            (qf, None)
2963        };
2964
2965        let mut qn = e.uninit(t * n_head * head_dim)?;
2966        e.rms_norm(
2967            &q,
2968            fa.q_norm.float_data(),
2969            &mut qn,
2970            head_dim,
2971            n_head * t,
2972            eps,
2973        )?;
2974        q = qn;
2975        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2976        e.rms_norm(
2977            &k,
2978            fa.k_norm.float_data(),
2979            &mut kn,
2980            head_dim,
2981            n_head_kv * t,
2982            eps,
2983        )?;
2984        k = kn;
2985        let rope_dims = geometry.n_rot as usize;
2986        e.rope_neox(
2987            &mut q,
2988            pos_d,
2989            head_dim,
2990            rope_dims,
2991            n_head,
2992            t,
2993            geometry.rope_base,
2994            1.0,
2995        )?;
2996        e.rope_neox(
2997            &mut k,
2998            pos_d,
2999            head_dim,
3000            rope_dims,
3001            n_head_kv,
3002            t,
3003            geometry.rope_base,
3004            1.0,
3005        )?;
3006
3007        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
3008        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
3009        {
3010            let kvl = cache.kv[il].as_mut().unwrap();
3011            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
3012            e.append_kv_quantized_rows(
3013                &k,
3014                &v,
3015                &mut kvl.k,
3016                &mut kvl.v,
3017                kvl.len,
3018                t,
3019                kvl.kv_dim_k,
3020                kvl.kv_dim_v,
3021                kvl.k_tok_bytes,
3022                kvl.v_tok_bytes,
3023                crate::Engine::kv_fp8_on(),
3024            )?;
3025            kvl.len += t;
3026            let new_len = kvl.len as i32;
3027            e.set_i32_one(&mut kvl.len_d, new_len)?;
3028        }
3029
3030        let base_len = {
3031            let kvl = cache.kv[il].as_ref().unwrap();
3032            kvl.len - t // KV rows present BEFORE this chunk's append above
3033        };
3034        Ok((AttnPre { q, k, v, gate }, base_len))
3035    }
3036
3037    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
3038    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
3039    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
3040    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
3041    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
3042    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
3043    #[allow(clippy::too_many_arguments)]
3044    fn full_attn_prime_fa_dispatch(
3045        &self,
3046        e: &Engine,
3047        q: &CudaSlice<f32>,
3048        k: &CudaSlice<f32>,
3049        v: &CudaSlice<f32>,
3050        attn: &mut CudaSlice<f32>,
3051        base_len: usize,
3052        t: usize,
3053        cache: &mut Cache,
3054        il: usize,
3055        head_dim: usize,
3056        n_head: usize,
3057        n_head_kv: usize,
3058        scale: f32,
3059    ) -> Result<(), Box<dyn std::error::Error>> {
3060        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
3061        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
3062        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
3063        // attend through the quantized cache exactly like every later chunk (quantize-then-
3064        // attend). One numeric class for every row => the chunk size cannot decide where a
3065        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
3066        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
3067        // pin-the-boundary approach).
3068        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
3069        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
3070        // with the fix unconditional, only re-introducing the class edge can prove the gate
3071        // still detects the mechanism. Never on in a measured default run.
3072        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
3073            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3074                e.sdpa_naive(
3075                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3076                )?;
3077            } else {
3078                e.fa_prefill(
3079                    q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3080                )?;
3081            }
3082            return Ok(());
3083        }
3084        let kvl = cache.kv[il].as_ref().unwrap();
3085        let t_kv = base_len + t;
3086        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3087        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3088        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
3089        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
3090        // same numeric class, so the uniform contract holds on the fallback too.
3091        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3092            e.sdpa_naive_quantized_view(
3093                q,
3094                &k_view,
3095                &v_view,
3096                attn,
3097                head_dim,
3098                n_head,
3099                n_head_kv,
3100                t,
3101                t_kv,
3102                scale,
3103                true,
3104                kvl.k_tok_bytes,
3105                kvl.v_tok_bytes,
3106            )?;
3107            return Ok(());
3108        }
3109        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
3110        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
3111        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
3112        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
3113        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
3114        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
3115        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
3116        let deqw = std::env::var("MEMRA_PRIME_DEQW")
3117            .map(|v| v != "0")
3118            .unwrap_or(true);
3119        if deqw {
3120            e.fa_prefill_view_ws(
3121                q,
3122                &k_view,
3123                &v_view,
3124                attn,
3125                head_dim,
3126                n_head,
3127                n_head_kv,
3128                t,
3129                t_kv,
3130                scale,
3131                true,
3132                kvl.k_tok_bytes,
3133                kvl.v_tok_bytes,
3134                crate::Engine::kv_fp8_on(),
3135            )?;
3136        } else {
3137            e.fa_prefill_view(
3138                q,
3139                &k_view,
3140                &v_view,
3141                attn,
3142                head_dim,
3143                n_head,
3144                n_head_kv,
3145                t,
3146                t_kv,
3147                scale,
3148                true,
3149                kvl.k_tok_bytes,
3150                kvl.v_tok_bytes,
3151                crate::Engine::kv_fp8_on(),
3152            )?;
3153        }
3154        Ok(())
3155    }
3156
3157    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
3158    /// (bit-identical composition) and hands wo its fp16 operand directly.
3159    fn full_attn_prime_post_fa(
3160        &self,
3161        e: &Engine,
3162        attn: CudaSlice<f32>,
3163        gate: &Option<CudaSlice<f32>>,
3164        t: usize,
3165        n_head: usize,
3166        head_dim: usize,
3167    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3168        let (attn_g, ag16) = match gate {
3169            Some(gate) => {
3170                let n = t * n_head * head_dim;
3171                let mut ag = e.uninit(n)?;
3172                if Self::f16out_on(e, t) {
3173                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
3174                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
3175                    (ag, Some(a16))
3176                } else {
3177                    let mut gsig = e.uninit(n)?;
3178                    e.sigmoid(gate, &mut gsig, n)?;
3179                    e.mul(&attn, &gsig, &mut ag, n)?;
3180                    (ag, None)
3181                }
3182            }
3183            None => (attn, None),
3184        };
3185        Ok((attn_g, ag16))
3186    }
3187
3188    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
3189    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
3190    /// carried THROUGH the cache like the spec verify does: carried-ring conv
3191    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
3192    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
3193    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
3194    fn linear_attn_prime(
3195        &self,
3196        e: &Engine,
3197        la: &LinearAttnLayer,
3198        h: &CudaSlice<f32>,
3199        hx: Option<&CudaSlice<u8>>,
3200        t: usize,
3201        cache: &mut Cache,
3202        il: usize,
3203    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3204        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
3205        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
3206        let g4 = match hx {
3207            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
3208            None => e.matmul_group(&ws, h, t)?,
3209        };
3210        self.linear_attn_prime_core(e, la, g4, t, cache, il)
3211    }
3212
3213    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
3214    fn linear_attn_prime_core(
3215        &self,
3216        e: &Engine,
3217        la: &LinearAttnLayer,
3218        mut g4: Vec<CudaSlice<f32>>,
3219        t: usize,
3220        cache: &mut Cache,
3221        il: usize,
3222    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3223        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
3224    }
3225
3226    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
3227    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
3228    /// conv ring writes back from the true tail. None = classic path, byte-identical.
3229    #[allow(clippy::too_many_arguments)]
3230    fn linear_attn_prime_core_pad_inner(
3231        &self,
3232        e: &Engine,
3233        la: &LinearAttnLayer,
3234        mut g4: Vec<CudaSlice<f32>>,
3235        t: usize,
3236        cache: &mut Cache,
3237        il: usize,
3238        pad_len: Option<&CudaSlice<i32>>,
3239    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3240        // shim over the view twin (task #16): full-range views of the owned buffers.
3241        let ssm = self.cfg.ssm.as_ref().unwrap();
3242        let d_state = ssm.state_size as usize;
3243        let num_k = ssm.group_count as usize;
3244        let num_v = ssm.time_step_rank as usize;
3245        let key_dim = d_state * num_k;
3246        let value_dim = d_state * num_v;
3247        let conv_dim = key_dim * 2 + value_dim;
3248        let alpha = g4.pop().unwrap(); // [T, num_v]
3249        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3250        let z = g4.pop().unwrap(); // [T, value_dim]
3251        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3252        self.linear_attn_prime_core_pad_view(
3253            e,
3254            la,
3255            &qkv_mixed.slice(0..t * conv_dim),
3256            &z.slice(0..t * value_dim),
3257            &beta_raw.slice(0..t * num_v),
3258            &alpha.slice(0..t * num_v),
3259            t,
3260            cache,
3261            il,
3262            pad_len,
3263        )
3264    }
3265
3266    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
3267    /// shared verbatim by the per-seq scan path and the varlen batched path.
3268    #[allow(clippy::too_many_arguments)]
3269    fn linear_attn_gdn_prep(
3270        &self,
3271        e: &Engine,
3272        la: &LinearAttnLayer,
3273        qkv_mixed: &cudarc::driver::CudaView<f32>,
3274        beta_raw: &cudarc::driver::CudaView<f32>,
3275        alpha: &cudarc::driver::CudaView<f32>,
3276        t: usize,
3277        cache: &mut Cache,
3278        il: usize,
3279        pad_len: Option<&CudaSlice<i32>>,
3280    ) -> Result<GdnPrep, Box<dyn std::error::Error>> {
3281        let cfg = &self.cfg;
3282        let ssm = cfg.ssm.as_ref().unwrap();
3283        let d_state = ssm.state_size as usize; // 128
3284        let num_k = ssm.group_count as usize; // 16
3285        let num_v = ssm.time_step_rank as usize; // 32
3286        let d_conv = ssm.conv_kernel as usize; // 4
3287        let key_dim = d_state * num_k; // 2048
3288        let value_dim = d_state * num_v; // 4096
3289        let conv_dim = key_dim * 2 + value_dim; // 8192
3290        let eps = cfg.rms_eps;
3291        debug_assert!(
3292            t >= d_conv - 1,
3293            "stateful conv needs T >= pad (PRIME_MIN_T gates)"
3294        );
3295
3296        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
3297        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
3298        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
3299        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
3300        let rl = cache.recur[il].as_mut().unwrap();
3301        let hk = Self::gdn_hk(e, t, num_v, num_k);
3302        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
3303        let hk = if conv_fuse { hk } else { num_v }; // de-broadcast rides the fused conv
3304        let mut q_g = e.uninit(d_state * hk * t)?;
3305        let mut k_g = e.uninit(d_state * hk * t)?;
3306        let mut v_g = e.uninit(d_state * num_v * t)?;
3307        if conv_fuse {
3308            e.ssm_conv1d_gdn_state_pad(
3309                qkv_mixed,
3310                &mut rl.conv_state,
3311                la.ssm_conv1d.float_data(),
3312                &mut q_g,
3313                &mut k_g,
3314                &mut v_g,
3315                conv_dim,
3316                t,
3317                d_conv,
3318                d_state,
3319                num_v,
3320                num_k,
3321                key_dim,
3322                hk,
3323                pad_len,
3324            )?;
3325        } else {
3326            let mut conv_out = e.uninit(conv_dim * t)?; // [conv_dim, T] channel-major, SiLU
3327            e.ssm_conv1d_tm_state_pad_v(
3328                qkv_mixed,
3329                &mut rl.conv_state,
3330                la.ssm_conv1d.float_data(),
3331                &mut conv_out,
3332                conv_dim,
3333                t,
3334                d_conv,
3335                pad_len,
3336            )?;
3337            e.qkv_to_gdn_repack(
3338                &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
3339            )?;
3340        }
3341        let mut q_l2 = e.uninit(d_state * hk * t)?;
3342        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
3343        // Emitted only where a consumer exists (the wgmma config) — on other arches the
3344        // alloc + epilogue stores would be pure waste.
3345        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
3346            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3347            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
3348            Some(qb)
3349        } else {
3350            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
3351            None
3352        };
3353        let mut k_l2 = e.uninit(d_state * hk * t)?;
3354        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
3355        let kb16 = if Engine::l2_v2_on(d_state) {
3356            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
3357            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
3358            Some(kb)
3359        } else {
3360            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
3361            None
3362        };
3363        let mut beta = e.uninit(t * num_v)?;
3364        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
3365        let mut g_log = e.uninit(t * num_v)?;
3366        e.gdn_glog_v(
3367            alpha,
3368            la.ssm_dt.float_data(),
3369            la.ssm_a.float_data(),
3370            &mut g_log,
3371            num_v,
3372            t,
3373        )?;
3374        if let Some(len_d) = pad_len {
3375            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
3376        }
3377        Ok(GdnPrep {
3378            hk,
3379            q_l2,
3380            k_l2,
3381            v_g,
3382            beta,
3383            g_log,
3384            kb16,
3385            qb16,
3386        })
3387    }
3388
3389    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
3390    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
3391    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
3392    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
3393    #[allow(clippy::too_many_arguments)]
3394    fn linear_attn_prime_core_batch(
3395        &self,
3396        e: &Engine,
3397        la: &LinearAttnLayer,
3398        g4: &[CudaSlice<f32>],
3399        offs: &[usize],
3400        ts: &[usize],
3401        caches: &mut [&mut Cache],
3402        il: usize,
3403    ) -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
3404        let ssm = self.cfg.ssm.as_ref().unwrap();
3405        let d_state = ssm.state_size as usize;
3406        let num_k = ssm.group_count as usize;
3407        let num_v = ssm.time_step_rank as usize;
3408        let key_dim = d_state * num_k;
3409        let value_dim = d_state * num_v;
3410        let conv_dim = key_dim * 2 + value_dim;
3411        let eps = self.cfg.rms_eps;
3412        let scale = 1.0 / (d_state as f32).sqrt();
3413        let b = ts.len();
3414        let c = Engine::gdn_chunk_size();
3415        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
3416        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
3417        let carried = caches.iter().any(|c| c.pos > 0);
3418        let use_vl = !carried
3419            && (2..=8).contains(&b)
3420            && Engine::gdn_chunked_enabled()
3421            && ts.iter().all(|&t| t >= 16)
3422            && e.gdn_mma_enabled(c)
3423            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
3424        if !use_vl {
3425            return (0..b)
3426                .map(|s| {
3427                    let (o, t) = (offs[s], ts[s]);
3428                    self.linear_attn_prime_core_pad_view(
3429                        e,
3430                        la,
3431                        &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
3432                        &g4[1].slice(o * value_dim..(o + t) * value_dim),
3433                        &g4[2].slice(o * num_v..(o + t) * num_v),
3434                        &g4[3].slice(o * num_v..(o + t) * num_v),
3435                        t,
3436                        caches[s],
3437                        il,
3438                        None,
3439                    )
3440                })
3441                .collect();
3442        }
3443        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
3444        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
3445        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
3446        struct SeqBufs {
3447            conv_out: CudaSlice<f32>,
3448            q_g: CudaSlice<f32>,
3449            k_g: CudaSlice<f32>,
3450            v_g: CudaSlice<f32>,
3451            q_l2: CudaSlice<f32>,
3452            k_l2: CudaSlice<f32>,
3453            beta: CudaSlice<f32>,
3454            g_log: CudaSlice<f32>,
3455            gn: CudaSlice<f32>,
3456            gn16: CudaSlice<u8>,
3457        }
3458        let d_conv = ssm.conv_kernel as usize;
3459        let f16o = Self::f16out_on(e, 16);
3460        let hk = Self::gdn_hk(e, 16, num_v, num_k); // vl path is always chunked+mma
3461        let mut sb = Vec::with_capacity(b);
3462        let mut pres = Vec::with_capacity(b);
3463        for &t in ts.iter().take(b) {
3464            sb.push(SeqBufs {
3465                conv_out: e.uninit(conv_dim * t)?,
3466                q_g: e.uninit(d_state * hk * t)?,
3467                k_g: e.uninit(d_state * hk * t)?,
3468                v_g: e.uninit(d_state * num_v * t)?,
3469                q_l2: e.uninit(d_state * hk * t)?,
3470                k_l2: e.uninit(d_state * hk * t)?,
3471                beta: e.uninit(t * num_v)?,
3472                g_log: e.uninit(t * num_v)?,
3473                gn: e.uninit(d_state * num_v * t)?,
3474                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
3475            });
3476            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
3477        }
3478        let prep_args: Vec<crate::GdnPrepVl> = (0..b)
3479            .map(|s| {
3480                let (o, t) = (offs[s], ts[s]);
3481                let rl = caches[s].recur[il].as_ref().unwrap();
3482                crate::GdnPrepVl {
3483                    qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
3484                    conv_state: e.addr_f32(&rl.conv_state),
3485                    conv_out: e.addr_f32(&sb[s].conv_out),
3486                    q_g: e.addr_f32(&sb[s].q_g),
3487                    k_g: e.addr_f32(&sb[s].k_g),
3488                    v_g: e.addr_f32(&sb[s].v_g),
3489                    q_l2: e.addr_f32(&sb[s].q_l2),
3490                    k_l2: e.addr_f32(&sb[s].k_l2),
3491                    beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
3492                    alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
3493                    beta: e.addr_f32(&sb[s].beta),
3494                    g_log: e.addr_f32(&sb[s].g_log),
3495                    o: e.addr_f32(&pres[s].o),
3496                    z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
3497                    gn: e.addr_f32(&sb[s].gn),
3498                    gn16: e.addr_u8(&sb[s].gn16),
3499                    kb16: if Engine::l2_v2_on(d_state) {
3500                        e.addr_u8(&pres[s].kb16)
3501                    } else {
3502                        0
3503                    },
3504                    qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) {
3505                        e.addr_u8(&pres[s].qb16)
3506                    } else {
3507                        0
3508                    },
3509                    t: t as i32,
3510                    pad: 0,
3511                }
3512            })
3513            .collect();
3514        let args: Vec<crate::GdnSeqVl> = (0..b)
3515            .map(|s| {
3516                let rl = caches[s].recur[il].as_ref().unwrap();
3517                crate::GdnSeqVl {
3518                    kb16: e.addr_u8(&pres[s].kb16),
3519                    gcum: e.addr_f32(&pres[s].gcum),
3520                    beta: e.addr_f32(&sb[s].beta),
3521                    u: e.addr_f32(&pres[s].u),
3522                    wb16: e.addr_u8(&pres[s].wb16),
3523                    y: e.addr_u8(&pres[s].y16),
3524                    ssnap: e.addr_u8(&pres[s].ssnap16),
3525                    state_in: e.addr_f32(&rl.ssm_state),
3526                    state_out: e.addr_f32(&rl.ssm_state_alt),
3527                    q: e.addr_f32(&sb[s].q_l2),
3528                    p: e.addr_f32(&pres[s].p),
3529                    o: e.addr_f32(&pres[s].o),
3530                    k: e.addr_f32(&sb[s].k_l2),
3531                    v: e.addr_f32(&sb[s].v_g),
3532                    g: e.addr_f32(&sb[s].g_log),
3533                    a: e.addr_f32(&pres[s].a),
3534                    w: e.addr_f32(&pres[s].w),
3535                    t: ts[s] as i32,
3536                    nc: pres[s].nc as i32,
3537                }
3538            })
3539            .collect();
3540        e.gdn_prep_vl8(
3541            &prep_args,
3542            la.ssm_conv1d.float_data(),
3543            la.ssm_dt.float_data(),
3544            la.ssm_a.float_data(),
3545            conv_dim,
3546            d_conv,
3547            d_state,
3548            num_v,
3549            num_k,
3550            key_dim,
3551            hk,
3552            eps,
3553        )?;
3554        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
3555        // both standalone mirror launches vanish on the default config.
3556        if !Engine::l2_v2_on(d_state) {
3557            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
3558        }
3559        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
3560        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
3561            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
3562            if !Engine::l2_v2_on(d_state) {
3563                for s in 0..b {
3564                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
3565                }
3566            }
3567            let mut wa = [crate::GdnWVl::default(); 8];
3568            for s in 0..b {
3569                wa[s] = crate::GdnWVl {
3570                    qb16: e.addr_u8(&pres[s].qb16),
3571                    pb16: e.addr_u8(&pres[s].pb16),
3572                };
3573            }
3574            Some(crate::GdnWVl8(wa))
3575        } else {
3576            None
3577        };
3578        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
3579        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
3580        if f16o {
3581            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
3582        }
3583        // per-seq state swap (+ non-f16out tail fallback)
3584        let mut out = Vec::with_capacity(b);
3585        for (s, bufs) in sb.into_iter().enumerate() {
3586            let rl = caches[s].recur[il].as_mut().unwrap();
3587            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3588            let (o, t) = (offs[s], ts[s]);
3589            let SeqBufs { mut gn, gn16, .. } = bufs;
3590            if f16o {
3591                out.push((gn, Some(gn16)));
3592            } else {
3593                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
3594                e.gated_rmsnorm_zv(
3595                    &pres[s].o,
3596                    la.ssm_norm.float_data(),
3597                    &z_v,
3598                    &mut gn,
3599                    d_state,
3600                    num_v * t,
3601                    eps,
3602                )?;
3603                out.push((gn, None));
3604            }
3605        }
3606        Ok(out)
3607    }
3608
3609    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
3610    /// views of the CONCAT projection outputs directly (no per-seq split copies).
3611    /// Same kernels, same values, byte-identical to the Vec shim above.
3612    #[allow(clippy::too_many_arguments)]
3613    fn linear_attn_prime_core_pad_view(
3614        &self,
3615        e: &Engine,
3616        la: &LinearAttnLayer,
3617        qkv_mixed: &cudarc::driver::CudaView<f32>,
3618        z: &cudarc::driver::CudaView<f32>,
3619        beta_raw: &cudarc::driver::CudaView<f32>,
3620        alpha: &cudarc::driver::CudaView<f32>,
3621        t: usize,
3622        cache: &mut Cache,
3623        il: usize,
3624        pad_len: Option<&CudaSlice<i32>>,
3625    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
3626        let cfg = &self.cfg;
3627        let ssm = cfg.ssm.as_ref().unwrap();
3628        let d_state = ssm.state_size as usize; // 128
3629        let num_v = ssm.time_step_rank as usize; // 32
3630        let eps = cfg.rms_eps;
3631        let scale = 1.0 / (d_state as f32).sqrt();
3632
3633        let prep =
3634            self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
3635
3636        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
3637        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
3638        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
3639        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
3640        // verify keep the sequential kernel).
3641        let mut o = e.uninit(d_state * num_v * t)?;
3642        let rl = cache.recur[il].as_mut().unwrap();
3643        {
3644            let crate::cache::RecurLayer {
3645                ssm_state,
3646                ssm_state_alt,
3647                ..
3648            } = rl;
3649            e.gdn_scan_prefill(
3650                &prep.q_l2,
3651                &prep.k_l2,
3652                &prep.v_g,
3653                &prep.g_log,
3654                &prep.beta,
3655                prep.kb16.as_ref(),
3656                prep.qb16.as_ref(),
3657                ssm_state,
3658                ssm_state_alt,
3659                &mut o,
3660                num_v,
3661                t,
3662                scale,
3663                prep.hk,
3664            )?;
3665        }
3666        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3667
3668        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
3669        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
3670        let mut gn = e.uninit(d_state * num_v * t)?;
3671        let gn16 = if Self::f16out_on(e, t) {
3672            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
3673            e.gated_rmsnorm_f16out_zv(
3674                &o,
3675                la.ssm_norm.float_data(),
3676                z,
3677                &mut gn,
3678                &mut g16,
3679                d_state,
3680                num_v * t,
3681                eps,
3682            )?;
3683            Some(g16)
3684        } else {
3685            e.gated_rmsnorm_zv(
3686                &o,
3687                la.ssm_norm.float_data(),
3688                z,
3689                &mut gn,
3690                d_state,
3691                num_v * t,
3692                eps,
3693            )?;
3694            None
3695        };
3696        Ok((gn, gn16))
3697    }
3698
3699    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
3700    #[allow(clippy::too_many_arguments)]
3701    fn linear_attn_prime_core_pad(
3702        &self,
3703        e: &Engine,
3704        la: &LinearAttnLayer,
3705        g4: Vec<CudaSlice<f32>>,
3706        t: usize,
3707        cache: &mut Cache,
3708        il: usize,
3709        pad_len: Option<&CudaSlice<i32>>,
3710    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3711        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
3712        if let Some(xh) = &gn16 {
3713            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
3714                return Ok(y);
3715            }
3716        }
3717        Ok(e.matmul(&la.ssm_out, &gn, t)?)
3718    }
3719
3720    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
3721    ///
3722    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
3723    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
3724    pub fn full_attn(
3725        &self,
3726        e: &Engine,
3727        fa: &FullAttnLayer,
3728        h: &CudaSlice<f32>,
3729        pos_d: &CudaSlice<i32>,
3730        t: usize,
3731        il: usize,
3732    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3733        if self.cfg.step35.is_some() {
3734            return self.step35_attn(e, fa, h, pos_d, t, il);
3735        }
3736        let cfg = &self.cfg;
3737        let _n_embd = cfg.n_embd as usize;
3738        let geometry = cfg.full_attention_geometry_at(il as u32);
3739        let n_head = geometry.n_head as usize;
3740        let n_head_kv = geometry.n_head_kv as usize;
3741        let head_dim = geometry.head_dim_k as usize;
3742        let eps = cfg.rms_eps;
3743        let scale = geometry.attention_scale();
3744
3745        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
3746        // gate — wq out = n_head*head_dim, no split (see prime-path note).
3747        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3748        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
3749        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
3750        let v = g3.pop().unwrap();
3751        let mut k = g3.pop().unwrap();
3752        let qf = g3.pop().unwrap();
3753        let (mut q, gate) = if gated {
3754            let mut q = e.uninit(t * n_head * head_dim)?;
3755            let mut gate = e.uninit(t * n_head * head_dim)?;
3756            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
3757            (q, Some(gate))
3758        } else {
3759            (qf, None)
3760        };
3761
3762        // QK-norm (per head_dim row), then partial RoPE.
3763        let mut qn = e.uninit(t * n_head * head_dim)?;
3764        e.rms_norm(
3765            &q,
3766            fa.q_norm.float_data(),
3767            &mut qn,
3768            head_dim,
3769            n_head * t,
3770            eps,
3771        )?;
3772        q = qn;
3773        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
3774        e.rms_norm(
3775            &k,
3776            fa.k_norm.float_data(),
3777            &mut kn,
3778            head_dim,
3779            n_head_kv * t,
3780            eps,
3781        )?;
3782        k = kn;
3783        let rope_dims = geometry.n_rot as usize;
3784        e.rope_neox(
3785            &mut q,
3786            pos_d,
3787            head_dim,
3788            rope_dims,
3789            n_head,
3790            t,
3791            geometry.rope_base,
3792            1.0,
3793        )?;
3794        e.rope_neox(
3795            &mut k,
3796            pos_d,
3797            head_dim,
3798            rope_dims,
3799            n_head_kv,
3800            t,
3801            geometry.rope_base,
3802            1.0,
3803        )?;
3804
3805        // SDPA
3806        let mut attn = e.uninit(t * n_head * head_dim)?;
3807        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
3808        // falls back to naive sdpa.
3809        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3810            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
3811            e.sdpa_naive(
3812                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3813            )?;
3814        } else {
3815            e.fa_prefill(
3816                &q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true,
3817            )?;
3818        }
3819
3820        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
3821        let attn_g = match &gate {
3822            Some(gate) => {
3823                let mut gsig = e.uninit(t * n_head * head_dim)?;
3824                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3825                let mut ag = e.uninit(t * n_head * head_dim)?;
3826                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3827                ag
3828            }
3829            None => attn,
3830        };
3831
3832        // o projection
3833        let o = e.matmul(&fa.wo, &attn_g, t)?;
3834        Ok(o)
3835    }
3836
3837    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
3838    pub fn linear_attn(
3839        &self,
3840        e: &Engine,
3841        la: &LinearAttnLayer,
3842        h: &CudaSlice<f32>,
3843        t: usize,
3844    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3845        let cfg = &self.cfg;
3846        let _n_embd = cfg.n_embd as usize;
3847        let ssm = cfg.ssm.as_ref().unwrap();
3848        let d_state = ssm.state_size as usize; // 128
3849        let num_k = ssm.group_count as usize; // 16
3850        let num_v = ssm.time_step_rank as usize; // 32
3851        let d_conv = ssm.conv_kernel as usize; // 4
3852        let head_k = d_state;
3853        let head_v = d_state;
3854        let key_dim = head_k * num_k; // 2048
3855        let value_dim = head_v * num_v; // 4096
3856        let conv_dim = key_dim * 2 + value_dim; // 8192
3857        let eps = cfg.rms_eps;
3858        let scale = 1.0 / (d_state as f32).sqrt();
3859
3860        // projections
3861        // grouped: one f16 activation convert feeds all four projections (matmul_group)
3862        let mut g4 = e.matmul_group(
3863            &[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
3864            h,
3865            t,
3866        )?;
3867        let alpha = g4.pop().unwrap(); // [T, num_v]
3868        let beta_raw = g4.pop().unwrap(); // [T, num_v]
3869        let z = g4.pop().unwrap(); // [T, value_dim]
3870        let qkv_mixed = g4.pop().unwrap(); // [T, conv_dim] token-major
3871
3872        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
3873        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
3874        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
3875        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
3876        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
3877        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
3878        let _ = (head_k, head_v);
3879        let mut q_g = e.uninit(d_state * num_v * t)?;
3880        let mut k_g = e.uninit(d_state * num_v * t)?;
3881        let mut v_g = e.uninit(d_state * num_v * t)?;
3882        e.ssm_conv1d_gdn(
3883            &qkv_mixed,
3884            la.ssm_conv1d.float_data(),
3885            &mut q_g,
3886            &mut k_g,
3887            &mut v_g,
3888            conv_dim,
3889            t,
3890            d_conv,
3891            d_state,
3892            num_v,
3893            num_k,
3894            key_dim,
3895        )?;
3896        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
3897        let mut q_l2 = e.uninit(d_state * num_v * t)?;
3898        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3899        let mut k_l2 = e.uninit(d_state * num_v * t)?;
3900        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3901        let v_gd = v_g;
3902
3903        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
3904        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
3905        let mut beta = e.uninit(t * num_v)?;
3906        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3907        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
3908        let mut g_log = e.uninit(t * num_v)?;
3909        e.gdn_glog(
3910            &alpha,
3911            la.ssm_dt.float_data(),
3912            la.ssm_a.float_data(),
3913            &mut g_log,
3914            num_v,
3915            t,
3916        )?;
3917
3918        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
3919        let state_in = e.zeros(d_state * d_state * num_v)?; // zero state (prefill)
3920        let mut state_out = e.zeros(d_state * d_state * num_v)?;
3921        let mut o = e.uninit(d_state * num_v * t)?;
3922        e.gdn_scan_prefill(
3923            &q_l2,
3924            &k_l2,
3925            &v_gd,
3926            &g_log,
3927            &beta,
3928            None,
3929            None,
3930            &state_in,
3931            &mut state_out,
3932            &mut o,
3933            num_v,
3934            t,
3935            scale,
3936            num_v,
3937        )?;
3938
3939        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
3940        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
3941        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
3942        // o rows are (t*num_v+vh) too. Good.
3943        let mut gn = e.uninit(d_state * num_v * t)?;
3944        e.gated_rmsnorm(
3945            &o,
3946            la.ssm_norm.float_data(),
3947            &z,
3948            &mut gn,
3949            d_state,
3950            num_v * t,
3951            eps,
3952        )?;
3953
3954        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
3955        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
3956        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
3957        let out = e.matmul(&la.ssm_out, &gn, t)?;
3958        Ok(out)
3959    }
3960}
3961
3962impl HybridModel {
3963    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
3964    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
3965    ///
3966    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
3967    /// different 860160-byte block than the same expert of layer 7).
3968    ///
3969    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
3970    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
3971    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
3972    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
3973    pub fn moe_ffn_il(
3974        &self,
3975        e: &Engine,
3976        m: &MoeWeights,
3977        z: &CudaSlice<f32>,
3978        t: usize,
3979        il: u16,
3980    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3981        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
3982    }
3983
3984    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
3985    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
3986    pub fn moe_ffn_il_prefill(
3987        &self,
3988        e: &Engine,
3989        m: &MoeWeights,
3990        z: &CudaSlice<f32>,
3991        t: usize,
3992        il: u16,
3993    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3994        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
3995    }
3996
3997    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
3998    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
3999    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
4000    pub fn moe_ffn_il_zq8(
4001        &self,
4002        e: &Engine,
4003        m: &MoeWeights,
4004        z: &CudaSlice<f32>,
4005        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4006        t: usize,
4007        il: u16,
4008    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4009        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false)
4010    }
4011
4012    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
4013    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
4014    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
4015    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
4016    ///
4017    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
4018    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
4019    pub(crate) fn moe_ffn(
4020        e: &Engine,
4021        m: &MoeWeights,
4022        z: &CudaSlice<f32>,
4023        t: usize,
4024        cfg: &ModelConfig,
4025        il: u16,
4026        max_block: usize,
4027    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4028        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
4029    }
4030
4031    #[allow(clippy::too_many_arguments)]
4032    pub(crate) fn moe_ffn_inner(
4033        e: &Engine,
4034        m: &MoeWeights,
4035        z: &CudaSlice<f32>,
4036        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4037        t: usize,
4038        cfg: &ModelConfig,
4039        il: u16,
4040        max_block: usize,
4041        prefill: bool,
4042    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4043        let worker_io = crate::spill_pread::worker_enabled();
4044        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
4045        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
4046            e.with_moe_cache(max_block, |cache, _| {
4047                cache.begin_forward_epoch(il, t);
4048                if worker_io {
4049                    cache.begin_worker_scope();
4050                }
4051                Ok(())
4052            })?;
4053        }
4054        if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
4055            let moe = cfg.moe.as_ref().unwrap();
4056            let n_expert = moe.expert_count as usize;
4057            let n_used = moe.expert_used_count as usize;
4058            let sigmoid = cfg.sigmoid_router().unwrap();
4059            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4060            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sigmoid)?;
4061            return Self::moe_ffn_sigmoid_dev(e, m, z, zq8, &logits, t, cfg, il, sigmoid);
4062        }
4063        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
4064        // current caller into this research arm; the naked default stays on the established path.
4065        if t > 1 && moe_grouped_enabled(cfg, prefill) {
4066            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
4067            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
4068            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
4069            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
4070            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
4071            if std::env::var("MEMRA_MOE_GATE").is_ok() {
4072                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
4073                let g_host = e.dtoh(&grouped_out)?;
4074                let s_host = e.dtoh(&seq_out)?;
4075                let g_bytes: &[u8] = unsafe {
4076                    std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4)
4077                };
4078                let s_bytes: &[u8] = unsafe {
4079                    std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4)
4080                };
4081                if g_bytes == s_bytes {
4082                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
4083                } else {
4084                    let diffs = g_host
4085                        .iter()
4086                        .zip(s_host.iter())
4087                        .enumerate()
4088                        .filter(|(_, (a, b))| a != b)
4089                        .count();
4090                    let maxdiff = g_host
4091                        .iter()
4092                        .zip(s_host.iter())
4093                        .map(|(a, b)| (a - b).abs())
4094                        .fold(0.0f32, f32::max);
4095                    panic!(
4096                        "moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}",
4097                        g_host.len()
4098                    );
4099                }
4100            }
4101            return Ok(grouped_out);
4102        }
4103        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
4104    }
4105
4106    fn sigmoid_resident_dev_eligible(e: &Engine, m: &MoeWeights, cfg: &ModelConfig) -> bool {
4107        let Some(moe) = cfg.moe.as_ref() else {
4108            return false;
4109        };
4110        // Cached once per process: this predicate runs per MoE layer per decode step, and five
4111        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
4112        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4113        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
4114            std::env::var("MEMRA_MOE_STATS").is_ok()
4115                || std::env::var("MEMRA_MOE_TRACE").is_ok()
4116                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4117                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
4118                || std::env::var("MEMRA_MOE_GATE").is_ok()
4119        });
4120        cfg.step35.is_some()
4121            && sigmoid_router_enabled()
4122            && moe_dev_enabled()
4123            && moe_slab_enabled()
4124            && !observation_mode
4125            && moe.expert_used_count <= 8
4126            && m.has_uniform_expert_layout()
4127            && m.gate_exps.macros.is_none()
4128            && m.up_exps.macros.is_none()
4129            && m.down_exps.macros.is_none()
4130            && !m.has_macros
4131            && moe_q8_enabled()
4132            && q8_expert_supported(m.gate_exps.qtype)
4133            && q8_expert_supported(m.up_exps.qtype)
4134            && q8_expert_supported(m.down_exps.qtype)
4135            && m.dev_exps
4136                .as_ref()
4137                .is_some_and(|dev| dev.dev == e.ctx().ordinal())
4138    }
4139
4140    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
4141    pub(crate) fn moe_ffn_sequential(
4142        e: &Engine,
4143        m: &MoeWeights,
4144        z: &CudaSlice<f32>,
4145        t: usize,
4146        cfg: &ModelConfig,
4147        il: u16,
4148        max_block: usize,
4149    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4150        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
4151    }
4152
4153    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
4154    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
4155    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
4156    fn moe_router_logits(
4157        e: &Engine,
4158        m: &MoeWeights,
4159        z: &CudaSlice<f32>,
4160        t: usize,
4161        cfg: &ModelConfig,
4162    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4163        if t < PRIME_MIN_T {
4164            // Decode and speculative verify use one fixed per-row reduction program.
4165            if crate::router_kernel_on() {
4166                e.router_gemv(
4167                    m.gate_inp.float_data(),
4168                    z,
4169                    cfg.n_embd as usize,
4170                    m.gate_exps.n_expert,
4171                    t,
4172                )
4173            } else {
4174                e.matmul_decode_exact(&m.gate_inp, z, t)
4175            }
4176        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
4177            e.router_gemv(
4178                m.gate_inp.float_data(),
4179                z,
4180                cfg.n_embd as usize,
4181                m.gate_exps.n_expert,
4182                t,
4183            )
4184        } else {
4185            e.matmul(&m.gate_inp, z, t)
4186        }
4187    }
4188
4189    /// Append the host-visible router selection for one layer/forward when calibration tracing is
4190    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
4191    /// trace is independent of the dispatch optimization selected for the forward.
4192    fn trace_moe_routes(
4193        il: u16,
4194        t: usize,
4195        sel_all: &[u32],
4196        weights: &[f32],
4197    ) -> Result<(), Box<dyn std::error::Error>> {
4198        use std::io::Write as _;
4199        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
4200            let mut f = std::fs::OpenOptions::new()
4201                .create(true)
4202                .append(true)
4203                .open(path)?;
4204            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
4205            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
4206        }
4207        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
4208            let mut f = std::fs::OpenOptions::new()
4209                .create(true)
4210                .append(true)
4211                .open(path)?;
4212            let pairs: Vec<String> = sel_all
4213                .iter()
4214                .zip(weights)
4215                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
4216                .collect();
4217            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
4218        }
4219        Ok(())
4220    }
4221
4222    #[allow(clippy::too_many_arguments)]
4223    fn trace_sigmoid_router_logits(
4224        e: &Engine,
4225        il: u16,
4226        t: usize,
4227        n_expert: usize,
4228        n_used: usize,
4229        logits: &CudaSlice<f32>,
4230        m: &MoeWeights,
4231        (scaling_factor, route_norm): (f32, bool),
4232    ) -> Result<(), Box<dyn std::error::Error>> {
4233        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
4234            return Ok(());
4235        }
4236        let logits = e.dtoh(logits)?;
4237        let active: Vec<u8> = m
4238            .active_experts
4239            .as_ref()
4240            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
4241            .unwrap_or_else(|| vec![1; n_expert]);
4242        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
4243        crate::sigrouter_contract::capture_served_logits(
4244            il as u32,
4245            t,
4246            n_expert,
4247            n_used,
4248            scaling_factor,
4249            route_norm,
4250            &active,
4251            &bias,
4252            &logits,
4253        )?;
4254        Ok(())
4255    }
4256
4257    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
4258    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
4259    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
4260    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
4261    fn trace_moe_input(
4262        e: &Engine,
4263        il: u16,
4264        t: usize,
4265        n_embd: usize,
4266        z: &CudaSlice<f32>,
4267    ) -> Result<(), Box<dyn std::error::Error>> {
4268        use std::io::Write as _;
4269        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else {
4270            return Ok(());
4271        };
4272        let host = e.dtoh(z)?;
4273        if host.len() != t * n_embd {
4274            return Err(format!(
4275                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
4276                host.len(),
4277                t,
4278                n_embd
4279            )
4280            .into());
4281        }
4282        let bytes = unsafe {
4283            std::slice::from_raw_parts(
4284                host.as_ptr().cast::<u8>(),
4285                host.len() * std::mem::size_of::<f32>(),
4286            )
4287        };
4288        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
4289        let mut state = state
4290            .lock()
4291            .map_err(|_| "MoE input trace writer lock is poisoned")?;
4292        if state.is_none() {
4293            let dir = std::path::PathBuf::from(&dir);
4294            std::fs::create_dir_all(&dir)?;
4295            let index = std::fs::OpenOptions::new()
4296                .create(true)
4297                .append(true)
4298                .open(dir.join("index.jsonl"))?;
4299            *state = Some(MoeInputTraceWriter {
4300                dir,
4301                index,
4302                payloads: std::collections::HashMap::new(),
4303            });
4304        }
4305        let writer = state.as_mut().unwrap();
4306        if writer.dir != std::path::Path::new(&dir) {
4307            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
4308        }
4309        let file_name = format!("layer-{il:03}.f32");
4310        if !writer.payloads.contains_key(&il) {
4311            let payload = std::fs::OpenOptions::new()
4312                .create(true)
4313                .append(true)
4314                .open(writer.dir.join(&file_name))?;
4315            let offset = payload.metadata()?.len();
4316            writer.payloads.insert(il, (payload, offset));
4317        }
4318        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
4319        let row_offset = *offset;
4320        payload.write_all(bytes)?;
4321        *offset += bytes.len() as u64;
4322        writeln!(
4323            writer.index,
4324            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
4325             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
4326             \"payload_bytes\":{}}}",
4327            bytes.len()
4328        )?;
4329        Ok(())
4330    }
4331
4332    #[allow(clippy::too_many_arguments)]
4333    pub(crate) fn moe_ffn_sequential_zq8(
4334        e: &Engine,
4335        m: &MoeWeights,
4336        z: &CudaSlice<f32>,
4337        zq8: Option<&(CudaSlice<i8>, 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        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4344        let moe = cfg.moe.as_ref().unwrap();
4345        let n_embd = cfg.n_embd as usize; // 2048 (gate/up in_f, down out_f)
4346        let n_expert = moe.expert_count as usize; // 256
4347        let n_used = moe.expert_used_count as usize; // 8
4348        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
4349
4350        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
4351        debug_assert_eq!(m.gate_exps.in_f, n_embd);
4352        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
4353        debug_assert_eq!(m.down_exps.in_f, n_ff_exp); // down is TRANSPOSED: in=512
4354        debug_assert_eq!(m.down_exps.out_f, n_embd); //                     out=2048
4355        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
4356
4357        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
4358        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
4359        let lim_exp = cfg.clamp_exp_at(il as u32);
4360        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4361        let use_cache = Engine::moe_cache_enabled();
4362        let uniform_experts = m.has_uniform_expert_layout();
4363        let moe_q8 = uniform_experts
4364            && moe_q8_enabled()
4365            && q8_expert_supported(m.gate_exps.qtype)
4366            && q8_expert_supported(m.up_exps.qtype)
4367            && q8_expert_supported(m.down_exps.qtype);
4368        // Experimental secondary backend: complete experts already resident in the SLRU stay on
4369        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
4370        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
4371        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
4372        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
4373        // commands and CI have no llama.cpp or OpenMP dependency.
4374        let cpu_expert_requested = crate::cpu_experts::configured();
4375        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
4376            return Err(std::io::Error::other(
4377                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
4378            )
4379            .into());
4380        }
4381        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
4382        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
4383        // Those backends are each deterministic but are different numeric configurations, so a
4384        // later prefill eviction can change greedy output. Freeze after the first real prefill;
4385        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
4386        // staging below and cannot change backend assignment.
4387        let freeze_cpu_residency = cpu_expert_requested
4388            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
4389        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
4390            .ok()
4391            .and_then(|value| value.parse::<usize>().ok())
4392            .is_some_and(|tokens| tokens > 0);
4393        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
4394            e.freeze_moe_cache();
4395        }
4396        let cache_frozen = use_cache && e.moe_cache_frozen();
4397        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
4398
4399        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
4400        // cannot change logits, selected expert ids, or routing weights.
4401        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
4402        if let Some(sig) = cfg.sigmoid_router() {
4403            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
4404        }
4405
4406        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
4407        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
4408        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
4409        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
4410        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
4411        // per-token host stall that dominated the 35B decode wall after stages 1+2.
4412        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
4413        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
4414        // only difference is where sel/w/pointers are READ from (device instead of params).
4415        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
4416        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
4417        // Any non-resident layer falls through to host routing + the gdec/sequential path.
4418        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
4419        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
4420        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
4421        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
4422        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
4423        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
4424        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
4425        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
4426        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
4427        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
4428        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
4429        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
4430        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
4431        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
4432        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
4433        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
4434        // now rides the dev loop below (same kernels per token as decode); pairs serves real
4435        // prefill (t >= 16, where spec never verifies).
4436        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
4437        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
4438        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
4439        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
4440        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
4441        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
4442        // ride the macro-aware sequential/staged paths below or every expert output is off by
4443        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
4444        let no_exp_macros = m.gate_exps.macros.is_none()
4445            && m.up_exps.macros.is_none()
4446            && m.down_exps.macros.is_none();
4447        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
4448        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
4449        // so it cannot even see the per-layer limit.
4450        if cfg.sigmoid_router().is_none()
4451            && cfg.m3.is_none()
4452            && cfg.hy3.is_none()
4453            && !cfg.swiglu_clamped_at(il as u32)
4454            && no_exp_macros
4455            && t >= PRIME_MIN_T
4456            && m.dev_exps.is_some()
4457            && moe_q8_enabled()
4458            && q8_expert_supported(m.gate_exps.qtype)
4459            && q8_expert_supported(m.up_exps.qtype)
4460            && q8_expert_supported(m.down_exps.qtype)
4461            && std::env::var("MEMRA_MOE_PAIRS")
4462                .map(|v| v != "0")
4463                .unwrap_or(true)
4464            && std::env::var("MEMRA_MOE_STATS").is_err()
4465        {
4466            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
4467        }
4468
4469        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
4470        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
4471        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
4472        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
4473        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
4474        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
4475        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
4476        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
4477        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
4478        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
4479        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
4480        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
4481        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
4482        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
4483        // Keyed off sigmoid_router() so arch #4 is denied by construction.
4484        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
4485        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
4486        let dev_ok = uniform_experts
4487            && cfg.sigmoid_router().is_none()
4488            && cfg.m3.is_none()
4489            && cfg.hy3.is_none()
4490            && !cfg.swiglu_clamped_at(il as u32);
4491        // Observation modes must route through the host-visible selection below. Otherwise a fully
4492        // resident layer returns through device dispatch before its trace/stats row is recorded,
4493        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
4494        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
4495            || std::env::var("MEMRA_MOE_TRACE").is_ok()
4496            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
4497            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
4498        if dev_ok
4499            && t < PRIME_MIN_T
4500            && m.dev_exps.is_some()
4501            && n_used <= 8
4502            && moe_dev_enabled()
4503            && !observe_routes
4504        {
4505            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4506        }
4507        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled() && !observe_routes {
4508            let row_ok = e.with_moe_cache(max_block, |c, eng| {
4509                if moe_prewarm_enabled() {
4510                    c.prewarm_layer(il, m, eng)?;
4511                }
4512                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
4513            })?;
4514            if row_ok {
4515                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
4516            }
4517        }
4518
4519        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
4520        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
4521            if cpu_hybrid {
4522                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
4523                    e,
4524                    &logits,
4525                    z,
4526                    t,
4527                    n_expert,
4528                    n_used,
4529                    m.exp_probs_b.as_deref(),
4530                    sig,
4531                    m.active_experts.as_deref(),
4532                )?;
4533                (sel, w, Some(input))
4534            } else {
4535                let (sel, w) =
4536                    Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?;
4537                (sel, w, None)
4538            }
4539        } else {
4540            let (sel, w) =
4541                Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?;
4542            (sel, w, None)
4543        };
4544        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
4545
4546        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
4547        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
4548        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
4549        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4550        Self::trace_moe_input(e, il, t, n_embd, z)?;
4551
4552        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
4553        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
4554        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
4555        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
4556        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
4557        // wait for each pending block, so later copies can overlap the earlier expert kernels while
4558        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
4559        // T=1; batched forwards can have token-local consumers still in flight between selections.
4560        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
4561        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
4562        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
4563        let worker_disk_prefetch =
4564            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
4565        let promote_worker_h2d =
4566            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
4567        if promote_worker_h2d {
4568            let mut selected_blocks = Vec::with_capacity(n_used * 3);
4569            for &ex in sel_all.iter().take(n_used) {
4570                let ex = ex as u16;
4571                selected_blocks.extend([
4572                    BlockId::new(il, PROJ_GATE, ex),
4573                    BlockId::new(il, PROJ_UP, ex),
4574                    BlockId::new(il, PROJ_DOWN, ex),
4575                ]);
4576            }
4577            for &ex in sel_all.iter().take(n_used) {
4578                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
4579            }
4580            e.with_moe_cache(max_block, |cache, eng| {
4581                cache.promote_worker_reads_at_safe_boundary(
4582                    &selected_blocks,
4583                    &selected_blocks,
4584                    eng,
4585                )?;
4586                Ok(())
4587            })?;
4588        }
4589
4590        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
4591        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
4592        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
4593            let mut cnt = vec![0u32; n_expert];
4594            for &s in sel_all.iter() {
4595                cnt[s as usize] += 1;
4596            }
4597            let total = sel_all.len() as f64;
4598            let mut h = 0.0f64;
4599            let mut active = 0usize;
4600            for &c in &cnt {
4601                if c > 0 {
4602                    active += 1;
4603                    let p = c as f64 / total;
4604                    h -= p * p.log2();
4605                }
4606            }
4607            let maxc = cnt.iter().copied().max().unwrap_or(0);
4608            println!(
4609                "moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
4610                il,
4611                t,
4612                sel_all.len(),
4613                active,
4614                n_expert,
4615                h,
4616                (n_expert as f64).log2(),
4617                total / active.max(1) as f64,
4618                maxc
4619            );
4620        }
4621
4622        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
4623        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
4624        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
4625        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
4626        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
4627        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
4628        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
4629        // zeroed-then-accumulated exactly as before (fallback).
4630        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
4631        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
4632        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
4633        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
4634        let gdec_may_fire = uniform_experts
4635            && use_cache
4636            && n_used <= 8
4637            && gdec_enabled()
4638            && !cfg.swiglu_clamped_at(il as u32);
4639        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
4640        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
4641        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
4642        // archs the slabs were uploaded but never read, and every expert went through the
4643        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
4644        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
4645        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
4646        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
4647        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
4648        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
4649        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
4650        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
4651        // strictly worse than staging); under PP-2 without the prime walker this admits
4652        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
4653        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
4654        let slab_local = m
4655            .dev_exps
4656            .as_ref()
4657            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
4658        let slab_bases = slab_local.map(|d| {
4659            use cudarc::driver::DevicePtr;
4660            let s = e.stream();
4661            let (pg, _g0) = d.gate.device_ptr(&s);
4662            let (pu, _g1) = d.up.device_ptr(&s);
4663            let (pd, _g2) = d.down.device_ptr(&s);
4664            (pg as u64, pu as u64, pd as u64)
4665        });
4666        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
4667        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
4668        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
4669        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
4670        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
4671        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
4672        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
4673        // all-resident tokens, staged loop for misses), which is a dispatch-class
4674        // comparison, not a provenance one.
4675        let slab_fused_may_fire = slab_bases.is_some()
4676            && n_used <= 8
4677            && gdec_enabled()
4678            && !cfg.swiglu_clamped_at(il as u32)
4679            && cfg.m3.is_none()
4680            && no_exp_macros
4681            && moe_q8;
4682        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
4683        // uninit; a token that falls through to any accumulating loop zeroes its own row.
4684        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
4685            e.uninit(t * n_embd)?
4686        } else {
4687            e.zeros(t * n_embd)?
4688        };
4689        // The router readback above already established a host boundary. Copy each small-t hidden
4690        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
4691        let cpu_input = if cpu_hybrid {
4692            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
4693        } else {
4694            None
4695        };
4696
4697        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
4698        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
4699        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
4700        // measured ~123 memsets/token of the decode wall).
4701        let g_len = m.gate_exps.max_expert_bytes(); // 860160 for the uniform 35B gate
4702        let u_len = m.up_exps.max_expert_bytes(); // 860160 for the uniform 35B up
4703        let d_len = m.down_exps.max_expert_bytes(); // 1114112 for the uniform 35B down
4704        let mut scratch_g: Option<CudaSlice<u8>> = None;
4705        let mut scratch_u: Option<CudaSlice<u8>> = None;
4706        let mut scratch_d: Option<CudaSlice<u8>> = None;
4707        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
4708        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
4709
4710        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
4711        // the copy stream before launching the current expert's compute. Pending slots stay invisible
4712        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
4713        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
4714        let page_window = moe_page_prefetch_window();
4715
4716        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
4717        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
4718        for tok in 0..t {
4719            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4720            let w = &w_all[tok * n_used..(tok + 1) * n_used];
4721            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd); // CudaView<f32>
4722            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4723
4724            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
4725            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
4726            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
4727            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
4728            // memcpy, zero admission, so no slot can move under the collected pointers) — any
4729            // miss falls through to the sequential loop below, which admits as before. In steady
4730            // state on a fully-resident rig every token-layer takes the grouped path.
4731            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
4732            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
4733            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
4734            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
4735            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
4736            // per-expert macro-scales the fused kernels don't fold — those fall through too.
4737            let no_macros = m.gate_exps.macros.is_none()
4738                && m.up_exps.macros.is_none()
4739                && m.down_exps.macros.is_none();
4740            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
4741            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
4742            // with pointers computed from the resident slab base + ex*stride instead of
4743            // collected SLRU slot addresses. No cache lock, no residency predicate — the
4744            // slab holds every expert by construction, so this arm never falls through
4745            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
4746            // staging both die). Bit-identity class: pointer provenance only, the same
4747            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
4748            // slab exists it is strictly better (no lock, no miss).
4749            if slab_fused_may_fire {
4750                let (pg, pu, pd) = slab_bases.unwrap();
4751                let mut gp = [0u64; 8];
4752                let mut up = [0u64; 8];
4753                let mut dp = [0u64; 8];
4754                for (j, &ex) in sel.iter().enumerate() {
4755                    let ex = ex as usize;
4756                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
4757                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
4758                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
4759                }
4760                let mut wv = [0f32; 8];
4761                wv[..n_used].copy_from_slice(w);
4762                if tok_q8.is_none() {
4763                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4764                }
4765                let (zq, zd) = tok_q8.as_ref().unwrap();
4766                let act = e.moe_gate_up_silu8_q8(
4767                    crate::WPtr8(gp),
4768                    crate::WPtr8(up),
4769                    zq,
4770                    zd,
4771                    n_embd,
4772                    n_ff_exp,
4773                    n_used,
4774                    m.gate_exps.qtype,
4775                    m.up_exps.qtype,
4776                    m.gate_exps.row_bytes,
4777                    m.up_exps.row_bytes,
4778                )?;
4779                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4780                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4781                e.moe_down8_fma_q8(
4782                    crate::WPtr8(dp),
4783                    crate::F32x8(wv),
4784                    &aq2,
4785                    &ad2,
4786                    &mut dst,
4787                    n_ff_exp,
4788                    n_embd,
4789                    n_used,
4790                    m.down_exps.qtype,
4791                    m.down_exps.row_bytes,
4792                )?;
4793                continue;
4794            }
4795            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
4796                if tok_q8.is_none() {
4797                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4798                }
4799                let (zq, zd) = tok_q8.as_ref().unwrap();
4800                if Self::moe_gdec_token_q8(
4801                    e,
4802                    m,
4803                    il,
4804                    max_block,
4805                    zq,
4806                    zd,
4807                    sel,
4808                    w,
4809                    &mut moe_out,
4810                    tok,
4811                    n_embd,
4812                    n_ff_exp,
4813                    n_used,
4814                )? {
4815                    continue;
4816                }
4817            } else if gdec_may_fire
4818                && cfg.m3.is_none()
4819                && no_macros
4820                && Self::moe_gdec_token(
4821                    e,
4822                    m,
4823                    il,
4824                    max_block,
4825                    &zt,
4826                    sel,
4827                    w,
4828                    &mut moe_out,
4829                    tok,
4830                    n_embd,
4831                    n_ff_exp,
4832                    n_used,
4833                )?
4834            {
4835                continue;
4836            }
4837
4838            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
4839            // slab pair could fire. This token fell through to a sequential axpy loop, which
4840            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
4841            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
4842            // has no fallible predicate), included for the allocation invariant's symmetry.
4843            if gdec_may_fire || slab_fused_may_fire {
4844                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4845                e.memset_zeros_view(&mut row)?;
4846            }
4847
4848            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
4849            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
4850            // stall this path exists to remove, while mixing projections would require another
4851            // activation round-trip. Weight addresses remain valid until this worker is joined at
4852            // the bottom of the token scope.
4853            let mut cpu_mask = vec![false; sel.len()];
4854            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
4855                let gpu_resident = if use_cache {
4856                    e.with_moe_cache(max_block, |cache, _| {
4857                        Ok(sel
4858                            .iter()
4859                            .map(|&expert| {
4860                                let expert = expert as u16;
4861                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
4862                                    .into_iter()
4863                                    .filter(|&projection| {
4864                                        cache
4865                                            .resident(BlockId::new(il, projection, expert))
4866                                            .is_some()
4867                                    })
4868                                    .count()
4869                            })
4870                            .collect::<Vec<_>>())
4871                    })?
4872                } else {
4873                    vec![0; sel.len()]
4874                };
4875                let mut cpu_selected = Vec::new();
4876                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
4877                    if gpu_resident[index] != 3 {
4878                        cpu_mask[index] = true;
4879                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
4880                        let expert = expert as usize;
4881                        cpu_selected.push((expert, route_weight));
4882                    }
4883                }
4884                if crate::cpu_experts::predictor_enabled() {
4885                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
4886                    // from this layer's MoE input and prefetches predicted-and-missing
4887                    // experts into the companion RAM cache. Never blocks this thread.
4888                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4889                    crate::cpu_experts::predictor_submit(il, row);
4890                }
4891                if cpu_selected.is_empty() {
4892                    None
4893                } else {
4894                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
4895                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
4896                        .map_err(std::io::Error::other)?;
4897                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
4898                }
4899            } else {
4900                None
4901            };
4902
4903            let worker_window = worker_disk_prefetch
4904                .then(worker_prefetch_window)
4905                .unwrap_or(0);
4906            for (j, &ex) in sel.iter().enumerate() {
4907                if cpu_mask[j] {
4908                    continue;
4909                }
4910                let ex = ex as usize;
4911                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
4912                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
4913                // fused form) and macro-carrying artifacts — still have their bytes in the
4914                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
4915                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
4916                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
4917                if let Some(d) = slab_local {
4918                    let gl = m.gate_exps.expert_layout(ex);
4919                    let ul = m.up_exps.expert_layout(ex);
4920                    let dl = m.down_exps.expert_layout(ex);
4921                    let (g0, u0, d0) = (
4922                        ex * m.gate_exps.expert_stride,
4923                        ex * m.up_exps.expert_stride,
4924                        ex * m.down_exps.expert_stride,
4925                    );
4926                    let (gate, up) = if moe_q8 {
4927                        if tok_q8.is_none() {
4928                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
4929                        }
4930                        let (zq, zd) = tok_q8.as_ref().unwrap();
4931                        (
4932                            e.qmatvec_expert_q8(
4933                                &d.gate,
4934                                g0..g0 + gl.len,
4935                                zq,
4936                                zd,
4937                                1,
4938                                m.gate_exps.in_f,
4939                                m.gate_exps.out_f,
4940                                gl.qtype,
4941                                gl.row_bytes,
4942                            )?,
4943                            e.qmatvec_expert_q8(
4944                                &d.up,
4945                                u0..u0 + ul.len,
4946                                zq,
4947                                zd,
4948                                1,
4949                                m.up_exps.in_f,
4950                                m.up_exps.out_f,
4951                                ul.qtype,
4952                                ul.row_bytes,
4953                            )?,
4954                        )
4955                    } else {
4956                        (
4957                            e.qmatvec_view(
4958                                &d.gate,
4959                                g0..g0 + gl.len,
4960                                &zt,
4961                                1,
4962                                m.gate_exps.in_f,
4963                                m.gate_exps.out_f,
4964                                gl.qtype,
4965                                gl.row_bytes,
4966                            )?,
4967                            e.qmatvec_view(
4968                                &d.up,
4969                                u0..u0 + ul.len,
4970                                &zt,
4971                                1,
4972                                m.up_exps.in_f,
4973                                m.up_exps.out_f,
4974                                ul.qtype,
4975                                ul.row_bytes,
4976                            )?,
4977                        )
4978                    };
4979                    let mut act = e.uninit(n_ff_exp)?;
4980                    Self::ffn_act_lim(
4981                        e,
4982                        cfg,
4983                        &gate,
4984                        &up,
4985                        m.gate_exps.macro_scale(ex),
4986                        m.up_exps.macro_scale(ex),
4987                        lim_exp,
4988                        &mut act,
4989                        n_ff_exp,
4990                    )?;
4991                    let y = if moe_q8 {
4992                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
4993                        e.qmatvec_expert_q8(
4994                            &d.down,
4995                            d0..d0 + dl.len,
4996                            &aq2,
4997                            &ad2,
4998                            1,
4999                            m.down_exps.in_f,
5000                            m.down_exps.out_f,
5001                            dl.qtype,
5002                            dl.row_bytes,
5003                        )?
5004                    } else {
5005                        let actv = act.slice(0..n_ff_exp);
5006                        e.qmatvec_view(
5007                            &d.down,
5008                            d0..d0 + dl.len,
5009                            &actv,
5010                            1,
5011                            m.down_exps.in_f,
5012                            m.down_exps.out_f,
5013                            dl.qtype,
5014                            dl.row_bytes,
5015                        )?
5016                    };
5017                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5018                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5019                    continue;
5020                }
5021                for next in page_prefetch_positions(j, sel.len(), page_window) {
5022                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
5023                }
5024                let keep = [
5025                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
5026                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
5027                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
5028                ];
5029                if worker_disk_prefetch && worker_window > 0 {
5030                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
5031                        Self::moe_prefetch_disk_expert(
5032                            e,
5033                            il,
5034                            sel[next] as usize,
5035                            m,
5036                            max_block,
5037                            &keep,
5038                        )?;
5039                    }
5040                } else if cache_dispatch
5041                    && !cpu_hybrid
5042                    && moe_prefetch_enabled()
5043                    && j + 1 < sel.len()
5044                {
5045                    let next = sel[j + 1] as usize;
5046                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
5047                }
5048                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
5049                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
5050                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
5051                    // layouts stay on the metadata-aware f32 path.
5052                    if (gate_q8 || up_q8) && tok_q8.is_none() {
5053                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
5054                    }
5055                    let gate = if gate_q8 {
5056                        let (zq, zd) = tok_q8.as_ref().unwrap();
5057                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
5058                    } else {
5059                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
5060                    };
5061                    let up = if up_q8 {
5062                        let (zq, zd) = tok_q8.as_ref().unwrap();
5063                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
5064                    } else {
5065                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
5066                    };
5067                    let mut act = e.uninit(n_ff_exp)?;
5068                    Self::ffn_act_lim(
5069                        e,
5070                        cfg,
5071                        &gate,
5072                        &up,
5073                        m.gate_exps.macro_scale(ex),
5074                        m.up_exps.macro_scale(ex),
5075                        lim_exp,
5076                        &mut act,
5077                        n_ff_exp,
5078                    )?;
5079                    let y = if down_q8 {
5080                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
5081                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
5082                    } else {
5083                        let actv = act.slice(0..n_ff_exp);
5084                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
5085                    };
5086                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5087                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
5088                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5089                } else if cache_dispatch {
5090                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
5091                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
5092                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
5093                    // only difference between HIT and MISS is whether the memcpy_htod ran.
5094                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
5095                    let up = Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?;
5096                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5097                    Self::ffn_act_lim(
5098                        e,
5099                        cfg,
5100                        &gate,
5101                        &up,
5102                        m.gate_exps.macro_scale(ex),
5103                        m.up_exps.macro_scale(ex),
5104                        lim_exp,
5105                        &mut act,
5106                        n_ff_exp,
5107                    )?;
5108                    let actv = act.slice(0..n_ff_exp);
5109                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
5110                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5111                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
5112                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5113                } else if cache_frozen {
5114                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
5115                    // first prime. Reuse every fixed resident projection directly and stage only a
5116                    // true miss through the ordinary scratch slot. This preserves the established
5117                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
5118                    let gate = Self::moe_frozen_gemm(
5119                        e,
5120                        il,
5121                        PROJ_GATE,
5122                        ex,
5123                        m,
5124                        max_block,
5125                        &zt,
5126                        &mut scratch_g,
5127                        g_len,
5128                    )?;
5129                    let up = Self::moe_frozen_gemm(
5130                        e,
5131                        il,
5132                        PROJ_UP,
5133                        ex,
5134                        m,
5135                        max_block,
5136                        &zt,
5137                        &mut scratch_u,
5138                        u_len,
5139                    )?;
5140                    let mut act = e.uninit(n_ff_exp)?;
5141                    Self::ffn_act_lim(
5142                        e,
5143                        cfg,
5144                        &gate,
5145                        &up,
5146                        m.gate_exps.macro_scale(ex),
5147                        m.up_exps.macro_scale(ex),
5148                        lim_exp,
5149                        &mut act,
5150                        n_ff_exp,
5151                    )?;
5152                    let actv = act.slice(0..n_ff_exp);
5153                    let y = Self::moe_frozen_gemm(
5154                        e,
5155                        il,
5156                        PROJ_DOWN,
5157                        ex,
5158                        m,
5159                        max_block,
5160                        &actv,
5161                        &mut scratch_d,
5162                        d_len,
5163                    )?;
5164                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5165                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5166                } else {
5167                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
5168                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
5169                    // fully overwrites the byte range the GEMM reads).
5170                    if scratch_g.is_none() {
5171                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
5172                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
5173                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
5174                    }
5175                    let (sg, su, sd) = (
5176                        scratch_g.as_mut().unwrap(),
5177                        scratch_u.as_mut().unwrap(),
5178                        scratch_d.as_mut().unwrap(),
5179                    );
5180                    let gl = m.gate_exps.expert_layout(ex);
5181                    let ul = m.up_exps.expert_layout(ex);
5182                    let dl = m.down_exps.expert_layout(ex);
5183                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5184                    let gate = e.qmatvec_view(
5185                        sg,
5186                        0..gl.len,
5187                        &zt,
5188                        1,
5189                        m.gate_exps.in_f,
5190                        m.gate_exps.out_f,
5191                        gl.qtype,
5192                        gl.row_bytes,
5193                    )?;
5194
5195                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5196                    let up = e.qmatvec_view(
5197                        su,
5198                        0..ul.len,
5199                        &zt,
5200                        1,
5201                        m.up_exps.in_f,
5202                        m.up_exps.out_f,
5203                        ul.qtype,
5204                        ul.row_bytes,
5205                    )?;
5206
5207                    let mut act = e.uninit(n_ff_exp)?; // activation fully overwrites
5208                    Self::ffn_act_lim(
5209                        e,
5210                        cfg,
5211                        &gate,
5212                        &up,
5213                        m.gate_exps.macro_scale(ex),
5214                        m.up_exps.macro_scale(ex),
5215                        lim_exp,
5216                        &mut act,
5217                        n_ff_exp,
5218                    )?;
5219
5220                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5221                    let actv = act.slice(0..n_ff_exp);
5222                    let y = e.qmatvec_view(
5223                        sd,
5224                        0..dl.len,
5225                        &actv,
5226                        1,
5227                        m.down_exps.in_f,
5228                        m.down_exps.out_f,
5229                        dl.qtype,
5230                        dl.row_bytes,
5231                    )?;
5232
5233                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5234                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
5235                }
5236            }
5237            if let Some(worker) = cpu_worker {
5238                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
5239                let cpu_output = e.htod(&cpu_output)?;
5240                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5241                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
5242            }
5243            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
5244                for (j, &ex) in sel.iter().enumerate() {
5245                    if cpu_mask[j] {
5246                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
5247                    }
5248                }
5249            }
5250        }
5251
5252        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
5253        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
5254        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5255        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5256        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5257            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5258        {
5259            let n_ff_sh = gate_shexp.out_features(); // 512
5260            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
5261            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
5262            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
5263            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
5264            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
5265            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
5266            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
5267            let verify_t = t > 1 && t < PRIME_MIN_T;
5268            let (sg_gate, sg_up) = if t == 1 {
5269                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5270                    Some(pair) => pair,
5271                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5272                }
5273            } else if verify_t {
5274                (
5275                    e.matmul_decode_exact(gate_shexp, z, t)?,
5276                    e.matmul_decode_exact(up_shexp, z, t)?,
5277                )
5278            } else {
5279                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?) // [T, 512] each
5280            };
5281            let mut sa = e.uninit(t * n_ff_sh)?; // activation fully overwrites
5282            Self::ffn_act_lim(
5283                e,
5284                cfg,
5285                &sg_gate,
5286                &sg_up,
5287                1.0,
5288                1.0,
5289                lim_shexp,
5290                &mut sa,
5291                t * n_ff_sh,
5292            )?;
5293            let sh = if verify_t {
5294                e.matmul_decode_exact(down_shexp, &sa, t)?
5295            } else {
5296                e.matmul(down_shexp, &sa, t)?
5297            }; // [T, n_embd]
5298
5299            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
5300            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
5301            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
5302            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
5303            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
5304            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
5305            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
5306            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
5307            // expert's contribution into every token's residual, so under cross-request
5308            // concat prefill a session's hidden state depended on its co-arrivals' token
5309            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
5310            let g = match &m.gate_inp_shexp {
5311                Some(gate_inp_shexp) => {
5312                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5313                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5314                    } else {
5315                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5316                        let mut g = e.uninit(t)?; // sigmoid fully overwrites
5317                        e.sigmoid(&gs, &mut g, t)?;
5318                        g
5319                    }
5320                }
5321                None => e.htod(&vec![1.0f32; t])?,
5322            };
5323            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
5324            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5325        }
5326
5327        Ok(moe_out)
5328    }
5329
5330    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
5331    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
5332    pub fn stage1_h2d_per_token(&self) -> u64 {
5333        use crate::hybrid::Ffn;
5334        let n_used = self
5335            .cfg
5336            .moe
5337            .as_ref()
5338            .map(|m| m.expert_used_count as u64)
5339            .unwrap_or(0);
5340        let mut bytes = 0u64;
5341        for l in self.layers.iter() {
5342            if let Ffn::Moe(m) = &l.ffn {
5343                bytes += n_used
5344                    * (m.gate_exps.max_expert_bytes()
5345                        + m.up_exps.max_expert_bytes()
5346                        + m.down_exps.max_expert_bytes()) as u64;
5347            }
5348        }
5349        bytes
5350    }
5351
5352    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
5353    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
5354    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
5355    pub(crate) fn max_moe_block(&self) -> usize {
5356        use crate::hybrid::Ffn;
5357        let mut mx = 0usize;
5358        let mut scan = |ffn: &Ffn| {
5359            if let Ffn::Moe(m) = ffn {
5360                mx = mx
5361                    .max(m.gate_exps.max_expert_bytes())
5362                    .max(m.up_exps.max_expert_bytes())
5363                    .max(m.down_exps.max_expert_bytes());
5364            }
5365        };
5366        for l in self.layers.iter() {
5367            scan(&l.ffn);
5368        }
5369        if let Some(mtp) = self.mtp.as_ref() {
5370            scan(&mtp.ffn);
5371        }
5372        mx
5373    }
5374
5375    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
5376    /// but have no bytes and therefore consume no residency slot.
5377    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
5378        use crate::hybrid::Ffn;
5379        let mut sizes = Vec::new();
5380        let mut scan = |ffn: &Ffn| {
5381            let Ffn::Moe(m) = ffn else { return };
5382            for ex in 0..m.gate_exps.n_expert {
5383                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
5384                    continue;
5385                }
5386                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
5387                    let len = exps.expert_layout(ex).len;
5388                    if len > 0 {
5389                        sizes.push(len);
5390                    }
5391                }
5392            }
5393        };
5394        for layer in &self.layers {
5395            scan(&layer.ffn);
5396        }
5397        if let Some(mtp) = &self.mtp {
5398            scan(&mtp.ffn);
5399        }
5400        sizes
5401    }
5402
5403    /// Persist the frozen residency set so a later process can restage it directly and skip
5404    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
5405    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
5406    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
5407    /// post-freeze argmax gate still validates the serving assignment.
5408    pub fn save_cpu_expert_residency_profile(
5409        &self,
5410        e: &Engine,
5411        path: &std::path::Path,
5412    ) -> Result<(), Box<dyn std::error::Error>> {
5413        let Some(ids) = e.export_moe_residency() else {
5414            return Err("no MoE residency cache to persist".into());
5415        };
5416        let mut body = format!(
5417            "memra-freeze-profile v1 max_block={} blocks={}\n",
5418            self.max_moe_block(),
5419            ids.len()
5420        );
5421        for (layer, proj, ex) in &ids {
5422            body.push_str(&format!("{layer} {proj} {ex}\n"));
5423        }
5424        let tmp = path.with_extension("tmp");
5425        std::fs::write(&tmp, body)?;
5426        std::fs::rename(&tmp, path)?;
5427        println!(
5428            "[moe-cache] freeze profile saved: {} blocks -> {}",
5429            ids.len(),
5430            path.display()
5431        );
5432        Ok(())
5433    }
5434
5435    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
5436    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
5437    /// missing or its header does not match this model's slot geometry.
5438    pub fn restore_cpu_expert_residency_profile(
5439        &self,
5440        e: &Engine,
5441        path: &std::path::Path,
5442    ) -> Result<bool, Box<dyn std::error::Error>> {
5443        use crate::hybrid::Ffn;
5444        use crate::moe_cache::BlockId;
5445        let Ok(content) = std::fs::read_to_string(path) else {
5446            return Ok(false);
5447        };
5448        let mut lines = content.lines();
5449        let Some(header) = lines.next() else {
5450            return Ok(false);
5451        };
5452        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
5453        if !header.starts_with(&expected) {
5454            println!(
5455                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
5456                path.display()
5457            );
5458            return Ok(false);
5459        }
5460        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
5461            std::collections::HashMap::new();
5462        for line in lines {
5463            let mut fields = line.split_whitespace();
5464            let (Some(layer), Some(proj), Some(ex)) = (fields.next(), fields.next(), fields.next())
5465            else {
5466                continue;
5467            };
5468            let (Ok(layer), Ok(proj), Ok(ex)) =
5469                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
5470            else {
5471                continue;
5472            };
5473            by_layer
5474                .entry(layer)
5475                .or_default()
5476                .push(BlockId::new(layer, proj, ex));
5477        }
5478        let requested: usize = by_layer.values().map(Vec::len).sum();
5479        if requested == 0 {
5480            return Ok(false);
5481        }
5482        let max_block = self.max_moe_block();
5483        let mut restaged = 0usize;
5484        let mut stage_layer =
5485            |layer_index: u16, ffn: &Ffn| -> Result<(), Box<dyn std::error::Error>> {
5486                let Ffn::Moe(m) = ffn else { return Ok(()) };
5487                let Some(ids) = by_layer.get(&layer_index) else {
5488                    return Ok(());
5489                };
5490                e.with_moe_cache(max_block, |cache, eng| {
5491                    for id in ids {
5492                        if cache.restage_block(*id, m, eng)? {
5493                            restaged += 1;
5494                        }
5495                    }
5496                    Ok(())
5497                })
5498            };
5499        for (index, layer) in self.layers.iter().enumerate() {
5500            stage_layer(index as u16, &layer.ffn)?;
5501        }
5502        if let Some(mtp) = self.mtp.as_ref() {
5503            stage_layer(u16::MAX, &mtp.ffn)?;
5504        }
5505        e.freeze_moe_cache();
5506        println!(
5507            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
5508            path.display()
5509        );
5510        Ok(true)
5511    }
5512
5513    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
5514    pub fn freeze_cpu_expert_residency(
5515        &self,
5516        e: &Engine,
5517    ) -> Result<(), Box<dyn std::error::Error>> {
5518        e.freeze_moe_cache();
5519        Ok(())
5520    }
5521
5522    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
5523    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
5524    /// the model's activation exactly.
5525    ///
5526    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
5527    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
5528    /// form for anything that can land on a clamped layer.
5529    pub fn ffn_act(
5530        e: &Engine,
5531        cfg: &ModelConfig,
5532        gate: &CudaSlice<f32>,
5533        up: &CudaSlice<f32>,
5534        act: &mut CudaSlice<f32>,
5535        n: usize,
5536    ) -> Result<(), Box<dyn std::error::Error>> {
5537        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
5538    }
5539
5540    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
5541    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
5542    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
5543    #[allow(clippy::too_many_arguments)]
5544    pub(crate) fn ffn_act_scaled(
5545        e: &Engine,
5546        cfg: &ModelConfig,
5547        gate: &CudaSlice<f32>,
5548        up: &CudaSlice<f32>,
5549        gs: f32,
5550        us: f32,
5551        act: &mut CudaSlice<f32>,
5552        n: usize,
5553    ) -> Result<(), Box<dyn std::error::Error>> {
5554        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
5555    }
5556
5557    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
5558    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
5559    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
5560    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
5561    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
5562    ///                 arrays are SEPARATE and a layer can have one without the other.
5563    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
5564    /// already known live.
5565    #[allow(clippy::too_many_arguments)]
5566    pub(crate) fn ffn_act_lim(
5567        e: &Engine,
5568        cfg: &ModelConfig,
5569        gate: &CudaSlice<f32>,
5570        up: &CudaSlice<f32>,
5571        gs: f32,
5572        us: f32,
5573        limit: Option<f32>,
5574        act: &mut CudaSlice<f32>,
5575        n: usize,
5576    ) -> Result<(), Box<dyn std::error::Error>> {
5577        if let Some(m3) = cfg.m3.as_ref() {
5578            debug_assert!(
5579                limit.is_none(),
5580                "m3 swigluoai and step35 clamp are different archs"
5581            );
5582            return e.swigluoai_mul_scaled(
5583                gate,
5584                up,
5585                gs,
5586                us,
5587                m3.swiglu_alpha,
5588                m3.swiglu_limit,
5589                act,
5590                n,
5591            );
5592        }
5593        if let Some(l) = limit {
5594            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
5595        }
5596        if gs == 1.0 && us == 1.0 {
5597            return e.silu_mul(gate, up, act, n);
5598        }
5599        e.silu_mul_scaled(gate, up, gs, us, act, n)
5600    }
5601
5602    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
5603    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
5604    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
5605    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
5606    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
5607    fn moe_route(
5608        e: &Engine,
5609        logits: &CudaSlice<f32>,
5610        t: usize,
5611        n_expert: usize,
5612        n_used: usize,
5613    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5614        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
5615    }
5616
5617    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
5618    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
5619    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
5620    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
5621    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
5622    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
5623    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
5624    #[allow(clippy::too_many_arguments)]
5625    fn moe_route_sigmoid_cfg(
5626        e: &Engine,
5627        logits: &CudaSlice<f32>,
5628        t: usize,
5629        n_expert: usize,
5630        n_used: usize,
5631        m: &MoeWeights,
5632        (sf, route_norm): (f32, bool),
5633    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5634        if sigmoid_router_enabled() {
5635            return e.moe_router_sigmoid_topk_host(
5636                logits,
5637                t,
5638                n_expert,
5639                n_used,
5640                m.active_count(),
5641                &m.exp_probs_b_dev,
5642                &m.active_experts_dev,
5643                sf,
5644                route_norm,
5645            );
5646        }
5647        let lg = e.dtoh(logits)?;
5648        Self::moe_route_sigmoid_host(
5649            &lg,
5650            t,
5651            n_expert,
5652            n_used,
5653            m.exp_probs_b.as_deref(),
5654            sf,
5655            route_norm,
5656            m.active_experts.as_deref(),
5657        )
5658    }
5659
5660    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
5661    /// the existing softmax device kernel has no mask input.
5662    fn moe_route_cfg(
5663        e: &Engine,
5664        logits: &CudaSlice<f32>,
5665        t: usize,
5666        n_expert: usize,
5667        n_used: usize,
5668        active: Option<&[bool]>,
5669    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5670        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
5671        // rollback) via the single-sync pinned readback — softmax arch only.
5672        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
5673            return e.moe_router_topk_host(logits, t, n_expert, n_used);
5674        }
5675        // Host oracle (the §D bit-identity reference).
5676        let lg = e.dtoh(logits)?; // [T*n_expert] host
5677        let mut sel = vec![0u32; t * n_used];
5678        let mut w_out = vec![0f32; t * n_used];
5679        for tok in 0..t {
5680            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5681            // softmax over ALL n_expert (stable: subtract max)
5682            let maxl = row
5683                .iter()
5684                .enumerate()
5685                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
5686                .map(|(_, &x)| x)
5687                .fold(f32::NEG_INFINITY, f32::max);
5688            let mut probs = vec![0f32; n_expert];
5689            let mut den = 0f32;
5690            for i in 0..n_expert {
5691                if active.is_some_and(|mask| !mask[i]) {
5692                    continue;
5693                }
5694                let x = (row[i] - maxl).exp();
5695                probs[i] = x;
5696                den += x;
5697            }
5698            for p in probs.iter_mut() {
5699                *p /= den;
5700            }
5701            // stable DESC sort: prob DESC, ascending-index tiebreak.
5702            let mut idx: Vec<usize> = (0..n_expert)
5703                .filter(|&i| active.is_none_or(|mask| mask[i]))
5704                .collect();
5705            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
5706            let sl = &idx[..n_used];
5707            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
5708            let mut ws: f32 = wv.iter().sum();
5709            ws = ws.max(6.103515625e-5_f32); // F16 smallest normal, clamp BEFORE divide
5710            for x in wv.iter_mut() {
5711                *x /= ws;
5712            }
5713            for j in 0..n_used {
5714                sel[tok * n_used + j] = sl[j] as u32;
5715                w_out[tok * n_used + j] = wv[j];
5716            }
5717        }
5718        Ok((sel, w_out))
5719    }
5720
5721    #[allow(clippy::too_many_arguments)]
5722    fn moe_route_sigmoid_with_input(
5723        e: &Engine,
5724        logits: &CudaSlice<f32>,
5725        input: &CudaSlice<f32>,
5726        t: usize,
5727        n_expert: usize,
5728        n_used: usize,
5729        bias: Option<&[f32]>,
5730        (sf, route_norm): (f32, bool),
5731        active: Option<&[bool]>,
5732    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
5733        let (lg, input) = e.dtoh_pair(logits, input)?;
5734        let (sel, w) =
5735            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
5736        Ok((sel, w, input))
5737    }
5738
5739    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
5740    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
5741    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
5742    /// active mask, prebuilt projection descriptors) so no model reference escapes.
5743    pub fn start_moe_prefetch_predictor(
5744        &self,
5745        e: &Engine,
5746        cfg: &ModelConfig,
5747    ) -> Result<(), Box<dyn std::error::Error>> {
5748        use crate::hybrid::Ffn;
5749        let Some(sig) = cfg.sigmoid_router() else {
5750            return Err("prefetch predictor requires a sigmoid-router arch".into());
5751        };
5752        let resident: std::collections::HashSet<(u16, u8, u16)> = e
5753            .export_moe_residency()
5754            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
5755            .into_iter()
5756            .collect();
5757        let mut layers = Vec::new();
5758        for (index, layer) in self.layers.iter().enumerate() {
5759            let Ffn::Moe(m) = &layer.ffn else { continue };
5760            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else {
5761                continue;
5762            };
5763            let router = e.dtoh(data)?;
5764            let n_expert = m.gate_exps.n_expert;
5765            let n_embd = m.gate_exps.in_f;
5766            if router.len() != n_embd * n_expert {
5767                continue;
5768            }
5769            let build = |exps: &crate::model::HostExps| {
5770                (0..n_expert)
5771                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
5772                    .collect::<Vec<_>>()
5773            };
5774            layers.push((
5775                index as u16,
5776                crate::cpu_experts::PredictLayerInit {
5777                    router,
5778                    bias: m.exp_probs_b.clone(),
5779                    active: m.active_experts.clone(),
5780                    n_embd,
5781                    n_used: cfg
5782                        .moe
5783                        .as_ref()
5784                        .map(|moe| moe.expert_used_count as usize)
5785                        .ok_or("prefetch predictor requires MoE config")?,
5786                    sig,
5787                    weights_n_expert: n_expert,
5788                    gate: build(&m.gate_exps),
5789                    up: build(&m.up_exps),
5790                    down: build(&m.down_exps),
5791                },
5792            ));
5793        }
5794        crate::cpu_experts::start_prefetch_predictor(layers, resident).map_err(|error| error.into())
5795    }
5796
5797    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
5798    /// selection math to the rollback runtime, applied to host-computed logits.
5799    #[allow(clippy::too_many_arguments)]
5800    pub fn moe_route_sigmoid_host_public(
5801        logits: &[f32],
5802        t: usize,
5803        n_expert: usize,
5804        n_used: usize,
5805        bias: Option<&[f32]>,
5806        sf: f32,
5807        route_norm: bool,
5808        active: Option<&[bool]>,
5809    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5810        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
5811    }
5812
5813    #[allow(clippy::too_many_arguments)]
5814    fn moe_route_sigmoid_host(
5815        lg: &[f32],
5816        t: usize,
5817        n_expert: usize,
5818        n_used: usize,
5819        bias: Option<&[f32]>,
5820        sf: f32,
5821        route_norm: bool,
5822        active: Option<&[bool]>,
5823    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5824        let active_count = active
5825            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
5826            .unwrap_or(n_expert);
5827        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5828        if lg.len() != t * n_expert {
5829            return Err(format!(
5830                "sigmoid router logits length mismatch: got {}, expected {}",
5831                lg.len(),
5832                t * n_expert,
5833            )
5834            .into());
5835        }
5836        let mut sel = vec![0u32; t * n_used];
5837        let mut w_out = vec![0f32; t * n_used];
5838        for tok in 0..t {
5839            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
5840            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
5841            // selection score = sigmoid + bias; weight = plain sigmoid.
5842            let selsc: Vec<f32> = match bias {
5843                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
5844                None => scores.clone(),
5845            };
5846            let mut idx: Vec<usize> = (0..n_expert)
5847                .filter(|&i| active.is_none_or(|mask| mask[i]))
5848                .collect();
5849            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
5850            let sl = &idx[..n_used];
5851            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
5852            if route_norm {
5853                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
5854                for x in wv.iter_mut() {
5855                    *x = *x / ws * sf;
5856                }
5857            } else {
5858                for x in wv.iter_mut() {
5859                    *x *= sf;
5860                }
5861            }
5862            for j in 0..n_used {
5863                sel[tok * n_used + j] = sl[j] as u32;
5864                w_out[tok * n_used + j] = wv[j];
5865            }
5866        }
5867        Ok((sel, w_out))
5868    }
5869
5870    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
5871    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
5872    /// macro-scaled experts, and observation modes are denied by the caller.
5873    #[allow(clippy::too_many_arguments)]
5874    fn moe_ffn_sigmoid_dev(
5875        e: &Engine,
5876        m: &MoeWeights,
5877        z: &CudaSlice<f32>,
5878        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
5879        logits: &CudaSlice<f32>,
5880        t: usize,
5881        cfg: &ModelConfig,
5882        il: u16,
5883        (scaling_factor, route_norm): (f32, bool),
5884    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5885        let moe = cfg.moe.as_ref().unwrap();
5886        let n_embd = cfg.n_embd as usize;
5887        let n_expert = moe.expert_count as usize;
5888        let n_used = moe.expert_used_count as usize;
5889        let n_ff_exp = moe.expert_ff_length as usize;
5890        let dev = m.dev_exps.as_ref().unwrap();
5891        debug_assert!(cfg.step35.is_some());
5892        debug_assert_eq!(dev.dev, e.ctx().ordinal());
5893        debug_assert!(m.has_uniform_expert_layout());
5894        debug_assert!(!m.has_macros);
5895
5896        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
5897            logits,
5898            t,
5899            n_expert,
5900            n_used,
5901            m.active_count(),
5902            &m.exp_probs_b_dev,
5903            &m.active_experts_dev,
5904            scaling_factor,
5905            route_norm,
5906        )?;
5907        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
5908        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
5909            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5910            (combined, combined)
5911        } else {
5912            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5913        };
5914        let (zq, zd) = match (t, zq8) {
5915            (1, Some((q, d))) => (q.clone(), d.clone()),
5916            _ => e.quantize_q8_1(z, t, n_embd)?,
5917        };
5918        let n_pairs = t * n_used;
5919        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
5920            // The final Step layers retain the established separate gate/up -> clamp -> down
5921            // arithmetic. Pair rows are derived from token position; selected expert ids and
5922            // routing weights remain the device router's buffers throughout.
5923            let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5924            let pair_tok_d = e.htod_i32(&pair_tok)?;
5925            let gate = e.moe_pairs_matvec_q8(
5926                &dev.ptr_row,
5927                0,
5928                &pair_tok_d,
5929                &sel_d,
5930                &zq,
5931                &zd,
5932                n_embd,
5933                n_ff_exp,
5934                n_expert,
5935                n_pairs,
5936                m.gate_exps.qtype,
5937                gate_row_bytes,
5938            )?;
5939            let up = e.moe_pairs_matvec_q8(
5940                &dev.ptr_row,
5941                1,
5942                &pair_tok_d,
5943                &sel_d,
5944                &zq,
5945                &zd,
5946                n_embd,
5947                n_ff_exp,
5948                n_expert,
5949                n_pairs,
5950                m.up_exps.qtype,
5951                up_row_bytes,
5952            )?;
5953            let mut act = e.uninit(n_pairs * n_ff_exp)?;
5954            Self::ffn_act_lim(
5955                e,
5956                cfg,
5957                &gate,
5958                &up,
5959                1.0,
5960                1.0,
5961                cfg.clamp_exp_at(il as u32),
5962                &mut act,
5963                n_pairs * n_ff_exp,
5964            )?;
5965            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5966            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5967            let pair_self_d = e.htod_i32(&pair_self)?;
5968            let down = e.moe_pairs_matvec_q8(
5969                &dev.ptr_row,
5970                2,
5971                &pair_self_d,
5972                &sel_d,
5973                &aq2,
5974                &ad2,
5975                n_ff_exp,
5976                n_embd,
5977                n_expert,
5978                n_pairs,
5979                m.down_exps.qtype,
5980                m.down_exps.row_bytes,
5981            )?;
5982            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5983            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5984            let tok_off_d = e.htod_i32(&tok_off)?;
5985            let tok_ids_d = e.htod_i32(&tok_ids)?;
5986            let mut output = e.uninit(t * n_embd)?;
5987            e.moe_pairs_scatter(&down, &w_d, &tok_off_d, &tok_ids_d, &mut output, t, n_embd)?;
5988            output
5989        } else {
5990            let act = e.moe_gate_up_silu8_dev_q8_rows(
5991                &dev.ptr_row,
5992                &sel_d,
5993                &zq,
5994                &zd,
5995                t,
5996                n_embd,
5997                n_ff_exp,
5998                n_used,
5999                n_expert,
6000                m.gate_exps.qtype,
6001                m.up_exps.qtype,
6002                gate_row_bytes,
6003                up_row_bytes,
6004                &m.dev_macros,
6005            )?;
6006            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6007            let mut output = e.uninit(t * n_embd)?;
6008            e.moe_down8_fma_dev_q8_rows_g(
6009                &dev.ptr_row,
6010                &sel_d,
6011                &w_d,
6012                &aq2,
6013                &ad2,
6014                &mut output,
6015                t,
6016                n_ff_exp,
6017                n_embd,
6018                n_used,
6019                n_expert,
6020                m.down_exps.qtype,
6021                m.down_exps.row_bytes,
6022            )?;
6023            output
6024        };
6025
6026        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
6027            eprintln!(
6028                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
6029                cfg.clamp_exp_at(il as u32).is_some(),
6030                dev.gu_il,
6031            );
6032        }
6033        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6034        Ok(moe_out)
6035    }
6036
6037    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
6038    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
6039    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
6040    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
6041    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
6042    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
6043    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
6044    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
6045    fn moe_ffn_pairs(
6046        e: &Engine,
6047        m: &MoeWeights,
6048        z: &CudaSlice<f32>,
6049        logits: &CudaSlice<f32>,
6050        t: usize,
6051        cfg: &ModelConfig,
6052    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6053        let moe = cfg.moe.as_ref().unwrap();
6054        let n_embd = cfg.n_embd as usize;
6055        let n_expert = moe.expert_count as usize;
6056        let n_used = moe.expert_used_count as usize;
6057        let n_ff_exp = moe.expert_ff_length as usize;
6058        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
6059        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
6060        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
6061        // that forgets the gate fails loudly in debug instead of returning wrong logits.
6062        debug_assert!(
6063            !cfg.swiglu_clamped_anywhere(),
6064            "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU"
6065        );
6066        let dev = m.dev_exps.as_ref().unwrap();
6067        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
6068        let (rbg_d, rbu_d) = if dev.gu_il {
6069            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6070            (sxx, sxx)
6071        } else {
6072            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6073        };
6074
6075        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
6076        let n_pairs = t * n_used;
6077        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
6078        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
6079        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
6080        let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
6081        let pair_w: Vec<f32> = w_all.clone();
6082        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
6083        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
6084        let pt = e.htod_i32(&pair_tok)?;
6085        let px = e.htod_i32(&pair_ex)?;
6086        let pw = e.htod(&pair_w)?;
6087        let toff = e.htod_i32(&tok_off)?;
6088        let tids = e.htod_i32(&tok_ids)?;
6089
6090        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
6091        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
6092        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
6093        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
6094        for p in 0..n_pairs {
6095            by_ex[pair_ex[p] as usize].push(p as i32);
6096        }
6097        let mut ex_ids: Vec<i32> = Vec::new();
6098        let mut ex_off: Vec<i32> = vec![0];
6099        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
6100        for (ex, list) in by_ex.iter().enumerate() {
6101            if list.is_empty() {
6102                continue;
6103            }
6104            ex_ids.push(ex as i32);
6105            ex_pairs.extend_from_slice(list);
6106            ex_off.push(ex_pairs.len() as i32);
6107        }
6108        let n_active = ex_ids.len();
6109        let exi = e.htod_i32(&ex_ids)?;
6110        let exo = e.htod_i32(&ex_off)?;
6111        let exp_d = e.htod_i32(&ex_pairs)?;
6112        let _ = &px; // pair-major twin keeps it; em path uses CSR
6113
6114        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
6115        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
6116        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
6117        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
6118        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
6119        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
6120        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
6121        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
6122        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
6123        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
6124        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
6125        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
6126        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
6127        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
6128        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
6129        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
6130        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
6131        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
6132        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6133        let mma_t = *MMA_T.get_or_init(|| {
6134            std::env::var("MEMRA_MOE_MMA_T")
6135                .ok()
6136                .and_then(|v| v.parse().ok())
6137                .unwrap_or(16)
6138        });
6139        let use_mma = std::env::var("MEMRA_MOE_MMA")
6140            .map(|v| v != "0")
6141            .unwrap_or(true)
6142            && t >= mma_t
6143            && q8_expert_dec_supported(m.gate_exps.qtype)
6144            && q8_expert_dec_supported(m.up_exps.qtype)
6145            && q8_expert_dec_supported(m.down_exps.qtype)
6146            && n_embd % 256 == 0
6147            && n_ff_exp % 256 == 0;
6148        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
6149        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
6150        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
6151        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
6152        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
6153        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
6154        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
6155        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
6156        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
6157        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
6158        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
6159        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
6160        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
6161        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
6162        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
6163        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
6164            && q8_expert_dec_supported(m.up_exps.qtype)
6165            && q8_expert_dec_supported(m.down_exps.qtype)
6166            && n_embd % 256 == 0
6167            && n_ff_exp % 256 == 0;
6168        let f16g_mode = crate::moe_f16g_mode();
6169        let f16g = f16g_mode != 0
6170            && t >= mma_t
6171            && (f16g_mode != 3 || !mma_capable)
6172            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
6173            && f16g_proj_ok(m.up_exps.qtype, n_embd)
6174            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
6175        if use_mma || f16g {
6176            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
6177            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
6178            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
6179            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
6180            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
6181            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
6182            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
6183            let y_down = if f16g {
6184                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
6185                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
6186                // permute at the very end back to pair-id order for the scatter.
6187                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
6188                let csr_tok_d = e.htod_i32(&csr_tok)?;
6189                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
6190                let g_csr = e.moe_f16_grouped(
6191                    &dev.ptr_row,
6192                    0,
6193                    n_expert,
6194                    &exi,
6195                    &ex_off,
6196                    &exo,
6197                    &z_f16,
6198                    &z_s,
6199                    n_embd,
6200                    n_ff_exp,
6201                    n_active,
6202                    n_pairs,
6203                    m.gate_exps.qtype,
6204                    rbg_d,
6205                )?;
6206                let u_csr = e.moe_f16_grouped(
6207                    &dev.ptr_row,
6208                    1,
6209                    n_expert,
6210                    &exi,
6211                    &ex_off,
6212                    &exo,
6213                    &z_f16,
6214                    &z_s,
6215                    n_embd,
6216                    n_ff_exp,
6217                    n_active,
6218                    n_pairs,
6219                    m.up_exps.qtype,
6220                    rbu_d,
6221                )?;
6222                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
6223                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
6224                let d_csr = e.moe_f16_grouped(
6225                    &dev.ptr_row,
6226                    2,
6227                    n_expert,
6228                    &exi,
6229                    &ex_off,
6230                    &exo,
6231                    &a_f16,
6232                    &a_s,
6233                    n_ff_exp,
6234                    n_embd,
6235                    n_active,
6236                    n_pairs,
6237                    m.down_exps.qtype,
6238                    m.down_exps.row_bytes,
6239                )?;
6240                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
6241            } else {
6242                // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
6243                let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
6244                let gate = e.mmq_iq_experts(
6245                    &dev.ptr_row,
6246                    0,
6247                    n_expert,
6248                    &exi,
6249                    &exo,
6250                    &exp_d,
6251                    &pt,
6252                    &z_scr,
6253                    n_embd,
6254                    n_ff_exp,
6255                    n_active,
6256                    n_pairs,
6257                    t,
6258                    m.gate_exps.qtype,
6259                    rbg_d,
6260                )?;
6261                let up = e.mmq_iq_experts(
6262                    &dev.ptr_row,
6263                    1,
6264                    n_expert,
6265                    &exi,
6266                    &exo,
6267                    &exp_d,
6268                    &pt,
6269                    &z_scr,
6270                    n_embd,
6271                    n_ff_exp,
6272                    n_active,
6273                    n_pairs,
6274                    t,
6275                    m.up_exps.qtype,
6276                    rbu_d,
6277                )?;
6278                // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
6279                // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
6280                // registers and writes ONLY the quantized scratch — the two-pass chain
6281                // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
6282                // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
6283                let a_scr = if crate::moe_fuse_actq_on() {
6284                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
6285                } else {
6286                    let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6287                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
6288                };
6289                let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6290                let pself = e.htod_i32(&pair_self)?;
6291                e.mmq_iq_experts(
6292                    &dev.ptr_row,
6293                    2,
6294                    n_expert,
6295                    &exi,
6296                    &exo,
6297                    &exp_d,
6298                    &pself,
6299                    &a_scr,
6300                    n_ff_exp,
6301                    n_embd,
6302                    n_active,
6303                    n_pairs,
6304                    n_pairs,
6305                    m.down_exps.qtype,
6306                    m.down_exps.row_bytes,
6307                )?
6308            };
6309            let mut moe_out = e.uninit(t * n_embd)?;
6310            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6311            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6312                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6313            {
6314                let n_ff_sh = gate_shexp.out_features();
6315                let sg_gate = e.matmul(gate_shexp, z, t)?;
6316                let sg_up = e.matmul(up_shexp, z, t)?;
6317                let mut sa = e.uninit(t * n_ff_sh)?;
6318                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6319                let sh = e.matmul(down_shexp, &sa, t)?;
6320                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6321                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
6322                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
6323                // i.e. the one real prefill actually takes on a resident-expert MoE model,
6324                // so the concat-prime isolation fix has to land here as well.
6325                let g = match &m.gate_inp_shexp {
6326                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6327                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6328                    }
6329                    Some(gate_inp_shexp) => {
6330                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6331                        let mut g = e.uninit(t)?;
6332                        e.sigmoid(&gs, &mut g, t)?;
6333                        g
6334                    }
6335                    None => e.htod(&vec![1.0f32; t])?,
6336                };
6337                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6338            }
6339            return Ok(moe_out);
6340        }
6341
6342        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
6343        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
6344        let dec = std::env::var("MEMRA_MOE_DEC")
6345            .map(|v| v != "0")
6346            .unwrap_or(true);
6347        let matvec = |proj,
6348                      exi: &_,
6349                      exo: &_,
6350                      exp_d: &_,
6351                      pt: &_,
6352                      aq: &_,
6353                      ad: &_,
6354                      inf,
6355                      outf,
6356                      qtype,
6357                      rb|
6358         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6359            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
6360            let dec = dec && q8_expert_dec_supported(qtype);
6361            if dec {
6362                e.moe_pairs_matvec_q8_dec(
6363                    &dev.ptr_row,
6364                    proj,
6365                    exi,
6366                    exo,
6367                    exp_d,
6368                    pt,
6369                    aq,
6370                    ad,
6371                    inf,
6372                    outf,
6373                    n_expert,
6374                    n_active,
6375                    n_pairs,
6376                    qtype,
6377                    rb,
6378                )
6379            } else {
6380                e.moe_pairs_matvec_q8_em(
6381                    &dev.ptr_row,
6382                    proj,
6383                    exi,
6384                    exo,
6385                    exp_d,
6386                    pt,
6387                    aq,
6388                    ad,
6389                    inf,
6390                    outf,
6391                    n_expert,
6392                    n_active,
6393                    n_pairs,
6394                    qtype,
6395                    rb,
6396                )
6397            }
6398        };
6399        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6400        let gate = matvec(
6401            0,
6402            &exi,
6403            &exo,
6404            &exp_d,
6405            &pt,
6406            &zq,
6407            &zd,
6408            n_embd,
6409            n_ff_exp,
6410            m.gate_exps.qtype,
6411            rbg_d,
6412        )?;
6413        let up = matvec(
6414            1,
6415            &exi,
6416            &exo,
6417            &exp_d,
6418            &pt,
6419            &zq,
6420            &zd,
6421            n_embd,
6422            n_ff_exp,
6423            m.up_exps.qtype,
6424            rbu_d,
6425        )?;
6426        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
6427        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6428        // down consumes PAIR-major activation rows: pair_tok = identity.
6429        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
6430        let pself = e.htod_i32(&pair_self)?;
6431        let y_down = matvec(
6432            2,
6433            &exi,
6434            &exo,
6435            &exp_d,
6436            &pself,
6437            &aq2,
6438            &ad2,
6439            n_ff_exp,
6440            n_embd,
6441            m.down_exps.qtype,
6442            m.down_exps.row_bytes,
6443        )?;
6444        let mut moe_out = e.uninit(t * n_embd)?; // scatter fully overwrites per (token,col)
6445        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
6446
6447        // SHARED EXPERT epilogue — same as the other paths.
6448        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
6449        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
6450        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6451            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6452        {
6453            let n_ff_sh = gate_shexp.out_features();
6454            // These decode-exact forms are required by the new Step resident arm. Keep the
6455            // established grouped shared-expert program for every other architecture: widening
6456            // this to Gemma changed its speculative acceptance despite green argmax gates.
6457            let step_exact = cfg.step35.is_some();
6458            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
6459            let (sg_gate, sg_up) = if step_exact && t == 1 {
6460                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6461                    Some(pair) => pair,
6462                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6463                }
6464            } else if verify_t {
6465                let mut fused = None;
6466                if crate::spec::spec_fused_t()
6467                    && (2..=4).contains(&t)
6468                    && e.uses_q8_1_fast(gate_shexp)
6469                    && e.uses_q8_1_fast(up_shexp)
6470                {
6471                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6472                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6473                }
6474                match fused {
6475                    Some(pair) => pair,
6476                    None => (
6477                        e.matmul_decode_exact(gate_shexp, z, t)?,
6478                        e.matmul_decode_exact(up_shexp, z, t)?,
6479                    ),
6480                }
6481            } else {
6482                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6483            };
6484            let mut sa = e.uninit(t * n_ff_sh)?;
6485            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6486            let sh = if verify_t {
6487                e.matmul_decode_exact(down_shexp, &sa, t)?
6488            } else {
6489                e.matmul(down_shexp, &sa, t)?
6490            };
6491            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6492            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
6493            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
6494            // dispatch choice cannot change bits.
6495            let g = match &m.gate_inp_shexp {
6496                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
6497                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6498                }
6499                Some(gate_inp_shexp) => {
6500                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6501                    let mut g = e.uninit(t)?;
6502                    e.sigmoid(&gs, &mut g, t)?;
6503                    g
6504                }
6505                None => e.htod(&vec![1.0f32; t])?,
6506            };
6507            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
6508        }
6509        Ok(moe_out)
6510    }
6511
6512    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
6513    #[allow(clippy::too_many_arguments)]
6514    #[allow(clippy::too_many_arguments)]
6515    fn moe_ffn_dev(
6516        e: &Engine,
6517        m: &MoeWeights,
6518        z: &CudaSlice<f32>,
6519        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
6520        logits: &CudaSlice<f32>,
6521        t: usize,
6522        cfg: &ModelConfig,
6523        il: u16,
6524        max_block: usize,
6525    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6526        let moe = cfg.moe.as_ref().unwrap();
6527        let n_embd = cfg.n_embd as usize;
6528        let n_expert = moe.expert_count as usize;
6529        let n_used = moe.expert_used_count as usize;
6530        let n_ff_exp = moe.expert_ff_length as usize;
6531        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
6532        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
6533        // clamped layers; assert both so a future caller that skips the gate fails loudly.
6534        debug_assert!(
6535            cfg.sigmoid_router().is_none(),
6536            "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts"
6537        );
6538        debug_assert!(
6539            !cfg.swiglu_clamped_at(il as u32),
6540            "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form"
6541        );
6542
6543        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
6544        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
6545        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
6546        // skipped entirely for macro-free experts (every k-quant GGUF).
6547        if m.has_macros {
6548            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
6549        }
6550
6551        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
6552        let mut moe_out = e.uninit(t * n_embd)?;
6553
6554        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
6555        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
6556        if let Some(dev) = m.dev_exps.as_ref() {
6557            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
6558            // the combined stride; up's base is offset in the ptr table. Down unchanged.
6559            let (rbg_d, rbu_d) = if dev.gu_il {
6560                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes;
6561                (sxx, sxx)
6562            } else {
6563                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
6564            };
6565            let q8 = moe_q8_enabled()
6566                && q8_expert_supported(m.gate_exps.qtype)
6567                && q8_expert_supported(m.up_exps.qtype)
6568                && q8_expert_supported(m.down_exps.qtype);
6569            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
6570            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
6571            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
6572            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
6573            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
6574            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
6575            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
6576            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
6577            let rows_arm = q8
6578                && t > 1
6579                && crate::spec::spec_m2()
6580                && n_ff_exp == 512
6581                && n_used <= 8
6582                && std::env::var("MEMRA_MOE_DEVQ8_GU")
6583                    .map(|v| v.is_empty() || v == "v")
6584                    .unwrap_or(true)
6585                && std::env::var("MEMRA_MOE_DEVQ8_DOWN")
6586                    .map(|v| v.is_empty() || v == "w8h2v")
6587                    .unwrap_or(true);
6588            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
6589            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
6590            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
6591            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
6592            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
6593            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
6594            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
6595            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
6596            let csr_mode = std::env::var("MEMRA_MOE_CSR")
6597                .ok()
6598                .and_then(|v| v.parse::<i32>().ok())
6599                .unwrap_or(1);
6600            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
6601            let csr_arm = rows_arm
6602                && csr_mode > 0
6603                && t <= 10
6604                && csr_qt(m.gate_exps.qtype)
6605                && csr_qt(m.up_exps.qtype)
6606                && csr_qt(m.down_exps.qtype);
6607            if csr_arm {
6608                if csr_mode == 2 {
6609                    static ENGAGED: std::sync::Once = std::sync::Once::new();
6610                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
6611                }
6612                let n_pairs = t * n_used;
6613                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6614                let act = e.moe_gate_up_silu8_dev_q8_csr(
6615                    &dev.ptr_row,
6616                    &sel_d,
6617                    &zq,
6618                    &zd,
6619                    n_pairs,
6620                    n_embd,
6621                    n_ff_exp,
6622                    n_used,
6623                    n_expert,
6624                    m.gate_exps.qtype,
6625                    m.up_exps.qtype,
6626                    rbg_d,
6627                    rbu_d,
6628                )?;
6629                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
6630                // down stays on the _rows twin — BOTH CSR down variants measured negative
6631                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
6632                // 16-group rows have too little decode to amortize any dedup structure.
6633                e.moe_down8_fma_dev_q8_rows(
6634                    &dev.ptr_row,
6635                    &sel_d,
6636                    &w_d,
6637                    &aq2,
6638                    &ad2,
6639                    &mut moe_out,
6640                    t,
6641                    n_ff_exp,
6642                    n_embd,
6643                    n_used,
6644                    n_expert,
6645                    m.down_exps.qtype,
6646                    m.down_exps.row_bytes,
6647                )?;
6648                if csr_mode == 2 {
6649                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
6650                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(
6651                        &dev.ptr_row,
6652                        &sel_d,
6653                        &zq,
6654                        &zd,
6655                        t,
6656                        n_embd,
6657                        n_ff_exp,
6658                        n_used,
6659                        n_expert,
6660                        m.gate_exps.qtype,
6661                        m.up_exps.qtype,
6662                        rbg_d,
6663                        rbu_d,
6664                        &m.dev_macros,
6665                    )?;
6666                    let mut out_r = e.uninit(t * n_embd)?;
6667                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
6668                    e.moe_down8_fma_dev_q8_rows(
6669                        &dev.ptr_row,
6670                        &sel_d,
6671                        &w_d,
6672                        &aq2r,
6673                        &ad2r,
6674                        &mut out_r,
6675                        t,
6676                        n_ff_exp,
6677                        n_embd,
6678                        n_used,
6679                        n_expert,
6680                        m.down_exps.qtype,
6681                        m.down_exps.row_bytes,
6682                    )?;
6683                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
6684                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
6685                    let ba = a1
6686                        .iter()
6687                        .zip(&a2)
6688                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6689                        .count();
6690                    let bo = o1
6691                        .iter()
6692                        .zip(&o2)
6693                        .filter(|(x, y)| x.to_bits() != y.to_bits())
6694                        .count();
6695                    if ba + bo > 0 {
6696                        eprintln!(
6697                            "[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
6698                            a1.len(),
6699                            o1.len()
6700                        );
6701                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
6702                        let sel_h = e.dtoh_i32(&sel_d)?;
6703                        let mut shown = 0;
6704                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
6705                            if x.to_bits() != y.to_bits() && shown < 4 {
6706                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
6707                                let ex = sel_h[p];
6708                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
6709                                eprintln!(
6710                                    "  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}"
6711                                );
6712                                shown += 1;
6713                            }
6714                        }
6715                        std::process::exit(3);
6716                    }
6717                }
6718            } else if rows_arm {
6719                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
6720                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
6721                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
6722                    use std::sync::atomic::{AtomicU64, Ordering};
6723                    static PAIRS: AtomicU64 = AtomicU64::new(0);
6724                    static UNIQ: AtomicU64 = AtomicU64::new(0);
6725                    static CALLS: AtomicU64 = AtomicU64::new(0);
6726                    let sel_h = e.dtoh_i32(&sel_d)?;
6727                    let mut u: Vec<i32> = sel_h.clone();
6728                    u.sort_unstable();
6729                    u.dedup();
6730                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
6731                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
6732                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
6733                    if c % 480 == 0 {
6734                        let p = PAIRS.load(Ordering::Relaxed);
6735                        let q = UNIQ.load(Ordering::Relaxed);
6736                        eprintln!(
6737                            "[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
6738                            q as f64 / p as f64
6739                        );
6740                    }
6741                }
6742                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6743                let act = e.moe_gate_up_silu8_dev_q8_rows(
6744                    &dev.ptr_row,
6745                    &sel_d,
6746                    &zq,
6747                    &zd,
6748                    t,
6749                    n_embd,
6750                    n_ff_exp,
6751                    n_used,
6752                    n_expert,
6753                    m.gate_exps.qtype,
6754                    m.up_exps.qtype,
6755                    rbg_d,
6756                    rbu_d,
6757                    &m.dev_macros,
6758                )?;
6759                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6760                e.moe_down8_fma_dev_q8_rows(
6761                    &dev.ptr_row,
6762                    &sel_d,
6763                    &w_d,
6764                    &aq2,
6765                    &ad2,
6766                    &mut moe_out,
6767                    t,
6768                    n_ff_exp,
6769                    n_embd,
6770                    n_used,
6771                    n_expert,
6772                    m.down_exps.qtype,
6773                    m.down_exps.row_bytes,
6774                )?;
6775            } else {
6776                for tok in 0..t {
6777                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6778                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6779                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6780                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6781                    if q8 {
6782                        let (zq, zd) = match (t, zq8) {
6783                            (1, Some((q, d))) => (q.clone(), d.clone()),
6784                            _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
6785                        };
6786                        let act = e.moe_gate_up_silu8_dev_q8(
6787                            &dev.ptr_row,
6788                            &selt,
6789                            &zq,
6790                            &zd,
6791                            n_embd,
6792                            n_ff_exp,
6793                            n_used,
6794                            n_expert,
6795                            m.gate_exps.qtype,
6796                            m.up_exps.qtype,
6797                            rbg_d,
6798                            rbu_d,
6799                            &m.dev_macros,
6800                        )?;
6801                        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6802                        e.moe_down8_fma_dev_q8(
6803                            &dev.ptr_row,
6804                            &selt,
6805                            &wt,
6806                            &aq2,
6807                            &ad2,
6808                            &mut dst,
6809                            n_ff_exp,
6810                            n_embd,
6811                            n_used,
6812                            n_expert,
6813                            m.down_exps.qtype,
6814                            m.down_exps.row_bytes,
6815                        )?;
6816                    } else {
6817                        let act = e.moe_gate_up_silu8_dev(
6818                            &dev.ptr_row,
6819                            &selt,
6820                            &zt,
6821                            n_embd,
6822                            n_ff_exp,
6823                            n_used,
6824                            n_expert,
6825                            m.gate_exps.qtype,
6826                            m.up_exps.qtype,
6827                            rbg_d,
6828                            rbu_d,
6829                            &m.dev_macros,
6830                        )?;
6831                        e.moe_down8_fma_dev(
6832                            &dev.ptr_row,
6833                            &selt,
6834                            &wt,
6835                            &act,
6836                            &mut dst,
6837                            n_ff_exp,
6838                            n_embd,
6839                            n_used,
6840                            n_expert,
6841                            m.down_exps.qtype,
6842                            m.down_exps.row_bytes,
6843                        )?;
6844                    }
6845                }
6846            }
6847        } else {
6848            // Launch under the cache lock: the row borrow lives as long as the closure, and the
6849            // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
6850            // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
6851            // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
6852            // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
6853            // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
6854            let q8 = moe_q8_enabled()
6855                && q8_expert_supported(m.gate_exps.qtype)
6856                && q8_expert_supported(m.up_exps.qtype)
6857                && q8_expert_supported(m.down_exps.qtype);
6858            e.with_moe_cache(max_block, |c, eng| {
6859                let row = c
6860                    .layer_dev_row(il, n_expert, eng)?
6861                    .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
6862                for tok in 0..t {
6863                    let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
6864                    let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
6865                    let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
6866                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
6867                    if q8 {
6868                        let (zq, zd) = match (t, zq8) {
6869                            (1, Some((q, d))) => (q.clone(), d.clone()),
6870                            _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
6871                        };
6872                        let act = eng.moe_gate_up_silu8_dev_q8(
6873                            row,
6874                            &selt,
6875                            &zq,
6876                            &zd,
6877                            n_embd,
6878                            n_ff_exp,
6879                            n_used,
6880                            n_expert,
6881                            m.gate_exps.qtype,
6882                            m.up_exps.qtype,
6883                            m.gate_exps.row_bytes,
6884                            m.up_exps.row_bytes,
6885                            &m.dev_macros,
6886                        )?;
6887                        let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
6888                        eng.moe_down8_fma_dev_q8(
6889                            row,
6890                            &selt,
6891                            &wt,
6892                            &aq2,
6893                            &ad2,
6894                            &mut dst,
6895                            n_ff_exp,
6896                            n_embd,
6897                            n_used,
6898                            n_expert,
6899                            m.down_exps.qtype,
6900                            m.down_exps.row_bytes,
6901                        )?;
6902                    } else {
6903                        let act = eng.moe_gate_up_silu8_dev(
6904                            row,
6905                            &selt,
6906                            &zt,
6907                            n_embd,
6908                            n_ff_exp,
6909                            n_used,
6910                            n_expert,
6911                            m.gate_exps.qtype,
6912                            m.up_exps.qtype,
6913                            m.gate_exps.row_bytes,
6914                            m.up_exps.row_bytes,
6915                            &m.dev_macros,
6916                        )?;
6917                        eng.moe_down8_fma_dev(
6918                            row,
6919                            &selt,
6920                            &wt,
6921                            &act,
6922                            &mut dst,
6923                            n_ff_exp,
6924                            n_embd,
6925                            n_used,
6926                            n_expert,
6927                            m.down_exps.qtype,
6928                            m.down_exps.row_bytes,
6929                        )?;
6930                    }
6931                }
6932                // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
6933                c.hits += (t * 3 * n_used) as u64;
6934                Ok(())
6935            })?;
6936        }
6937
6938        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
6939        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
6940        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
6941        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
6942        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6943            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6944        {
6945            let n_ff_sh = gate_shexp.out_features();
6946            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
6947            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
6948            let verify_t = t > 1 && t < PRIME_MIN_T;
6949            let (sg_gate, sg_up) = if t == 1 {
6950                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
6951                    Some(pair) => pair,
6952                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
6953                }
6954            } else if verify_t {
6955                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
6956                // rides one shared quantize + one fused2 batched launch instead of two
6957                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
6958                let mut fused = None;
6959                if crate::spec::spec_fused_t()
6960                    && (2..=4).contains(&t)
6961                    && e.uses_q8_1_fast(gate_shexp)
6962                    && e.uses_q8_1_fast(up_shexp)
6963                {
6964                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
6965                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
6966                }
6967                match fused {
6968                    Some(pair) => pair,
6969                    None => (
6970                        e.matmul_decode_exact(gate_shexp, z, t)?,
6971                        e.matmul_decode_exact(up_shexp, z, t)?,
6972                    ),
6973                }
6974            } else {
6975                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
6976            };
6977            let mut sa = e.uninit(t * n_ff_sh)?; // silu_mul fully overwrites
6978            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
6979            let sh = if verify_t {
6980                e.matmul_decode_exact(down_shexp, &sa, t)?
6981            } else {
6982                e.matmul(down_shexp, &sa, t)?
6983            };
6984            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
6985            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
6986            // between the two arms; prefill keeps the batched cuBLASLt linear).
6987            let g = match &m.gate_inp_shexp {
6988                Some(gate_inp_shexp) => {
6989                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
6990                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
6991                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
6992                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
6993                    } else {
6994                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
6995                        let mut g = e.uninit(t)?;
6996                        e.sigmoid(&gs, &mut g, t)?;
6997                        g
6998                    }
6999                }
7000                None => e.htod(&vec![1.0f32; t])?,
7001            };
7002            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
7003        }
7004
7005        Ok(moe_out)
7006    }
7007
7008    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
7009    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
7010    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
7011    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
7012    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
7013    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
7014    /// the collected raw pointers cannot move between collection and launch (single-threaded
7015    /// decode; the lock is held only for collection, launches are stream-ordered after any
7016    /// prior same-stream staging writes).
7017    #[allow(clippy::too_many_arguments)]
7018    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
7019    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
7020    #[allow(clippy::too_many_arguments)]
7021    fn moe_gdec_token_q8(
7022        e: &Engine,
7023        m: &MoeWeights,
7024        il: u16,
7025        max_block: usize,
7026        zq: &CudaSlice<i8>,
7027        zd: &CudaSlice<f32>,
7028        sel: &[u32],
7029        w: &[f32],
7030        moe_out: &mut CudaSlice<f32>,
7031        tok: usize,
7032        n_embd: usize,
7033        n_ff_exp: usize,
7034        n_used: usize,
7035    ) -> Result<bool, Box<dyn std::error::Error>> {
7036        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7037        use cudarc::driver::DevicePtr;
7038        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7039            let mut g = [0u64; 8];
7040            let mut u = [0u64; 8];
7041            let mut d = [0u64; 8];
7042            for (j, &ex) in sel.iter().enumerate() {
7043                let ex = ex as u16;
7044                let (Some(sg), Some(su), Some(sd)) = (
7045                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7046                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7047                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7048                ) else {
7049                    return Ok(None);
7050                };
7051                let __s = eng.stream();
7052                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7053                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7054                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7055                g[j] = pg as u64;
7056                u[j] = pu as u64;
7057                d[j] = pd as u64;
7058            }
7059            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7060                for &ex in sel {
7061                    let ex = ex as u16;
7062                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7063                        c.note_profile_hit(BlockId::new(il, proj, ex));
7064                    }
7065                }
7066            }
7067            c.hits += (3 * n_used) as u64;
7068            Ok(Some((g, u, d)))
7069        })?;
7070        let Some((g, u, d)) = ptrs else {
7071            return Ok(false);
7072        };
7073        let mut wv = [0f32; 8];
7074        wv[..n_used].copy_from_slice(w);
7075        let act = e.moe_gate_up_silu8_q8(
7076            crate::WPtr8(g),
7077            crate::WPtr8(u),
7078            zq,
7079            zd,
7080            n_embd,
7081            n_ff_exp,
7082            n_used,
7083            m.gate_exps.qtype,
7084            m.up_exps.qtype,
7085            m.gate_exps.row_bytes,
7086            m.up_exps.row_bytes,
7087        )?;
7088        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
7089        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7090        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7091        e.moe_down8_fma_q8(
7092            crate::WPtr8(d),
7093            crate::F32x8(wv),
7094            &aq2,
7095            &ad2,
7096            &mut dst,
7097            n_ff_exp,
7098            n_embd,
7099            n_used,
7100            m.down_exps.qtype,
7101            m.down_exps.row_bytes,
7102        )?;
7103        Ok(true)
7104    }
7105
7106    fn moe_gdec_token(
7107        e: &Engine,
7108        m: &MoeWeights,
7109        il: u16,
7110        max_block: usize,
7111        zt: &cudarc::driver::CudaView<f32>,
7112        sel: &[u32],
7113        w: &[f32],
7114        moe_out: &mut CudaSlice<f32>,
7115        tok: usize,
7116        n_embd: usize,
7117        n_ff_exp: usize,
7118        n_used: usize,
7119    ) -> Result<bool, Box<dyn std::error::Error>> {
7120        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7121        use cudarc::driver::DevicePtr;
7122        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
7123        let ptrs = e.with_moe_cache(max_block, |c, eng| {
7124            let mut g = [0u64; 8];
7125            let mut u = [0u64; 8];
7126            let mut d = [0u64; 8];
7127            for (j, &ex) in sel.iter().enumerate() {
7128                let ex = ex as u16;
7129                let (Some(sg), Some(su), Some(sd)) = (
7130                    c.resident(BlockId::new(il, PROJ_GATE, ex)),
7131                    c.resident(BlockId::new(il, PROJ_UP, ex)),
7132                    c.resident(BlockId::new(il, PROJ_DOWN, ex)),
7133                ) else {
7134                    return Ok(None);
7135                };
7136                let __s = eng.stream();
7137                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
7138                let (pu, _e1) = c.slot(su).device_ptr(&__s);
7139                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
7140                g[j] = pg as u64;
7141                u[j] = pu as u64;
7142                d[j] = pd as u64;
7143            }
7144            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
7145                for &ex in sel {
7146                    let ex = ex as u16;
7147                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
7148                        c.note_profile_hit(BlockId::new(il, proj, ex));
7149                    }
7150                }
7151            }
7152            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
7153            Ok(Some((g, u, d)))
7154        })?;
7155        let Some((g, u, d)) = ptrs else {
7156            return Ok(false);
7157        };
7158        let mut wv = [0f32; 8];
7159        wv[..n_used].copy_from_slice(w);
7160        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
7161        let act = e.moe_gate_up_silu8(
7162            crate::WPtr8(g),
7163            crate::WPtr8(u),
7164            zt,
7165            n_embd,
7166            n_ff_exp,
7167            n_used,
7168            m.gate_exps.qtype,
7169            m.up_exps.qtype,
7170            m.gate_exps.row_bytes,
7171            m.up_exps.row_bytes,
7172        )?;
7173        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7174        e.moe_down8_fma_into(
7175            crate::WPtr8(d),
7176            crate::F32x8(wv),
7177            &act,
7178            &mut dst,
7179            n_ff_exp,
7180            n_embd,
7181            n_used,
7182            m.down_exps.qtype,
7183            m.down_exps.row_bytes,
7184        )?;
7185        Ok(true)
7186    }
7187
7188    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
7189    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
7190    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
7191    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
7192    fn moe_cached_gemm_q8(
7193        e: &Engine,
7194        il: u16,
7195        proj: u8,
7196        ex: usize,
7197        m: &MoeWeights,
7198        max_block: usize,
7199        aq: &CudaSlice<i8>,
7200        ad: &CudaSlice<f32>,
7201    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7202        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7203        let exps = match proj {
7204            PROJ_GATE => &m.gate_exps,
7205            PROJ_UP => &m.up_exps,
7206            _ => &m.down_exps,
7207        };
7208        let layout = exps.expert_layout(ex);
7209        let id = BlockId::new(il, proj, ex as u16);
7210        let source = exps.expert_source(ex);
7211        e.with_moe_cache(max_block, |c, eng| {
7212            let slot = c.dispatch_source(id, source, eng)?;
7213            let DispatchSlot::Resident(sl) = slot;
7214            let buf = c.slot(sl);
7215            eng.qmatvec_expert_q8(
7216                buf,
7217                0..layout.len,
7218                aq,
7219                ad,
7220                1,
7221                exps.in_f,
7222                exps.out_f,
7223                layout.qtype,
7224                layout.row_bytes,
7225            )
7226        })
7227    }
7228
7229    fn moe_cached_gemm(
7230        e: &Engine,
7231        il: u16,
7232        proj: u8,
7233        ex: usize,
7234        m: &MoeWeights,
7235        max_block: usize,
7236        x: &cudarc::driver::CudaView<f32>,
7237    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7238        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
7239        let exps = match proj {
7240            PROJ_GATE => &m.gate_exps,
7241            PROJ_UP => &m.up_exps,
7242            _ => &m.down_exps,
7243        };
7244        let layout = exps.expert_layout(ex);
7245        let id = BlockId::new(il, proj, ex as u16);
7246        let source = exps.expert_source(ex);
7247        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
7248        e.with_moe_cache(max_block, |c, eng| {
7249            let slot = c.dispatch_source(id, source, eng)?;
7250            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
7251            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
7252            let DispatchSlot::Resident(sl) = slot;
7253            let buf = c.slot(sl);
7254            eng.qmatvec_view(
7255                buf,
7256                0..layout.len,
7257                x,
7258                1,
7259                exps.in_f,
7260                exps.out_f,
7261                layout.qtype,
7262                layout.row_bytes,
7263            )
7264        })
7265    }
7266
7267    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
7268    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
7269    /// so the current forward's backend assignment and output remain unchanged.
7270    fn moe_profile_admit_expert(
7271        e: &Engine,
7272        il: u16,
7273        ex: usize,
7274        m: &MoeWeights,
7275        max_block: usize,
7276    ) -> Result<(), Box<dyn std::error::Error>> {
7277        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7278        e.with_moe_cache(max_block, |cache, eng| {
7279            for (proj, exps) in [
7280                (PROJ_GATE, &m.gate_exps),
7281                (PROJ_UP, &m.up_exps),
7282                (PROJ_DOWN, &m.down_exps),
7283            ] {
7284                let id = BlockId::new(il, proj, ex as u16);
7285                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
7286            }
7287            Ok(())
7288        })
7289    }
7290
7291    /// Read a projection from the immutable residency set when present; otherwise use one
7292    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
7293    #[allow(clippy::too_many_arguments)]
7294    fn moe_frozen_gemm(
7295        e: &Engine,
7296        il: u16,
7297        proj: u8,
7298        ex: usize,
7299        m: &MoeWeights,
7300        max_block: usize,
7301        x: &cudarc::driver::CudaView<f32>,
7302        scratch: &mut Option<CudaSlice<u8>>,
7303        scratch_len: usize,
7304    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7305        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
7306        let exps = match proj {
7307            PROJ_GATE => &m.gate_exps,
7308            PROJ_UP => &m.up_exps,
7309            _ => &m.down_exps,
7310        };
7311        let layout = exps.expert_layout(ex);
7312        let id = BlockId::new(il, proj, ex as u16);
7313        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
7314            let Some(slot) = cache.resident(id) else {
7315                return Ok(None);
7316            };
7317            let buf = cache.slot(slot);
7318            Ok(Some(eng.qmatvec_view(
7319                buf,
7320                0..layout.len,
7321                x,
7322                1,
7323                exps.in_f,
7324                exps.out_f,
7325                layout.qtype,
7326                layout.row_bytes,
7327            )?))
7328        })? {
7329            return Ok(output);
7330        }
7331        if scratch.is_none() {
7332            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
7333        }
7334        let scratch = scratch.as_mut().unwrap();
7335        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
7336        e.qmatvec_view(
7337            scratch,
7338            0..layout.len,
7339            x,
7340            1,
7341            exps.in_f,
7342            exps.out_f,
7343            layout.qtype,
7344            layout.row_bytes,
7345        )
7346    }
7347
7348    fn moe_prefetch_expert(
7349        e: &Engine,
7350        il: u16,
7351        ex: usize,
7352        m: &MoeWeights,
7353        max_block: usize,
7354        keep: &[crate::moe_cache::BlockId],
7355    ) -> Result<(), Box<dyn std::error::Error>> {
7356        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7357        e.with_moe_cache(max_block, |c, eng| {
7358            for (proj, exps) in [
7359                (PROJ_GATE, &m.gate_exps),
7360                (PROJ_UP, &m.up_exps),
7361                (PROJ_DOWN, &m.down_exps),
7362            ] {
7363                let id = BlockId::new(il, proj, ex as u16);
7364                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
7365            }
7366            Ok(())
7367        })
7368    }
7369
7370    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
7371    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
7372    fn moe_prefetch_disk_expert(
7373        e: &Engine,
7374        il: u16,
7375        ex: usize,
7376        m: &MoeWeights,
7377        max_block: usize,
7378        keep: &[crate::moe_cache::BlockId],
7379    ) -> Result<(), Box<dyn std::error::Error>> {
7380        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7381        e.with_moe_cache(max_block, |c, eng| {
7382            for (proj, exps) in [
7383                (PROJ_GATE, &m.gate_exps),
7384                (PROJ_UP, &m.up_exps),
7385                (PROJ_DOWN, &m.down_exps),
7386            ] {
7387                let source = exps.expert_source(ex);
7388                if let crate::model::ExpertSource::Disk { .. } = &source {
7389                    let id = BlockId::new(il, proj, ex as u16);
7390                    let _ = c.prefetch_source(id, source, keep, eng)?;
7391                }
7392            }
7393            Ok(())
7394        })
7395    }
7396
7397    #[inline]
7398    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
7399        let _ = m.gate_exps.prefetch_expert_pages(ex);
7400        let _ = m.up_exps.prefetch_expert_pages(ex);
7401        let _ = m.down_exps.prefetch_expert_pages(ex);
7402    }
7403}
7404
7405// ================================================================================================
7406// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
7407//
7408// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
7409// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
7410// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
7411//
7412// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
7413// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
7414// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
7415// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
7416// identical to the per-token loop regardless of expert processing order.
7417//
7418// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
7419// ================================================================================================
7420
7421impl HybridModel {
7422    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
7423    /// sequential fused q8 program over the token axis; clamped layers use the separate
7424    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
7425    #[allow(clippy::too_many_arguments)]
7426    fn moe_ffn_grouped_resident_q8(
7427        e: &Engine,
7428        m: &MoeWeights,
7429        z: &CudaSlice<f32>,
7430        t: usize,
7431        cfg: &ModelConfig,
7432        il: u16,
7433        sel_all: &[u32],
7434        w_all: &[f32],
7435        table: &CudaSlice<u64>,
7436        gu_il: bool,
7437    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7438        let moe = cfg.moe.as_ref().unwrap();
7439        let n_embd = cfg.n_embd as usize;
7440        let n_expert = moe.expert_count as usize;
7441        let n_used = moe.expert_used_count as usize;
7442        let n_ff_exp = moe.expert_ff_length as usize;
7443        let n_pairs = t * n_used;
7444        debug_assert_eq!(sel_all.len(), n_pairs);
7445        debug_assert_eq!(w_all.len(), n_pairs);
7446        debug_assert!(
7447            m.gate_exps.macros.is_none()
7448                && m.up_exps.macros.is_none()
7449                && m.down_exps.macros.is_none(),
7450            "resident grouped q8 does not fold per-expert macro scales",
7451        );
7452
7453        // The rows twins run the resident sequential program verbatim on grid.z = token:
7454        // fused gate/up/SiLU per slot, batched activation quantization, then the original
7455        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
7456        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
7457        // never enter the softmax router.
7458        if !cfg.swiglu_clamped_at(il as u32) {
7459            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7460            let sel_d = e.htod_i32(&sel)?;
7461            let w_d = e.htod(w_all)?;
7462            let (gate_row_bytes, up_row_bytes) = if gu_il {
7463                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7464                (combined, combined)
7465            } else {
7466                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7467            };
7468            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7469            let act = e.moe_gate_up_silu8_dev_q8_rows(
7470                table,
7471                &sel_d,
7472                &zq,
7473                &zd,
7474                t,
7475                n_embd,
7476                n_ff_exp,
7477                n_used,
7478                n_expert,
7479                m.gate_exps.qtype,
7480                m.up_exps.qtype,
7481                gate_row_bytes,
7482                up_row_bytes,
7483                &m.dev_macros,
7484            )?;
7485            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7486            let mut moe_out = e.uninit(t * n_embd)?;
7487            e.moe_down8_fma_dev_q8_rows_g(
7488                table,
7489                &sel_d,
7490                &w_d,
7491                &aq2,
7492                &ad2,
7493                &mut moe_out,
7494                t,
7495                n_ff_exp,
7496                n_embd,
7497                n_used,
7498                n_expert,
7499                m.down_exps.qtype,
7500                m.down_exps.row_bytes,
7501            )?;
7502
7503            if std::env::var("MEMRA_MOE_STATS").is_ok() {
7504                let mut counts = vec![0usize; n_expert];
7505                for &expert in sel_all {
7506                    counts[expert as usize] += 1;
7507                }
7508                let mut sizes: Vec<usize> =
7509                    counts.into_iter().filter(|&count| count != 0).collect();
7510                sizes.sort_unstable();
7511                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7512                println!(
7513                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
7514                     m_e: min={} median={} mean={mean:.1} max={}",
7515                    sizes.len(),
7516                    n_expert,
7517                    sizes.first().copied().unwrap_or(0),
7518                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7519                    sizes.last().copied().unwrap_or(0),
7520                );
7521            }
7522            return Ok(moe_out);
7523        }
7524
7525        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
7526        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
7527        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
7528        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
7529        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
7530        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7531        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7532
7533        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7534        for (pair, &expert) in pair_ex.iter().enumerate() {
7535            by_expert[expert as usize].push(pair as i32);
7536        }
7537
7538        let pair_tok_d = e.htod_i32(&pair_tok)?;
7539        let pair_ex_d = e.htod_i32(&pair_ex)?;
7540        let pair_w_d = e.htod(w_all)?;
7541        let tok_off_d = e.htod_i32(&tok_off)?;
7542        let tok_ids_d = e.htod_i32(&tok_ids)?;
7543
7544        let matvec = |proj: i32,
7545                      pair_rows: &CudaSlice<i32>,
7546                      aq: &CudaSlice<i8>,
7547                      ad: &CudaSlice<f32>,
7548                      in_f: usize,
7549                      out_f: usize,
7550                      qtype: i32,
7551                      row_bytes: usize|
7552         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7553            e.moe_pairs_matvec_q8(
7554                table, proj, pair_rows, &pair_ex_d, aq, ad, in_f, out_f, n_expert, n_pairs, qtype,
7555                row_bytes,
7556            )
7557        };
7558
7559        let (gate_row_bytes, up_row_bytes) = if gu_il {
7560            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
7561            (combined, combined)
7562        } else {
7563            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
7564        };
7565        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
7566        let gate = matvec(
7567            0,
7568            &pair_tok_d,
7569            &zq,
7570            &zd,
7571            n_embd,
7572            n_ff_exp,
7573            m.gate_exps.qtype,
7574            gate_row_bytes,
7575        )?;
7576        let up = matvec(
7577            1,
7578            &pair_tok_d,
7579            &zq,
7580            &zd,
7581            n_embd,
7582            n_ff_exp,
7583            m.up_exps.qtype,
7584            up_row_bytes,
7585        )?;
7586        let mut act = e.uninit(n_pairs * n_ff_exp)?;
7587        Self::ffn_act_lim(
7588            e,
7589            cfg,
7590            &gate,
7591            &up,
7592            1.0,
7593            1.0,
7594            cfg.clamp_exp_at(il as u32),
7595            &mut act,
7596            n_pairs * n_ff_exp,
7597        )?;
7598        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7599        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7600        let pair_self_d = e.htod_i32(&pair_self)?;
7601        let down = matvec(
7602            2,
7603            &pair_self_d,
7604            &aq2,
7605            &ad2,
7606            n_ff_exp,
7607            n_embd,
7608            m.down_exps.qtype,
7609            m.down_exps.row_bytes,
7610        )?;
7611        let mut moe_out = e.uninit(t * n_embd)?;
7612        e.moe_pairs_scatter(
7613            &down,
7614            &pair_w_d,
7615            &tok_off_d,
7616            &tok_ids_d,
7617            &mut moe_out,
7618            t,
7619            n_embd,
7620        )?;
7621
7622        if std::env::var("MEMRA_MOE_STATS").is_ok() {
7623            let mut sizes: Vec<usize> = by_expert
7624                .iter()
7625                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
7626                .collect();
7627            sizes.sort_unstable();
7628            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
7629            println!(
7630                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
7631                 m_e: min={} median={} mean={mean:.1} max={}",
7632                sizes.len(),
7633                n_expert,
7634                sizes.first().copied().unwrap_or(0),
7635                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
7636                sizes.last().copied().unwrap_or(0),
7637            );
7638        }
7639        Ok(moe_out)
7640    }
7641
7642    fn moe_ffn_grouped_add_shared(
7643        e: &Engine,
7644        m: &MoeWeights,
7645        z: &CudaSlice<f32>,
7646        t: usize,
7647        cfg: &ModelConfig,
7648        il: u16,
7649        moe_out: &mut CudaSlice<f32>,
7650    ) -> Result<(), Box<dyn std::error::Error>> {
7651        let n_embd = cfg.n_embd as usize;
7652        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
7653            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
7654        {
7655            let n_ff_sh = gate_shexp.out_features();
7656            let sg_gate = e.matmul(gate_shexp, z, t)?;
7657            let sg_up = e.matmul(up_shexp, z, t)?;
7658            let mut sa = e.uninit(t * n_ff_sh)?;
7659            Self::ffn_act_lim(
7660                e,
7661                cfg,
7662                &sg_gate,
7663                &sg_up,
7664                1.0,
7665                1.0,
7666                cfg.clamp_shexp_at(il as u32),
7667                &mut sa,
7668                t * n_ff_sh,
7669            )?;
7670            let sh = e.matmul(down_shexp, &sa, t)?;
7671            let gate = match &m.gate_inp_shexp {
7672                Some(gate_inp_shexp) => {
7673                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
7674                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
7675                    } else {
7676                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
7677                        let mut gate = e.uninit(t)?;
7678                        e.sigmoid(&raw, &mut gate, t)?;
7679                        gate
7680                    }
7681                }
7682                None => e.htod(&vec![1.0f32; t])?,
7683            };
7684            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
7685        }
7686        Ok(())
7687    }
7688
7689    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
7690    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
7691    pub(crate) fn moe_ffn_grouped(
7692        e: &Engine,
7693        m: &MoeWeights,
7694        z: &CudaSlice<f32>,
7695        t: usize,
7696        cfg: &ModelConfig,
7697        il: u16,
7698        max_block: usize,
7699    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7700        let moe = cfg.moe.as_ref().unwrap();
7701        let n_embd = cfg.n_embd as usize;
7702        let n_expert = moe.expert_count as usize;
7703        let n_used = moe.expert_used_count as usize;
7704        let n_ff_exp = moe.expert_ff_length as usize;
7705        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
7706        let lim_exp = cfg.clamp_exp_at(il as u32);
7707
7708        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
7709        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
7710        // enters the softmax-only pairs/dev router.
7711        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
7712        if let Some(sig) = cfg.sigmoid_router() {
7713            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
7714        }
7715        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
7716            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
7717        } else {
7718            Self::moe_route_cfg(e, &logits, t, n_expert, n_used, m.active_experts.as_deref())?
7719        };
7720        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
7721        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
7722        Self::trace_moe_input(e, il, t, n_embd, z)?;
7723
7724        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
7725        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
7726        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
7727        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
7728        let no_exp_macros = m.gate_exps.macros.is_none()
7729            && m.up_exps.macros.is_none()
7730            && m.down_exps.macros.is_none();
7731        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
7732            m.has_uniform_expert_layout()
7733                && no_exp_macros
7734                && moe_q8_enabled()
7735                && q8_expert_supported(m.gate_exps.qtype)
7736                && q8_expert_supported(m.up_exps.qtype)
7737                && q8_expert_supported(m.down_exps.qtype)
7738                && moe_slab_enabled()
7739                && dev.dev == e.ctx().ordinal()
7740        });
7741        if let Some(dev) = resident_q8 {
7742            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
7743                e,
7744                m,
7745                z,
7746                t,
7747                cfg,
7748                il,
7749                &sel_all,
7750                &w_all,
7751                &dev.ptr_row,
7752                dev.gu_il,
7753            )?;
7754            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
7755            return Ok(moe_out);
7756        }
7757
7758        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
7759        // For each expert e, we need: which tokens use it, their positions in z, their top-k
7760        // slot index (for bit-identical accumulation), and their weights.
7761        struct ExpertGroup {
7762            tok_indices: Vec<i32>,  // indices into z rows (0..T-1)
7763            slot_indices: Vec<i32>, // top-k slot (0..n_used-1) for that token-expert pair
7764            weights: Vec<f32>,      // renormalized weight for that token-expert pair
7765        }
7766        let mut groups: Vec<ExpertGroup> = (0..n_expert)
7767            .map(|_| ExpertGroup {
7768                tok_indices: Vec::new(),
7769                slot_indices: Vec::new(),
7770                weights: Vec::new(),
7771            })
7772            .collect();
7773
7774        for tok in 0..t {
7775            for j in 0..n_used {
7776                let ex = sel_all[tok * n_used + j] as usize;
7777                let w = w_all[tok * n_used + j];
7778                groups[ex].tok_indices.push(tok as i32);
7779                groups[ex].slot_indices.push(j as i32);
7780                groups[ex].weights.push(w);
7781            }
7782        }
7783
7784        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
7785        // Each token's 8 expert contributions land in their respective slots.
7786        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
7787        let mut wbuf = e.zeros(t * n_used)?; // [T, n_used] weight buffer for FMA reduce
7788
7789        // Expert weight dimensions (used in both cache and staging paths).
7790        let g_len = m.gate_exps.max_expert_bytes();
7791        let u_len = m.up_exps.max_expert_bytes();
7792        let d_len = m.down_exps.max_expert_bytes();
7793        let moe_q8 = m.has_uniform_expert_layout()
7794            && moe_q8_enabled()
7795            && q8_expert_supported(m.gate_exps.qtype)
7796            && q8_expert_supported(m.up_exps.qtype)
7797            && q8_expert_supported(m.down_exps.qtype);
7798        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
7799        // Interleaved GU slabs require the pointer-table fast path above.
7800        let slab_local = m
7801            .dev_exps
7802            .as_ref()
7803            .filter(|dev| !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal());
7804        let use_cache =
7805            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
7806        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
7807        // also does: a local resident slab or a live SLRU dispatch.
7808        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
7809
7810        // GPU scratch for staging (only allocated without a local slab or cache).
7811        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
7812            (
7813                Some(e.alloc_u8(g_len)?),
7814                Some(e.alloc_u8(u_len)?),
7815                Some(e.alloc_u8(d_len)?),
7816            )
7817        } else {
7818            (None, None, None)
7819        };
7820
7821        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
7822        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
7823        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
7824        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
7825        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
7826        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
7827        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
7828        // at long prompts where every expert stages regardless. Order is FREE to change without
7829        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
7830        // regardless of expert processing order (the whole point of the slots).
7831        let mut order: Vec<usize> = (0..n_expert)
7832            .filter(|&ex| !groups[ex].tok_indices.is_empty())
7833            .collect();
7834        order.sort_by(|&a, &b| {
7835            groups[b]
7836                .tok_indices
7837                .len()
7838                .cmp(&groups[a].tok_indices.len())
7839                .then(a.cmp(&b))
7840        });
7841        let mut m_dist: Vec<usize> = Vec::new(); // for stats
7842        let page_window = moe_page_prefetch_window();
7843        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
7844        if worker_disk_prefetch {
7845            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
7846                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
7847            }
7848        }
7849        for (order_pos, &ex) in order.iter().enumerate() {
7850            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
7851                Self::moe_prefetch_host_expert(order[next], m);
7852            }
7853            if worker_disk_prefetch {
7854                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
7855                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7856                    let keep = [
7857                        BlockId::new(il, PROJ_GATE, ex as u16),
7858                        BlockId::new(il, PROJ_UP, ex as u16),
7859                        BlockId::new(il, PROJ_DOWN, ex as u16),
7860                    ];
7861                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
7862                }
7863            }
7864            let grp = &groups[ex];
7865            let m_e = grp.tok_indices.len();
7866            m_dist.push(m_e);
7867            let gl = m.gate_exps.expert_layout(ex);
7868            let ul = m.up_exps.expert_layout(ex);
7869            let dl = m.down_exps.expert_layout(ex);
7870
7871            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
7872            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
7873            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
7874            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
7875            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
7876            let dmac = m.down_exps.macro_scale(ex);
7877            let weight_d = if dmac == 1.0 {
7878                e.htod(&grp.weights)?
7879            } else {
7880                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
7881                e.htod(&scaled)?
7882            };
7883
7884            // GATHER: collect m_e activation rows from z into a contiguous buffer.
7885            let mut gathered = e.zeros(m_e * n_embd)?;
7886            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
7887            let gv = gathered.slice(0..m_e * n_embd);
7888
7889            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
7890            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
7891            let y = if let Some(dev) = slab_local {
7892                let gate_start = ex * m.gate_exps.expert_stride;
7893                let up_start = ex * m.up_exps.expert_stride;
7894                let down_start = ex * m.down_exps.expert_stride;
7895                if grouped_q8 {
7896                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7897                    let gate = e.qmatvec_expert_q8(
7898                        &dev.gate,
7899                        gate_start..gate_start + gl.len,
7900                        &zq,
7901                        &zd,
7902                        m_e,
7903                        m.gate_exps.in_f,
7904                        m.gate_exps.out_f,
7905                        gl.qtype,
7906                        gl.row_bytes,
7907                    )?;
7908                    let up = e.qmatvec_expert_q8(
7909                        &dev.up,
7910                        up_start..up_start + ul.len,
7911                        &zq,
7912                        &zd,
7913                        m_e,
7914                        m.up_exps.in_f,
7915                        m.up_exps.out_f,
7916                        ul.qtype,
7917                        ul.row_bytes,
7918                    )?;
7919                    let mut act = e.uninit(m_e * n_ff_exp)?;
7920                    Self::ffn_act_lim(
7921                        e,
7922                        cfg,
7923                        &gate,
7924                        &up,
7925                        m.gate_exps.macro_scale(ex),
7926                        m.up_exps.macro_scale(ex),
7927                        lim_exp,
7928                        &mut act,
7929                        m_e * n_ff_exp,
7930                    )?;
7931                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
7932                    e.qmatvec_expert_q8(
7933                        &dev.down,
7934                        down_start..down_start + dl.len,
7935                        &aq2,
7936                        &ad2,
7937                        m_e,
7938                        m.down_exps.in_f,
7939                        m.down_exps.out_f,
7940                        dl.qtype,
7941                        dl.row_bytes,
7942                    )?
7943                } else {
7944                    let gate = e.qmatvec_view(
7945                        &dev.gate,
7946                        gate_start..gate_start + gl.len,
7947                        &gv,
7948                        m_e,
7949                        m.gate_exps.in_f,
7950                        m.gate_exps.out_f,
7951                        gl.qtype,
7952                        gl.row_bytes,
7953                    )?;
7954                    let up = e.qmatvec_view(
7955                        &dev.up,
7956                        up_start..up_start + ul.len,
7957                        &gv,
7958                        m_e,
7959                        m.up_exps.in_f,
7960                        m.up_exps.out_f,
7961                        ul.qtype,
7962                        ul.row_bytes,
7963                    )?;
7964                    let mut act = e.uninit(m_e * n_ff_exp)?;
7965                    Self::ffn_act_lim(
7966                        e,
7967                        cfg,
7968                        &gate,
7969                        &up,
7970                        m.gate_exps.macro_scale(ex),
7971                        m.up_exps.macro_scale(ex),
7972                        lim_exp,
7973                        &mut act,
7974                        m_e * n_ff_exp,
7975                    )?;
7976                    let actv = act.slice(0..m_e * n_ff_exp);
7977                    e.qmatvec_view(
7978                        &dev.down,
7979                        down_start..down_start + dl.len,
7980                        &actv,
7981                        m_e,
7982                        m.down_exps.in_f,
7983                        m.down_exps.out_f,
7984                        dl.qtype,
7985                        dl.row_bytes,
7986                    )?
7987                }
7988            } else if use_cache {
7989                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
7990                if grouped_q8 {
7991                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
7992                    let gate = e.with_moe_cache(max_block, |cache, eng| {
7993                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
7994                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
7995                        eng.qmatvec_expert_q8(
7996                            cache.buf(slot),
7997                            0..gl.len,
7998                            &zq,
7999                            &zd,
8000                            m_e,
8001                            m.gate_exps.in_f,
8002                            m.gate_exps.out_f,
8003                            gl.qtype,
8004                            gl.row_bytes,
8005                        )
8006                    })?;
8007                    let up = e.with_moe_cache(max_block, |cache, eng| {
8008                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8009                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8010                        eng.qmatvec_expert_q8(
8011                            cache.buf(slot),
8012                            0..ul.len,
8013                            &zq,
8014                            &zd,
8015                            m_e,
8016                            m.up_exps.in_f,
8017                            m.up_exps.out_f,
8018                            ul.qtype,
8019                            ul.row_bytes,
8020                        )
8021                    })?;
8022                    let mut act = e.uninit(m_e * n_ff_exp)?;
8023                    Self::ffn_act_lim(
8024                        e,
8025                        cfg,
8026                        &gate,
8027                        &up,
8028                        m.gate_exps.macro_scale(ex),
8029                        m.up_exps.macro_scale(ex),
8030                        lim_exp,
8031                        &mut act,
8032                        m_e * n_ff_exp,
8033                    )?;
8034                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8035                    e.with_moe_cache(max_block, |cache, eng| {
8036                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8037                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8038                        eng.qmatvec_expert_q8(
8039                            cache.buf(slot),
8040                            0..dl.len,
8041                            &aq2,
8042                            &ad2,
8043                            m_e,
8044                            m.down_exps.in_f,
8045                            m.down_exps.out_f,
8046                            dl.qtype,
8047                            dl.row_bytes,
8048                        )
8049                    })?
8050                } else {
8051                    let gate = e.with_moe_cache(max_block, |cache, eng| {
8052                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
8053                        let slot = cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
8054                        eng.qmatvec_view(
8055                            cache.buf(slot),
8056                            0..gl.len,
8057                            &gv,
8058                            m_e,
8059                            m.gate_exps.in_f,
8060                            m.gate_exps.out_f,
8061                            gl.qtype,
8062                            gl.row_bytes,
8063                        )
8064                    })?;
8065                    let up = e.with_moe_cache(max_block, |cache, eng| {
8066                        let id = BlockId::new(il, PROJ_UP, ex as u16);
8067                        let slot = cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
8068                        eng.qmatvec_view(
8069                            cache.buf(slot),
8070                            0..ul.len,
8071                            &gv,
8072                            m_e,
8073                            m.up_exps.in_f,
8074                            m.up_exps.out_f,
8075                            ul.qtype,
8076                            ul.row_bytes,
8077                        )
8078                    })?;
8079                    let mut act = e.uninit(m_e * n_ff_exp)?;
8080                    Self::ffn_act_lim(
8081                        e,
8082                        cfg,
8083                        &gate,
8084                        &up,
8085                        m.gate_exps.macro_scale(ex),
8086                        m.up_exps.macro_scale(ex),
8087                        lim_exp,
8088                        &mut act,
8089                        m_e * n_ff_exp,
8090                    )?;
8091                    let actv = act.slice(0..m_e * n_ff_exp);
8092                    e.with_moe_cache(max_block, |cache, eng| {
8093                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
8094                        let slot = cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
8095                        eng.qmatvec_view(
8096                            cache.buf(slot),
8097                            0..dl.len,
8098                            &actv,
8099                            m_e,
8100                            m.down_exps.in_f,
8101                            m.down_exps.out_f,
8102                            dl.qtype,
8103                            dl.row_bytes,
8104                        )
8105                    })?
8106                }
8107            } else {
8108                let sg = scratch_g.as_mut().unwrap();
8109                let su = scratch_u.as_mut().unwrap();
8110                let sd = scratch_d.as_mut().unwrap();
8111                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
8112                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
8113                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
8114                if grouped_q8 {
8115                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
8116                    let gate = e.qmatvec_expert_q8(
8117                        sg,
8118                        0..gl.len,
8119                        &zq,
8120                        &zd,
8121                        m_e,
8122                        m.gate_exps.in_f,
8123                        m.gate_exps.out_f,
8124                        gl.qtype,
8125                        gl.row_bytes,
8126                    )?;
8127                    let up = e.qmatvec_expert_q8(
8128                        su,
8129                        0..ul.len,
8130                        &zq,
8131                        &zd,
8132                        m_e,
8133                        m.up_exps.in_f,
8134                        m.up_exps.out_f,
8135                        ul.qtype,
8136                        ul.row_bytes,
8137                    )?;
8138                    let mut act = e.uninit(m_e * n_ff_exp)?;
8139                    Self::ffn_act_lim(
8140                        e,
8141                        cfg,
8142                        &gate,
8143                        &up,
8144                        m.gate_exps.macro_scale(ex),
8145                        m.up_exps.macro_scale(ex),
8146                        lim_exp,
8147                        &mut act,
8148                        m_e * n_ff_exp,
8149                    )?;
8150                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
8151                    e.qmatvec_expert_q8(
8152                        sd,
8153                        0..dl.len,
8154                        &aq2,
8155                        &ad2,
8156                        m_e,
8157                        m.down_exps.in_f,
8158                        m.down_exps.out_f,
8159                        dl.qtype,
8160                        dl.row_bytes,
8161                    )?
8162                } else {
8163                    let gate = e.qmatvec_view(
8164                        sg,
8165                        0..gl.len,
8166                        &gv,
8167                        m_e,
8168                        m.gate_exps.in_f,
8169                        m.gate_exps.out_f,
8170                        gl.qtype,
8171                        gl.row_bytes,
8172                    )?;
8173                    let up = e.qmatvec_view(
8174                        su,
8175                        0..ul.len,
8176                        &gv,
8177                        m_e,
8178                        m.up_exps.in_f,
8179                        m.up_exps.out_f,
8180                        ul.qtype,
8181                        ul.row_bytes,
8182                    )?;
8183                    let mut act = e.uninit(m_e * n_ff_exp)?;
8184                    Self::ffn_act_lim(
8185                        e,
8186                        cfg,
8187                        &gate,
8188                        &up,
8189                        m.gate_exps.macro_scale(ex),
8190                        m.up_exps.macro_scale(ex),
8191                        lim_exp,
8192                        &mut act,
8193                        m_e * n_ff_exp,
8194                    )?;
8195                    let actv = act.slice(0..m_e * n_ff_exp);
8196                    e.qmatvec_view(
8197                        sd,
8198                        0..dl.len,
8199                        &actv,
8200                        m_e,
8201                        m.down_exps.in_f,
8202                        m.down_exps.out_f,
8203                        dl.qtype,
8204                        dl.row_bytes,
8205                    )?
8206                }
8207            };
8208
8209            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
8210            e.scatter_slot(
8211                &y,
8212                &tok_idx_d,
8213                &slot_idx_d,
8214                &weight_d,
8215                &mut slot_buf,
8216                &mut wbuf,
8217                n_embd,
8218                n_used,
8219                m_e,
8220            )?;
8221        }
8222
8223        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
8224        let mut moe_out = e.zeros(t * n_embd)?;
8225        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
8226
8227        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
8228        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
8229            m_dist.sort_unstable();
8230            let active = m_dist.len();
8231            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
8232            let median = m_dist[active / 2];
8233            let max_m = *m_dist.last().unwrap();
8234            let min_m = m_dist[0];
8235            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
8236            println!(
8237                "moe-grouped il={il} t={t} active={active}/{n_expert} \
8238                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
8239                      above_gemm_threshold(>=16)={above16}/{active}"
8240            );
8241        }
8242
8243        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
8244        Ok(moe_out)
8245    }
8246
8247    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
8248    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
8249    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
8250    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
8251    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
8252    /// expert-sum order identical to the sequential path.
8253    pub(crate) fn moe_ffn_lockstep(
8254        &self,
8255        e: &Engine,
8256        m: &MoeWeights,
8257        zbatch: &CudaSlice<f32>,
8258        mrows: usize,
8259        il: u16,
8260        max_block: usize,
8261    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8262        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
8263        let cfg = &self.cfg;
8264        let moe = cfg.moe.as_ref().unwrap();
8265        let n_embd = cfg.n_embd as usize;
8266        let n_expert = moe.expert_count as usize;
8267        let n_used = moe.expert_used_count as usize;
8268        let n_ff_exp = moe.expert_ff_length as usize;
8269        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
8270        let lim_exp = cfg.clamp_exp_at(il as u32);
8271        let lim_shexp = cfg.clamp_shexp_at(il as u32);
8272
8273        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
8274        if let Some(sig) = cfg.sigmoid_router() {
8275            Self::trace_sigmoid_router_logits(e, il, mrows, n_expert, n_used, &logits, m, sig)?;
8276        }
8277        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
8278            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
8279        } else {
8280            Self::moe_route_cfg(
8281                e,
8282                &logits,
8283                mrows,
8284                n_expert,
8285                n_used,
8286                m.active_experts.as_deref(),
8287            )?
8288        };
8289        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
8290
8291        // Residency split at whole-expert granularity against the (frozen) cache.
8292        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
8293            Ok((0..n_expert)
8294                .map(|ex| {
8295                    [PROJ_GATE, PROJ_UP, PROJ_DOWN]
8296                        .into_iter()
8297                        .all(|p| c.resident(BlockId::new(il, p, ex as u16)).is_some())
8298                })
8299                .collect())
8300        })?;
8301
8302        struct Group {
8303            rows: Vec<i32>,
8304            slots: Vec<i32>,
8305            weights: Vec<f32>,
8306        }
8307        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
8308        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
8309        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
8310            Default::default();
8311        for row in 0..mrows {
8312            for j in 0..n_used {
8313                let ex = sel_all[row * n_used + j] as usize;
8314                let w = w_all[row * n_used + j];
8315                if resident_expert[ex] {
8316                    let group = groups.entry(ex).or_insert_with(|| Group {
8317                        rows: Vec::new(),
8318                        slots: Vec::new(),
8319                        weights: Vec::new(),
8320                    });
8321                    group.rows.push(row as i32);
8322                    group.slots.push(j as i32);
8323                    group.weights.push(w);
8324                } else {
8325                    crate::cpu_experts::record_incomplete_gpu_residency(0);
8326                    cpu_rows[row].push((ex, w));
8327                    cpu_by_expert.entry(ex).or_default().push((row, w));
8328                }
8329            }
8330        }
8331
8332        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
8333        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
8334        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
8335        // order per row differs from the sequential single-call chunk — part of the
8336        // documented lockstep numeric class.
8337        let host_rows = e.dtoh(zbatch)?;
8338        let rows_ok = crate::cpu_experts::rows_supported();
8339        enum CpuPart {
8340            Single { row: usize },
8341            Rows { rows: Vec<usize> },
8342        }
8343        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
8344        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
8345        if rows_ok {
8346            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
8347                .into_iter()
8348                .filter(|(_, rows)| rows.len() >= 2)
8349                .collect();
8350            shared.sort_by_key(|(ex, _)| *ex);
8351            for (ex, mut row_weights) in shared {
8352                row_weights.sort_by_key(|(row, _)| *row);
8353                let inputs: Vec<(&[f32], f32)> = row_weights
8354                    .iter()
8355                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
8356                    .collect();
8357                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
8358                    .map_err(std::io::Error::other)?;
8359                for &(row, _) in &row_weights {
8360                    rows_served.insert((row, ex));
8361                }
8362                tickets.push((
8363                    CpuPart::Rows {
8364                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
8365                    },
8366                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
8367                ));
8368            }
8369        }
8370        for (row, selected) in cpu_rows.iter().enumerate() {
8371            let leftover: Vec<(usize, f32)> = selected
8372                .iter()
8373                .copied()
8374                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
8375                .collect();
8376            if leftover.is_empty() {
8377                continue;
8378            }
8379            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
8380            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
8381                .map_err(std::io::Error::other)?;
8382            tickets.push((
8383                CpuPart::Single { row },
8384                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
8385            ));
8386        }
8387
8388        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
8389        let mut wbuf = e.zeros(mrows * n_used)?;
8390        let mut order: Vec<usize> = groups.keys().copied().collect();
8391        order.sort_by(|&a, &b| {
8392            groups[&b]
8393                .rows
8394                .len()
8395                .cmp(&groups[&a].rows.len())
8396                .then(a.cmp(&b))
8397        });
8398        for &ex in &order {
8399            let group = &groups[&ex];
8400            let m_e = group.rows.len();
8401            let gl = m.gate_exps.expert_layout(ex);
8402            let ul = m.up_exps.expert_layout(ex);
8403            let dl = m.down_exps.expert_layout(ex);
8404            let row_idx_d = e.htod_i32(&group.rows)?;
8405            let slot_idx_d = e.htod_i32(&group.slots)?;
8406            let dmac = m.down_exps.macro_scale(ex);
8407            let weight_d = if dmac == 1.0 {
8408                e.htod(&group.weights)?
8409            } else {
8410                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
8411                e.htod(&scaled)?
8412            };
8413            let mut gathered = e.zeros(m_e * n_embd)?;
8414            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
8415            let gv = gathered.slice(0..m_e * n_embd);
8416            let gate = e.with_moe_cache(max_block, |c, eng| {
8417                let slot = c
8418                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
8419                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8420                eng.qmatvec_view(
8421                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8422                    0..gl.len,
8423                    &gv,
8424                    m_e,
8425                    m.gate_exps.in_f,
8426                    m.gate_exps.out_f,
8427                    gl.qtype,
8428                    gl.row_bytes,
8429                )
8430            })?;
8431            let up = e.with_moe_cache(max_block, |c, eng| {
8432                let slot = c
8433                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
8434                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8435                eng.qmatvec_view(
8436                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8437                    0..ul.len,
8438                    &gv,
8439                    m_e,
8440                    m.up_exps.in_f,
8441                    m.up_exps.out_f,
8442                    ul.qtype,
8443                    ul.row_bytes,
8444                )
8445            })?;
8446            let mut act = e.zeros(m_e * n_ff_exp)?;
8447            Self::ffn_act_lim(
8448                e,
8449                cfg,
8450                &gate,
8451                &up,
8452                m.gate_exps.macro_scale(ex),
8453                m.up_exps.macro_scale(ex),
8454                lim_exp,
8455                &mut act,
8456                m_e * n_ff_exp,
8457            )?;
8458            let actv = act.slice(0..m_e * n_ff_exp);
8459            let y = e.with_moe_cache(max_block, |c, eng| {
8460                let slot = c
8461                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
8462                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
8463                eng.qmatvec_view(
8464                    c.buf(crate::moe_cache::DispatchSlot::Resident(slot)),
8465                    0..dl.len,
8466                    &actv,
8467                    m_e,
8468                    m.down_exps.in_f,
8469                    m.down_exps.out_f,
8470                    dl.qtype,
8471                    dl.row_bytes,
8472                )
8473            })?;
8474            e.scatter_slot(
8475                &y,
8476                &row_idx_d,
8477                &slot_idx_d,
8478                &weight_d,
8479                &mut slot_buf,
8480                &mut wbuf,
8481                n_embd,
8482                n_used,
8483                m_e,
8484            )?;
8485        }
8486        let mut moe_out = e.zeros(mrows * n_embd)?;
8487        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
8488
8489        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
8490        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
8491        for (part, ticket) in tickets {
8492            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
8493            let mut add_row = |row: usize, chunk: &[f32]| {
8494                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
8495                for (accumulator, value) in sum.iter_mut().zip(chunk) {
8496                    *accumulator += value;
8497                }
8498            };
8499            match part {
8500                CpuPart::Single { row } => add_row(row, &cpu_output),
8501                CpuPart::Rows { rows } => {
8502                    for (slot, row) in rows.into_iter().enumerate() {
8503                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
8504                    }
8505                }
8506            }
8507        }
8508        for (row, sum) in row_sums.into_iter().enumerate() {
8509            let Some(sum) = sum else { continue };
8510            let cpu_output = e.htod(&sum)?;
8511            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
8512            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
8513        }
8514
8515        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
8516            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
8517        {
8518            let n_ff_sh = gate_shexp.out_features();
8519            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
8520            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
8521            let mut sa = e.zeros(mrows * n_ff_sh)?;
8522            Self::ffn_act_lim(
8523                e,
8524                cfg,
8525                &sg_gate,
8526                &sg_up,
8527                1.0,
8528                1.0,
8529                lim_shexp,
8530                &mut sa,
8531                mrows * n_ff_sh,
8532            )?;
8533            let sh = e.matmul(down_shexp, &sa, mrows)?;
8534            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
8535            // decode matches the single-sequence decode chain bit-for-bit.
8536            let g = match &m.gate_inp_shexp {
8537                Some(gate_inp_shexp) => {
8538                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
8539                }
8540                None => e.htod(&vec![1.0f32; mrows])?,
8541            };
8542            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
8543        }
8544
8545        Ok(moe_out)
8546    }
8547}
8548
8549// ============================ gemma4 (R8 verified wiring) ==================================
8550// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
8551// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
8552// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
8553// gemma variants after the correctness gate).
8554impl HybridModel {
8555    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
8556    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
8557        let g = self.cfg.gemma4.as_ref().unwrap();
8558        let swa = g.swa_pattern[il];
8559        let hd = if swa {
8560            g.key_length_swa
8561        } else {
8562            g.key_length_global
8563        } as usize;
8564        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
8565        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
8566        // rows exact (softmax over one element) while every later position drifted).
8567        (
8568            hd,
8569            g.head_count_kv[il] as usize,
8570            self.cfg.n_head as usize,
8571            if swa {
8572                g.rope_base_swa
8573            } else {
8574                g.rope_base_global
8575            },
8576            1.0,
8577            swa,
8578        )
8579    }
8580
8581    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
8582    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
8583    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
8584    fn gemma4_suppress(
8585        &self,
8586        e: &Engine,
8587        ld: &mut CudaSlice<f32>,
8588        t: usize,
8589    ) -> Result<(), Box<dyn std::error::Error>> {
8590        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
8591            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
8592            // stage as primary, and this tail runs only after the last stage). The assert turns
8593            // that argued invariant into a checked one: any topology violating primary==head
8594            // trips here in debug instead of silently peer-reading a device-0 buffer.
8595            #[cfg(debug_assertions)]
8596            crate::debug_assert_tensor_stream_device(
8597                ids,
8598                &e.stream(),
8599                "gemma4_suppress.suppress_d",
8600            );
8601            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
8602        }
8603        Ok(())
8604    }
8605
8606    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
8607    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
8608    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
8609    /// only (v0): attends within `tokens` via the f32 sdpa.
8610    fn gemma4_attn_prime(
8611        &self,
8612        e: &Engine,
8613        fa: &crate::hybrid::FullAttnLayer,
8614        il: usize,
8615        h: &CudaSlice<f32>,
8616        pos_d: &CudaSlice<i32>,
8617        t: usize,
8618        cache: Option<&mut Cache>,
8619    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8620        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8621        let eps = self.cfg.rms_eps;
8622        let aux = self.gemma4_aux.as_ref().unwrap();
8623        let ones = aux.ones(e);
8624        #[cfg(debug_assertions)]
8625        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_attn_prime.ones");
8626
8627        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
8628        // (h stays borrowed across the triple, so the cache key can't go stale).
8629        e.mmq_act_begin();
8630        let q0 = e.matmul(&fa.wq, h, t)?; // [t, nh*hd]
8631        let k0 = e.matmul(&fa.wk, h, t)?; // [t, nkv*hd]
8632        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
8633        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
8634        let v0 = if swa {
8635            e.matmul(&fa.wv, h, t)?
8636        } else {
8637            e.clone_dtod(&k0)?
8638        };
8639
8640        let mut q = e.uninit(t * nh * hd)?;
8641        let mut k = e.uninit(t * nkv * hd)?;
8642        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
8643        let mut v = e.uninit(t * nkv * hd)?;
8644        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
8645        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
8646        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
8647        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8648        let emit = t >= 16
8649            && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
8650            && *EMIT.get_or_init(|| {
8651                std::env::var("MEMRA_FA_EMIT")
8652                    .map(|s| s != "0")
8653                    .unwrap_or(true)
8654            });
8655        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
8656        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8657        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
8658        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
8659        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
8660        let v_f16 = emit
8661            && crate::fa_f16pv_on()
8662            && match hd {
8663                512 => true,
8664                256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
8665                _ => false,
8666            };
8667        if emit {
8668            e.rms_norm_qkv_w4b(
8669                &q0,
8670                &k0,
8671                &v0,
8672                fa.q_norm.float_data(),
8673                fa.k_norm.float_data(),
8674                ones,
8675                &mut q,
8676                &mut k,
8677                &mut v,
8678                &mut vb,
8679                hd,
8680                nh * t,
8681                nkv * t,
8682                eps,
8683                v_f16,
8684            )?;
8685        } else {
8686            e.rms_norm_qkv(
8687                &q0,
8688                &k0,
8689                &v0,
8690                fa.q_norm.float_data(),
8691                fa.k_norm.float_data(),
8692                ones,
8693                &mut q,
8694                &mut k,
8695                &mut v,
8696                hd,
8697                nh * t,
8698                nkv * t,
8699                eps,
8700            )?;
8701        }
8702
8703        let ff = if swa {
8704            None
8705        } else {
8706            Some(
8707                aux.rope_freqs(e)
8708                    .expect("gemma4 global rope needs rope_freqs.weight"),
8709            )
8710        };
8711        #[cfg(debug_assertions)]
8712        if let Some(ff) = ff {
8713            crate::debug_assert_tensor_stream_device(
8714                ff,
8715                &e.stream(),
8716                "gemma4_attn_prime.rope_freqs",
8717            );
8718        }
8719        if emit {
8720            e.rope_neox2_bf16e(
8721                &mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff,
8722            )?;
8723        } else {
8724            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
8725        }
8726
8727        if let Some(cache) = cache {
8728            let kvl = cache.kv[il].as_mut().unwrap();
8729            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
8730            e.append_kv_quantized_rows(
8731                &k,
8732                &v,
8733                &mut kvl.k,
8734                &mut kvl.v,
8735                kvl.len,
8736                t,
8737                kvl.kv_dim_k,
8738                kvl.kv_dim_v,
8739                kvl.k_tok_bytes,
8740                kvl.v_tok_bytes,
8741                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
8742            )?;
8743            kvl.len += t;
8744        }
8745        let mut attn = e.zeros(t * nh * hd)?;
8746        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
8747        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
8748        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
8749        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8750        if swa && t > win {
8751            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8752                if emit {
8753                    e.fa_prefill_w_pre(
8754                        &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, win, v_f16,
8755                    )?;
8756                } else {
8757                    e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8758                }
8759            } else {
8760                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
8761            }
8762        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
8763            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8764        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
8765            if emit {
8766                e.fa_prefill_hd512_pre(
8767                    &qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t, scale, true, v_f16,
8768                )?;
8769            } else {
8770                e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8771            }
8772        } else {
8773            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
8774        }
8775        Ok(e.matmul(&fa.wo, &attn, t)?)
8776    }
8777
8778    /// Back-compat wrapper (pure prefill, no cache).
8779    fn gemma4_attn(
8780        &self,
8781        e: &Engine,
8782        fa: &crate::hybrid::FullAttnLayer,
8783        il: usize,
8784        h: &CudaSlice<f32>,
8785        pos_d: &CudaSlice<i32>,
8786        t: usize,
8787    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8788        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
8789    }
8790
8791    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
8792    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
8793    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
8794    /// the q8z epilogue is quantize_q8_1 verbatim).
8795    fn gemma4_moe_q8(
8796        &self,
8797        e: &Engine,
8798        m: &crate::hybrid::MoeWeights,
8799        bits: &crate::hybrid::Gemma4MoeBits,
8800        mq: &(CudaSlice<i8>, CudaSlice<f32>),
8801        router_in: &CudaSlice<f32>,
8802        t: usize,
8803    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8804        let cfg = &self.cfg;
8805        let moe = cfg.moe.as_ref().unwrap();
8806        let n_embd = cfg.n_embd as usize;
8807        let n_expert = moe.expert_count as usize;
8808        let n_used = moe.expert_used_count as usize;
8809        let n_ff_exp = moe.expert_ff_length as usize;
8810        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
8811        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
8812        // the pair's 12us is kernel time, not launch gaps.
8813        let logits = if crate::router_kernel_on() {
8814            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8815        } else {
8816            e.matmul(&m.gate_inp, router_in, t)?
8817        };
8818        let dev = m.dev_exps.as_ref().unwrap();
8819        let (sel_d, w_d) =
8820            e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8821        let (zq, zd) = mq;
8822        if t == 1 {
8823            let selv = sel_d.slice(0..n_used);
8824            let wv = w_d.slice(0..n_used);
8825            let act = e.moe_gate_up_gelu8_dev_q8(
8826                &dev.ptr_row,
8827                &selv,
8828                zq,
8829                zd,
8830                n_embd,
8831                n_ff_exp,
8832                n_used,
8833                n_expert,
8834                m.gate_exps.qtype,
8835                m.up_exps.qtype,
8836                m.gate_exps.row_bytes,
8837                m.up_exps.row_bytes,
8838            )?;
8839            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8840            let mut moe_out = e.uninit(n_embd)?;
8841            e.moe_down8_fma_dev_q8(
8842                &dev.ptr_row,
8843                &selv,
8844                &wv,
8845                &aq2,
8846                &ad2,
8847                &mut moe_out.slice_mut(0..n_embd),
8848                n_ff_exp,
8849                n_embd,
8850                n_used,
8851                n_expert,
8852                m.down_exps.qtype,
8853                m.down_exps.row_bytes,
8854            )?;
8855            return Ok(moe_out);
8856        }
8857        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8858        let act = if csr {
8859            e.moe_gate_up_gelu8_dev_q8_csr(
8860                &dev.ptr_row,
8861                &sel_d,
8862                zq,
8863                zd,
8864                t * n_used,
8865                n_embd,
8866                n_ff_exp,
8867                n_used,
8868                n_expert,
8869                m.gate_exps.qtype,
8870                m.up_exps.qtype,
8871                m.gate_exps.row_bytes,
8872                m.up_exps.row_bytes,
8873            )?
8874        } else {
8875            e.moe_gate_up_gelu8_dev_q8_rows(
8876                &dev.ptr_row,
8877                &sel_d,
8878                zq,
8879                zd,
8880                t,
8881                n_embd,
8882                n_ff_exp,
8883                n_used,
8884                n_expert,
8885                m.gate_exps.qtype,
8886                m.up_exps.qtype,
8887                m.gate_exps.row_bytes,
8888                m.up_exps.row_bytes,
8889            )?
8890        };
8891        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
8892        let mut moe_out = e.uninit(t * n_embd)?;
8893        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
8894        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
8895        e.moe_down8_fma_dev_q8_rows_g(
8896            &dev.ptr_row,
8897            &sel_d,
8898            &w_d,
8899            &aq2,
8900            &ad2,
8901            &mut moe_out,
8902            t,
8903            n_ff_exp,
8904            n_embd,
8905            n_used,
8906            n_expert,
8907            m.down_exps.qtype,
8908            m.down_exps.row_bytes,
8909        )?;
8910        Ok(moe_out)
8911    }
8912
8913    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
8914    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
8915    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
8916    fn gemma4_moe(
8917        &self,
8918        e: &Engine,
8919        m: &crate::hybrid::MoeWeights,
8920        bits: &crate::hybrid::Gemma4MoeBits,
8921        moe_in: &CudaSlice<f32>,
8922        router_in: &CudaSlice<f32>,
8923        t: usize,
8924    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8925        let cfg = &self.cfg;
8926        let moe = cfg.moe.as_ref().unwrap();
8927        let n_embd = cfg.n_embd as usize;
8928        let n_expert = moe.expert_count as usize;
8929        let n_used = moe.expert_used_count as usize;
8930        let n_ff_exp = moe.expert_ff_length as usize;
8931
8932        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
8933        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
8934        // batched matmul only at real prefill.
8935        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
8936            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
8937        } else {
8938            e.matmul(&m.gate_inp, router_in, t)?
8939        };
8940
8941        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
8942        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
8943        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
8944        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
8945        if t < PRIME_MIN_T
8946            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
8947            && expert_dp4a_supported(m.gate_exps.qtype)
8948            && expert_dp4a_supported(m.up_exps.qtype)
8949            && expert_dp4a_supported(m.down_exps.qtype)
8950            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
8951        {
8952            let dev = m.dev_exps.as_ref().unwrap();
8953            let (sel_d, w_d) =
8954                e.moe_router_topk_scaled(&logits, t, n_expert, n_used, &bits.per_expert_scale_d)?;
8955            if t == 1 {
8956                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
8957                let selv = sel_d.slice(0..n_used);
8958                let wv = w_d.slice(0..n_used);
8959                let act = e.moe_gate_up_gelu8_dev_q8(
8960                    &dev.ptr_row,
8961                    &selv,
8962                    &zq,
8963                    &zd,
8964                    n_embd,
8965                    n_ff_exp,
8966                    n_used,
8967                    n_expert,
8968                    m.gate_exps.qtype,
8969                    m.up_exps.qtype,
8970                    m.gate_exps.row_bytes,
8971                    m.up_exps.row_bytes,
8972                )?;
8973                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
8974                let mut moe_out = e.uninit(n_embd)?;
8975                e.moe_down8_fma_dev_q8(
8976                    &dev.ptr_row,
8977                    &selv,
8978                    &wv,
8979                    &aq2,
8980                    &ad2,
8981                    &mut moe_out.slice_mut(0..n_embd),
8982                    n_ff_exp,
8983                    n_embd,
8984                    n_used,
8985                    n_expert,
8986                    m.down_exps.qtype,
8987                    m.down_exps.row_bytes,
8988                )?;
8989                return Ok(moe_out);
8990            }
8991            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
8992            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
8993            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
8994            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
8995            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
8996            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
8997            let act = if csr {
8998                e.moe_gate_up_gelu8_dev_q8_csr(
8999                    &dev.ptr_row,
9000                    &sel_d,
9001                    &zq,
9002                    &zd,
9003                    t * n_used,
9004                    n_embd,
9005                    n_ff_exp,
9006                    n_used,
9007                    n_expert,
9008                    m.gate_exps.qtype,
9009                    m.up_exps.qtype,
9010                    m.gate_exps.row_bytes,
9011                    m.up_exps.row_bytes,
9012                )?
9013            } else {
9014                e.moe_gate_up_gelu8_dev_q8_rows(
9015                    &dev.ptr_row,
9016                    &sel_d,
9017                    &zq,
9018                    &zd,
9019                    t,
9020                    n_embd,
9021                    n_ff_exp,
9022                    n_used,
9023                    n_expert,
9024                    m.gate_exps.qtype,
9025                    m.up_exps.qtype,
9026                    m.gate_exps.row_bytes,
9027                    m.up_exps.row_bytes,
9028                )?
9029            };
9030            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
9031            let mut moe_out = e.uninit(t * n_embd)?;
9032            e.moe_down8_fma_dev_q8_rows_g(
9033                &dev.ptr_row,
9034                &sel_d,
9035                &w_d,
9036                &aq2,
9037                &ad2,
9038                &mut moe_out,
9039                t,
9040                n_ff_exp,
9041                n_embd,
9042                n_used,
9043                n_expert,
9044                m.down_exps.qtype,
9045                m.down_exps.row_bytes,
9046            )?;
9047            return Ok(moe_out);
9048        }
9049
9050        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
9051        for (i, &sx) in sel_all.iter().enumerate() {
9052            w_all[i] *= bits.per_expert_scale[sx as usize];
9053        }
9054
9055        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
9056        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
9057        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
9058        if t >= PRIME_MIN_T
9059            && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9060            && expert_dp4a_supported(m.gate_exps.qtype)
9061            && expert_dp4a_supported(m.up_exps.qtype)
9062            && expert_dp4a_supported(m.down_exps.qtype)
9063            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0")
9064        {
9065            let dev = m.dev_exps.as_ref().unwrap();
9066            let n_pairs = t * n_used;
9067            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
9068            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
9069            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
9070            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
9071            let pt = e.htod_i32(&pair_tok)?;
9072            let pw = e.htod(&w_all)?;
9073            let toff = e.htod_i32(&tok_off)?;
9074            let tids = e.htod_i32(&tok_ids)?;
9075            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
9076            for p in 0..n_pairs {
9077                by_ex[pair_ex[p] as usize].push(p as i32);
9078            }
9079            let mut ex_ids: Vec<i32> = Vec::new();
9080            let mut ex_off: Vec<i32> = vec![0];
9081            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
9082            for (ex, list) in by_ex.iter().enumerate() {
9083                if list.is_empty() {
9084                    continue;
9085                }
9086                ex_ids.push(ex as i32);
9087                ex_pairs.extend_from_slice(list);
9088                ex_off.push(ex_pairs.len() as i32);
9089            }
9090            let n_active = ex_ids.len();
9091            let exi = e.htod_i32(&ex_ids)?;
9092            let exo = e.htod_i32(&ex_off)?;
9093            let exp_d = e.htod_i32(&ex_pairs)?;
9094            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
9095            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
9096            // end-to-end (gelu is elementwise), one row permute before the scatter. The
9097            // ragged down k (704) needs no padding here — cublas takes any k.
9098            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
9099            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
9100            // Hopper default — see moe_f16g_gemma_on.
9101            if crate::moe_f16g_gemma_on()
9102                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
9103                && f16g_proj_ok(m.up_exps.qtype, n_embd)
9104                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp)
9105            {
9106                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
9107                let csr_tok_d = e.htod_i32(&csr_tok)?;
9108                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
9109                let g_csr = e.moe_f16_grouped(
9110                    &dev.ptr_row,
9111                    0,
9112                    n_expert,
9113                    &exi,
9114                    &ex_off,
9115                    &exo,
9116                    &z_f16,
9117                    &z_s,
9118                    n_embd,
9119                    n_ff_exp,
9120                    n_active,
9121                    n_pairs,
9122                    m.gate_exps.qtype,
9123                    m.gate_exps.row_bytes,
9124                )?;
9125                let u_csr = e.moe_f16_grouped(
9126                    &dev.ptr_row,
9127                    1,
9128                    n_expert,
9129                    &exi,
9130                    &ex_off,
9131                    &exo,
9132                    &z_f16,
9133                    &z_s,
9134                    n_embd,
9135                    n_ff_exp,
9136                    n_active,
9137                    n_pairs,
9138                    m.up_exps.qtype,
9139                    m.up_exps.row_bytes,
9140                )?;
9141                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
9142                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
9143                let d_csr = e.moe_f16_grouped(
9144                    &dev.ptr_row,
9145                    2,
9146                    n_expert,
9147                    &exi,
9148                    &ex_off,
9149                    &exo,
9150                    &a_f16,
9151                    &a_s,
9152                    n_ff_exp,
9153                    n_embd,
9154                    n_active,
9155                    n_pairs,
9156                    m.down_exps.qtype,
9157                    m.down_exps.row_bytes,
9158                )?;
9159                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
9160                let mut moe_out = e.uninit(t * n_embd)?;
9161                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9162                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
9163                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
9164                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
9165                    eprintln!(
9166                        "[f16g-debug] post-permute bad={} post-scatter bad={}",
9167                        scan(&yd),
9168                        scan(&mo)
9169                    );
9170                }
9171                return Ok(moe_out);
9172            }
9173            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
9174            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
9175            let mma =
9176                n_embd % 256 == 0 && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
9177            let (gate, up) = if mma {
9178                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
9179                (
9180                    e.mmq_iq_experts(
9181                        &dev.ptr_row,
9182                        0,
9183                        n_expert,
9184                        &exi,
9185                        &exo,
9186                        &exp_d,
9187                        &pt,
9188                        &z_scr,
9189                        n_embd,
9190                        n_ff_exp,
9191                        n_active,
9192                        n_pairs,
9193                        t,
9194                        m.gate_exps.qtype,
9195                        m.gate_exps.row_bytes,
9196                    )?,
9197                    e.mmq_iq_experts(
9198                        &dev.ptr_row,
9199                        1,
9200                        n_expert,
9201                        &exi,
9202                        &exo,
9203                        &exp_d,
9204                        &pt,
9205                        &z_scr,
9206                        n_embd,
9207                        n_ff_exp,
9208                        n_active,
9209                        n_pairs,
9210                        t,
9211                        m.up_exps.qtype,
9212                        m.up_exps.row_bytes,
9213                    )?,
9214                )
9215            } else {
9216                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
9217                (
9218                    e.moe_pairs_matvec_q8_dec(
9219                        &dev.ptr_row,
9220                        0,
9221                        &exi,
9222                        &exo,
9223                        &exp_d,
9224                        &pt,
9225                        &zq,
9226                        &zd,
9227                        n_embd,
9228                        n_ff_exp,
9229                        n_expert,
9230                        n_active,
9231                        n_pairs,
9232                        m.gate_exps.qtype,
9233                        m.gate_exps.row_bytes,
9234                    )?,
9235                    e.moe_pairs_matvec_q8_dec(
9236                        &dev.ptr_row,
9237                        1,
9238                        &exi,
9239                        &exo,
9240                        &exp_d,
9241                        &pt,
9242                        &zq,
9243                        &zd,
9244                        n_embd,
9245                        n_ff_exp,
9246                        n_expert,
9247                        n_active,
9248                        n_pairs,
9249                        m.up_exps.qtype,
9250                        m.up_exps.row_bytes,
9251                    )?,
9252                )
9253            };
9254            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
9255            let pself = e.htod_i32(&pair_self)?;
9256            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
9257            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
9258            // to the 256-val superblock (768) while the act quantizer's zero padding
9259            // makes every padded-k product exactly zero (weight overread bytes multiply
9260            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
9261            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
9262            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
9263            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
9264            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
9265            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
9266            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
9267            let y_down = if mma {
9268                let in_pad = n_ff_exp.div_ceil(256) * 256;
9269                let a_scr = if crate::moe_fuse_actq_on() {
9270                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
9271                } else {
9272                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9273                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
9274                };
9275                e.mmq_iq_experts(
9276                    &dev.ptr_row,
9277                    2,
9278                    n_expert,
9279                    &exi,
9280                    &exo,
9281                    &exp_d,
9282                    &pself,
9283                    &a_scr,
9284                    in_pad,
9285                    n_embd,
9286                    n_active,
9287                    n_pairs,
9288                    n_pairs,
9289                    m.down_exps.qtype,
9290                    m.down_exps.row_bytes,
9291                )?
9292            } else {
9293                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
9294                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
9295                e.moe_pairs_matvec_q8_dec(
9296                    &dev.ptr_row,
9297                    2,
9298                    &exi,
9299                    &exo,
9300                    &exp_d,
9301                    &pself,
9302                    &aq2,
9303                    &ad2,
9304                    n_ff_exp,
9305                    n_embd,
9306                    n_expert,
9307                    n_active,
9308                    n_pairs,
9309                    m.down_exps.qtype,
9310                    m.down_exps.row_bytes,
9311                )?
9312            };
9313            let mut moe_out = e.uninit(t * n_embd)?;
9314            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
9315            return Ok(moe_out);
9316        }
9317
9318        let g_len = m.gate_exps.expert_stride;
9319        let u_len = m.up_exps.expert_stride;
9320        let d_len = m.down_exps.expert_stride;
9321        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
9322        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
9323        // the spill fallback.
9324        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
9325        let (mut sg, mut su, mut sd) = if dev.is_some() {
9326            (None, None, None)
9327        } else {
9328            (
9329                Some(e.alloc_u8_uninit(g_len)?),
9330                Some(e.alloc_u8_uninit(u_len)?),
9331                Some(e.alloc_u8_uninit(d_len)?),
9332            )
9333        };
9334        let mut moe_out = e.zeros(t * n_embd)?;
9335        for tok in 0..t {
9336            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
9337            let w = &w_all[tok * n_used..(tok + 1) * n_used];
9338            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
9339            for (j, &ex) in sel.iter().enumerate() {
9340                let ex = ex as usize;
9341                let gate = match dev {
9342                    Some(d) => e.qmatvec_view(
9343                        &d.gate,
9344                        ex * g_len..(ex + 1) * g_len,
9345                        &zt,
9346                        1,
9347                        m.gate_exps.in_f,
9348                        m.gate_exps.out_f,
9349                        m.gate_exps.qtype,
9350                        m.gate_exps.row_bytes,
9351                    )?,
9352                    None => {
9353                        let sg = sg.as_mut().unwrap();
9354                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
9355                        e.qmatvec_view(
9356                            sg,
9357                            0..g_len,
9358                            &zt,
9359                            1,
9360                            m.gate_exps.in_f,
9361                            m.gate_exps.out_f,
9362                            m.gate_exps.qtype,
9363                            m.gate_exps.row_bytes,
9364                        )?
9365                    }
9366                };
9367                let up = match dev {
9368                    Some(d) => e.qmatvec_view(
9369                        &d.up,
9370                        ex * u_len..(ex + 1) * u_len,
9371                        &zt,
9372                        1,
9373                        m.up_exps.in_f,
9374                        m.up_exps.out_f,
9375                        m.up_exps.qtype,
9376                        m.up_exps.row_bytes,
9377                    )?,
9378                    None => {
9379                        let su = su.as_mut().unwrap();
9380                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
9381                        e.qmatvec_view(
9382                            su,
9383                            0..u_len,
9384                            &zt,
9385                            1,
9386                            m.up_exps.in_f,
9387                            m.up_exps.out_f,
9388                            m.up_exps.qtype,
9389                            m.up_exps.row_bytes,
9390                        )?
9391                    }
9392                };
9393                let mut act = e.uninit(n_ff_exp)?;
9394                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
9395                let actv = act.slice(0..n_ff_exp);
9396                let y = match dev {
9397                    Some(d) => e.qmatvec_view(
9398                        &d.down,
9399                        ex * d_len..(ex + 1) * d_len,
9400                        &actv,
9401                        1,
9402                        m.down_exps.in_f,
9403                        m.down_exps.out_f,
9404                        m.down_exps.qtype,
9405                        m.down_exps.row_bytes,
9406                    )?,
9407                    None => {
9408                        let sd = sd.as_mut().unwrap();
9409                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
9410                        e.qmatvec_view(
9411                            sd,
9412                            0..d_len,
9413                            &actv,
9414                            1,
9415                            m.down_exps.in_f,
9416                            m.down_exps.out_f,
9417                            m.down_exps.qtype,
9418                            m.down_exps.row_bytes,
9419                        )?
9420                    }
9421                };
9422                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
9423                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
9424            }
9425        }
9426        Ok(moe_out)
9427    }
9428
9429    /// One gemma4 trunk layer (R8): x -> x_next.
9430    fn gemma4_layer(
9431        &self,
9432        e: &Engine,
9433        il: usize,
9434        layer: &crate::hybrid::HybridLayer,
9435        x: &CudaSlice<f32>,
9436        pos_d: &CudaSlice<i32>,
9437        t: usize,
9438    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9439        let n_embd = self.cfg.n_embd as usize;
9440        let eps = self.cfg.rms_eps;
9441
9442        let mut h = e.zeros(t * n_embd)?;
9443        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9444        let Mixer::Full(fa) = &layer.mixer else {
9445            panic!("gemma4 layer {il} not full-attn")
9446        };
9447        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
9448        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
9449        let mut cur = e.zeros(t * n_embd)?;
9450        e.rms_norm(
9451            &o,
9452            layer.post_attn_norm.float_data(),
9453            &mut cur,
9454            n_embd,
9455            t,
9456            eps,
9457        )?;
9458        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
9459    }
9460
9461    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
9462    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
9463    /// layer scale — shared verbatim by the prefill, decode and verify paths.
9464    fn gemma4_layer_tail_add(
9465        &self,
9466        e: &Engine,
9467        layer: &crate::hybrid::HybridLayer,
9468        cur: &CudaSlice<f32>,
9469        x: &CudaSlice<f32>,
9470        t: usize,
9471    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9472        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
9473    }
9474
9475    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
9476    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
9477    fn gemma4_layer_tail_add_n(
9478        &self,
9479        e: &Engine,
9480        layer: &crate::hybrid::HybridLayer,
9481        cur: &CudaSlice<f32>,
9482        x: &CudaSlice<f32>,
9483        t: usize,
9484        next_norm: Option<&CudaSlice<f32>>,
9485    ) -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
9486        let n_embd = self.cfg.n_embd as usize;
9487        let bits = layer.gemma4.as_ref().unwrap();
9488        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9489        let mut xn = e.uninit(t * n_embd)?;
9490        match next_norm {
9491            Some(w) => {
9492                let mut hn = e.uninit(t * n_embd)?;
9493                e.add_scale_rms_norm(
9494                    &sn,
9495                    &attn_out,
9496                    bits.layer_scale,
9497                    w,
9498                    &mut xn,
9499                    &mut hn,
9500                    n_embd,
9501                    t,
9502                    self.cfg.rms_eps,
9503                )?;
9504                Ok((xn, Some(hn)))
9505            }
9506            None => {
9507                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9508                Ok((xn, None))
9509            }
9510        }
9511    }
9512
9513    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
9514    /// norm — returns (sn, attn_out) for the closing add+scale variants.
9515    fn gemma4_layer_tail_core(
9516        &self,
9517        e: &Engine,
9518        layer: &crate::hybrid::HybridLayer,
9519        cur: &CudaSlice<f32>,
9520        x: &CudaSlice<f32>,
9521        t: usize,
9522    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9523        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
9524    }
9525
9526    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
9527    /// means `cur` is the RAW attention output and the dense entry runs
9528    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
9529    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
9530    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
9531    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
9532    fn gemma4_layer_tail_core_pn(
9533        &self,
9534        e: &Engine,
9535        layer: &crate::hybrid::HybridLayer,
9536        cur: &CudaSlice<f32>,
9537        x: &CudaSlice<f32>,
9538        t: usize,
9539        pre_norm: Option<&CudaSlice<f32>>,
9540        defer_post_norm: bool,
9541    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9542        let n_embd = self.cfg.n_embd as usize;
9543        let eps = self.cfg.rms_eps;
9544        let bits = layer.gemma4.as_ref().unwrap();
9545
9546        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
9547        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
9548        let Some(mbits) = bits.moe_bits.as_ref() else {
9549            let crate::hybrid::Ffn::Dense {
9550                ffn_gate,
9551                ffn_up,
9552                ffn_down,
9553            } = &layer.ffn
9554            else {
9555                panic!("gemma4 dense layer without Dense ffn")
9556            };
9557            let mut attn_out = e.uninit(t * n_embd)?;
9558            let mut zsh = e.uninit(t * n_embd)?;
9559            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
9560            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
9561            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9562            match pre_norm {
9563                Some(wa) if t == 1 => {
9564                    zpair = Some(e.rms_pre_add_rms_norm_q8z(
9565                        cur,
9566                        wa,
9567                        x,
9568                        bits.ffn_norm.float_data(),
9569                        &mut attn_out,
9570                        &mut zsh,
9571                        n_embd,
9572                        t,
9573                        eps,
9574                    )?);
9575                }
9576                Some(wa) => e.rms_pre_add_rms_norm(
9577                    cur,
9578                    wa,
9579                    x,
9580                    bits.ffn_norm.float_data(),
9581                    &mut attn_out,
9582                    &mut zsh,
9583                    n_embd,
9584                    t,
9585                    eps,
9586                )?,
9587                None => e.add_rms_norm(
9588                    cur,
9589                    x,
9590                    bits.ffn_norm.float_data(),
9591                    &mut attn_out,
9592                    &mut zsh,
9593                    n_embd,
9594                    t,
9595                    eps,
9596                )?,
9597            }
9598            let n_ff = ffn_gate.out_features();
9599            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
9600            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
9601            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
9602            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
9603            // rescue segment C — the megakernel front is closed for the dense tail.
9604            let (gate, up) = if t == 1 {
9605                let (zq, zd) = match zpair {
9606                    Some(p) => p,
9607                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
9608                };
9609                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
9610                    Some(p) => p,
9611                    None => (
9612                        e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
9613                        e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?,
9614                    ),
9615                }
9616            } else {
9617                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
9618                // launch for the verify's gate+up — the up segment's blocks fill SMs as
9619                // the gate segment drains (the launch-tail mechanism behind the b-tier
9620                // plateau; first positive after six falsified in-kernel variants).
9621                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9622                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
9623                let fused = if f2b {
9624                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
9625                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
9626                } else {
9627                    None
9628                };
9629                match fused {
9630                    Some(p) => p,
9631                    None => {
9632                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
9633                        e.mmq_act_begin();
9634                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
9635                    }
9636                }
9637            };
9638            let mut act = e.uninit(t * n_ff)?;
9639            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
9640            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
9641            let f0 = if e.uses_q8_1_fast(ffn_down) {
9642                let upv = e.view(&up, t * n_ff);
9643                let up_all = upv.slice(0..t * n_ff);
9644                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
9645                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
9646            } else {
9647                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9648                e.matmul(ffn_down, &act, t)?
9649            };
9650            if defer_post_norm {
9651                return Ok((f0, attn_out));
9652            }
9653            let mut sn = e.uninit(t * n_embd)?;
9654            e.rms_norm(
9655                &f0,
9656                bits.post_ffw_norm.float_data(),
9657                &mut sn,
9658                n_embd,
9659                t,
9660                eps,
9661            )?;
9662            return Ok((sn, attn_out));
9663        };
9664
9665        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
9666        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
9667        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
9668        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
9669        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
9670        let mut attn_out = e.uninit(t * n_embd)?;
9671        let mut router_in = e.uninit(t * n_embd)?;
9672        let fast_moe = match &layer.ffn {
9673            crate::hybrid::Ffn::Moe(m) => {
9674                m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
9675                    && expert_dp4a_supported(m.gate_exps.qtype)
9676                    && expert_dp4a_supported(m.up_exps.qtype)
9677                    && expert_dp4a_supported(m.down_exps.qtype)
9678                    && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0")
9679            }
9680            _ => false,
9681        };
9682        let q8z = t < PRIME_MIN_T && fast_moe;
9683        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
9684            let (z0, m2) = e.add_rms_norm3_q8z(
9685                cur,
9686                x,
9687                bits.ffn_norm.float_data(),
9688                &mbits.router_scale_pre,
9689                mbits.pre_ffw_norm_2.float_data(),
9690                &mut attn_out,
9691                &mut router_in,
9692                n_embd,
9693                t,
9694                eps,
9695            )?;
9696            (None, Some(z0), Some(m2))
9697        } else {
9698            let mut zsh = e.uninit(t * n_embd)?;
9699            let mut moe_in = e.uninit(t * n_embd)?;
9700            e.add_rms_norm3(
9701                cur,
9702                x,
9703                bits.ffn_norm.float_data(),
9704                &mbits.router_scale_pre,
9705                mbits.pre_ffw_norm_2.float_data(),
9706                &mut attn_out,
9707                &mut zsh,
9708                &mut router_in,
9709                &mut moe_in,
9710                n_embd,
9711                t,
9712                eps,
9713            )?;
9714            (Some((zsh, moe_in)), None, None)
9715        };
9716        let attn_out2 = attn_out;
9717        #[allow(unused_variables)]
9718        let attn_out = &attn_out2;
9719        let n_ff = mbits.shared_gate.out_features();
9720        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
9721            if t == 1 {
9722                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
9723                    Some(p) => p,
9724                    None => {
9725                        let h0 = e.zeros(0)?;
9726                        (
9727                            e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
9728                            e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?,
9729                        )
9730                    }
9731                }
9732            } else {
9733                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
9734                let h0 = e.zeros(0)?;
9735                (
9736                    e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
9737                    e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?,
9738                )
9739            }
9740        } else {
9741            let (zsh, _) = zsh_f32.as_ref().unwrap();
9742            (
9743                e.matmul(&mbits.shared_gate, zsh, t)?,
9744                e.matmul(&mbits.shared_up, zsh, t)?,
9745            )
9746        };
9747        let mut act = e.uninit(t * n_ff)?;
9748        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
9749        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
9750        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else {
9751            panic!("gemma4 layer not MoE")
9752        };
9753        let moe0 = match (&moe_q8, &zsh_f32) {
9754            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
9755            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
9756            _ => unreachable!(),
9757        };
9758        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
9759        let mut mlp = e.uninit(t * n_embd)?;
9760        let mut moe = e.uninit(t * n_embd)?;
9761        e.rms_norm2x(
9762            &mlp0,
9763            &moe0,
9764            mbits.post_ffw_norm_1.float_data(),
9765            mbits.post_ffw_norm_2.float_data(),
9766            &mut mlp,
9767            &mut moe,
9768            n_embd,
9769            t,
9770            eps,
9771        )?;
9772
9773        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
9774        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
9775        let mut sum = e.uninit(t * n_embd)?;
9776        let mut sn = e.uninit(t * n_embd)?;
9777        e.add_rms_norm(
9778            &mlp,
9779            &moe,
9780            bits.post_ffw_norm.float_data(),
9781            &mut sum,
9782            &mut sn,
9783            n_embd,
9784            t,
9785            eps,
9786        )?;
9787        Ok((sn, attn_out2))
9788    }
9789
9790    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
9791    fn gemma4_layer_tail_add_nq(
9792        &self,
9793        e: &Engine,
9794        layer: &crate::hybrid::HybridLayer,
9795        cur: &CudaSlice<f32>,
9796        x: &CudaSlice<f32>,
9797        t: usize,
9798        next_norm: Option<&CudaSlice<f32>>,
9799    ) -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>>
9800    {
9801        let n_embd = self.cfg.n_embd as usize;
9802        let bits = layer.gemma4.as_ref().unwrap();
9803        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
9804        let mut xn = e.uninit(t * n_embd)?;
9805        match next_norm {
9806            Some(w) => {
9807                let pair = e.add_scale_rms_norm_q8_1(
9808                    &sn,
9809                    &attn_out,
9810                    bits.layer_scale,
9811                    w,
9812                    &mut xn,
9813                    n_embd,
9814                    t,
9815                    self.cfg.rms_eps,
9816                )?;
9817                Ok((xn, Some(pair)))
9818            }
9819            None => {
9820                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
9821                Ok((xn, None))
9822            }
9823        }
9824    }
9825
9826    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
9827    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
9828    fn gemma4_forward(
9829        &self,
9830        e: &Engine,
9831        tokens: &[u32],
9832        last_only: bool,
9833    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
9834        // E4B routes to its own forward regardless of the caller's entry point (forward /
9835        // forward_last / prime paths all funnel here for gemma4).
9836        if self.is_gemma4_e4b() {
9837            return self.gemma4_e4b_forward(e, tokens, last_only);
9838        }
9839        let n_embd = self.cfg.n_embd as usize;
9840        let t = tokens.len();
9841        let pos: Vec<i32> = (0..t as i32).collect();
9842        let pos_d = e.htod_i32(&pos)?;
9843
9844        let mut x = self.embed(e, tokens)?;
9845        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9846        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
9847        // the bring-up bisect vs llama-eval-callback node stats.
9848        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
9849        let stat =
9850            |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
9851                let h = e.dtoh(x)?;
9852                let bad = h.iter().filter(|v| !v.is_finite()).count();
9853                let mx = h
9854                    .iter()
9855                    .filter(|v| v.is_finite())
9856                    .fold(0.0f32, |m, v| m.max(v.abs()));
9857                eprintln!(
9858                    "[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}",
9859                    &h[..3]
9860                );
9861                Ok(())
9862            };
9863        if probe {
9864            stat(e, &x, "embed")?;
9865        }
9866        for (il, layer) in self.layers.iter().enumerate() {
9867            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
9868            if probe {
9869                stat(e, &x, &format!("L{il}"))?;
9870            }
9871        }
9872        let mut hn = e.zeros(t * n_embd)?;
9873        e.rms_norm(
9874            &x,
9875            self.output_norm.float_data(),
9876            &mut hn,
9877            n_embd,
9878            t,
9879            self.cfg.rms_eps,
9880        )?;
9881        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9882        let n_vocab = self.output.out_features();
9883        let logits = if last_only {
9884            let hv = e.view(&hn, t * n_embd);
9885            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
9886            let mut hlast = e.zeros(n_embd)?;
9887            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
9888            let mut ld = e.matmul(&self.output, &hlast, 1)?;
9889            e.softcap(&mut ld, cap, n_vocab)?;
9890            self.gemma4_suppress(e, &mut ld, 1)?;
9891            e.dtoh(&ld)?
9892        } else {
9893            let mut ld = e.matmul(&self.output, &hn, t)?;
9894            e.softcap(&mut ld, cap, t * n_vocab)?;
9895            self.gemma4_suppress(e, &mut ld, t)?;
9896            e.dtoh(&ld)?
9897        };
9898        Ok(logits)
9899    }
9900
9901    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
9902    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
9903    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
9904    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
9905    pub(crate) fn gemma4_prime(
9906        &self,
9907        e: &Engine,
9908        tokens: &[u32],
9909        cache: &mut Cache,
9910    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9911        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
9912        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
9913        // whole worker process on this line. The worker now primes gemma4 monolithically and
9914        // routes continuation suffixes tokenwise; this is the per-request backstop.
9915        if cache.pos != 0 {
9916            return Err(
9917                "gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
9918                        — prime the full prompt in one call or decode tokenwise"
9919                    .into(),
9920            );
9921        }
9922        let n_embd = self.cfg.n_embd as usize;
9923        let eps = self.cfg.rms_eps;
9924        let t = tokens.len();
9925        let pos: Vec<i32> = (0..t as i32).collect();
9926        let pos_d = e.htod_i32(&pos)?;
9927        let mut x = self.embed(e, tokens)?;
9928        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9929        for (il, layer) in self.layers.iter().enumerate() {
9930            let mut h = e.zeros(t * n_embd)?;
9931            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
9932            let Mixer::Full(fa) = &layer.mixer else {
9933                panic!("gemma4 layer not full-attn")
9934            };
9935            let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
9936            let mut cur = e.zeros(t * n_embd)?;
9937            e.rms_norm(
9938                &o,
9939                layer.post_attn_norm.float_data(),
9940                &mut cur,
9941                n_embd,
9942                t,
9943                eps,
9944            )?;
9945            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
9946            self.dflash_tap(e, cache, il, &x, t)?;
9947        }
9948        cache.pos += t;
9949        let hiddens = e.clone_dtod(&x)?;
9950        let xv = e.view(&x, t * n_embd);
9951        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
9952        let mut h_seed = e.zeros(n_embd)?;
9953        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
9954        let mut hn = e.uninit(n_embd)?;
9955        e.rms_norm(
9956            &h_seed,
9957            self.output_norm.float_data(),
9958            &mut hn,
9959            n_embd,
9960            1,
9961            eps,
9962        )?;
9963        let mut ld = e.matmul(&self.output, &hn, 1)?;
9964        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9965        e.softcap(&mut ld, cap, self.output.out_features())?;
9966        self.gemma4_suppress(e, &mut ld, 1)?;
9967        let logits = e.dtoh(&ld)?;
9968        Ok((logits, h_seed, hiddens))
9969    }
9970
9971    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
9972    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
9973    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
9974    /// fused norm emits q8 directly — the f32 h never materializes).
9975    fn gemma4_decode_attn(
9976        &self,
9977        e: &Engine,
9978        fa: &crate::hybrid::FullAttnLayer,
9979        il: usize,
9980        hq: &CudaSlice<i8>,
9981        hdq: &CudaSlice<f32>,
9982        pos_d: &CudaSlice<i32>,
9983        cache: &mut Cache,
9984    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9985        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
9986        let eps = self.cfg.rms_eps;
9987        let aux = self.gemma4_aux.as_ref().unwrap();
9988        let ones = aux.ones(e);
9989        #[cfg(debug_assertions)]
9990        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn.ones");
9991        let (hq, hdq) = (hq, hdq);
9992        let h0 = e.zeros(0)?;
9993        let h = &h0;
9994        let (q0, k0, v0) = if swa {
9995            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
9996                Some(t3) => t3,
9997                None => (
9998                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
9999                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10000                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
10001                ),
10002            }
10003        } else {
10004            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
10005                Some(p) => p,
10006                None => (
10007                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
10008                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
10009                ),
10010            };
10011            let v0 = e.clone_dtod(&k0)?;
10012            (q0, k0, v0)
10013        };
10014        let mut q = e.uninit(nh * hd)?;
10015        let mut k = e.uninit(nkv * hd)?;
10016        let mut v = e.uninit(nkv * hd)?;
10017        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
10018        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
10019        let ff = if swa {
10020            None
10021        } else {
10022            Some(
10023                aux.rope_freqs(e)
10024                    .expect("gemma4 global rope needs rope_freqs.weight"),
10025            )
10026        };
10027        #[cfg(debug_assertions)]
10028        if let Some(ff) = ff {
10029            crate::debug_assert_tensor_stream_device(
10030                ff,
10031                &e.stream(),
10032                "gemma4_decode_attn.rope_freqs",
10033            );
10034        }
10035        e.rms_norm_qkv_rope(
10036            &q0,
10037            &k0,
10038            &v0,
10039            fa.q_norm.float_data(),
10040            fa.k_norm.float_data(),
10041            ones,
10042            &mut q,
10043            &mut k,
10044            &mut v,
10045            hd,
10046            nh,
10047            nkv,
10048            pos_d,
10049            nh,
10050            nkv,
10051            base,
10052            1.0,
10053            ff,
10054            eps,
10055        )?;
10056        let kvl = cache.kv[il].as_mut().unwrap();
10057        e.append_kv_quantized(
10058            &k,
10059            &v,
10060            &mut kvl.k,
10061            &mut kvl.v,
10062            kvl.len,
10063            kvl.kv_dim_k,
10064            kvl.kv_dim_v,
10065            kvl.k_tok_bytes,
10066            kvl.v_tok_bytes,
10067            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
10068        )?;
10069        kvl.len += 1;
10070        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
10071        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
10072        // positional). Globals attend the full history.
10073        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10074        let mut attn = e.uninit(nh * hd)?;
10075        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
10076        if !swa
10077            && hd == 512
10078            && kvl.len >= crate::fa512_min_tkv()
10079            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10080        {
10081            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10082            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10083            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
10084            let base = kvl.len as i32;
10085            e.i32_set_k(&mut kvl.len_d, base)?;
10086            e.fa_decode_rows(
10087                &q,
10088                &kp,
10089                &vp,
10090                &mut attn,
10091                hd,
10092                nh,
10093                nkv,
10094                kvl.len - 1,
10095                1,
10096                scale,
10097                kvl.k_tok_bytes,
10098                kvl.v_tok_bytes,
10099                Some((&kvl.len_d, -1)),
10100                false,
10101                false,
10102                None,
10103            )?;
10104            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10105        }
10106        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
10107        if swa
10108            && kvl.len > win
10109            && hd == 256
10110            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10111        {
10112            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10113            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10114            let base = kvl.len as i32;
10115            e.i32_set_k(&mut kvl.len_d, base)?;
10116            e.fa_decode_rows_w(
10117                &q,
10118                &kp,
10119                &vp,
10120                &mut attn,
10121                hd,
10122                nh,
10123                nkv,
10124                &kvl.len_d,
10125                -1,
10126                1,
10127                scale,
10128                win,
10129                kvl.k_tok_bytes,
10130                kvl.v_tok_bytes,
10131                None,
10132            )?;
10133            return Ok(e.matmul(&fa.wo, &attn, 1)?);
10134        }
10135        let (off_tok, t_kv) = if swa && kvl.len > win {
10136            (kvl.len - win, win)
10137        } else {
10138            (0, kvl.len)
10139        };
10140        let k_view = e.view_u8_range(
10141            &kvl.k,
10142            off_tok * kvl.k_tok_bytes,
10143            (off_tok + t_kv) * kvl.k_tok_bytes,
10144        );
10145        let v_view = e.view_u8_range(
10146            &kvl.v,
10147            off_tok * kvl.v_tok_bytes,
10148            (off_tok + t_kv) * kvl.v_tok_bytes,
10149        );
10150        e.fa_decode_kvmod(
10151            &q,
10152            &k_view,
10153            &v_view,
10154            &mut attn,
10155            hd,
10156            nh,
10157            nkv,
10158            t_kv,
10159            scale,
10160            kvl.k_tok_bytes,
10161            kvl.v_tok_bytes,
10162            swa && crate::Engine::wkv_on(),
10163        )?;
10164        Ok(e.matmul(&fa.wo, &attn, 1)?)
10165    }
10166
10167    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
10168    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
10169    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
10170    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
10171    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
10172    /// in-graph; the driver gates).
10173    #[allow(clippy::too_many_arguments)]
10174    pub fn gemma4_decode_step_dc(
10175        &self,
10176        e: &Engine,
10177        token_d: &CudaSlice<u32>,
10178        pos_d: &mut CudaSlice<i32>,
10179        embd_gpu: &CudaSlice<u8>,
10180        embd_qt: i32,
10181        embd_rb: usize,
10182        cache: &mut Cache,
10183        n_vocab: usize,
10184        cap_bucket_max: Option<(usize, usize)>,
10185    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
10186        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
10187        self.gemma4_decode_step_dc_into(
10188            e,
10189            token_d,
10190            pos_d,
10191            embd_gpu,
10192            embd_qt,
10193            embd_rb,
10194            cache,
10195            n_vocab,
10196            cap_bucket_max,
10197            &mut tok_out,
10198        )?;
10199        Ok(tok_out)
10200    }
10201
10202    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
10203    /// every replay; pass `token_d` itself for the self-feeding graph loop).
10204    #[allow(clippy::too_many_arguments)]
10205    pub fn gemma4_decode_step_dc_into(
10206        &self,
10207        e: &Engine,
10208        token_d: &CudaSlice<u32>,
10209        pos_d: &mut CudaSlice<i32>,
10210        embd_gpu: &CudaSlice<u8>,
10211        embd_qt: i32,
10212        embd_rb: usize,
10213        cache: &mut Cache,
10214        n_vocab: usize,
10215        cap_bucket_max: Option<(usize, usize)>,
10216        tok_out: &mut CudaSlice<u32>,
10217    ) -> Result<(), Box<dyn std::error::Error>> {
10218        let n_embd = self.cfg.n_embd as usize;
10219        let eps = self.cfg.rms_eps;
10220        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
10221        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
10222        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10223        let n_layers = self.layers.len();
10224        for (il, layer) in self.layers.iter().enumerate() {
10225            let (hq, hdq) = match h_carry.take() {
10226                Some(p) => p,
10227                None => {
10228                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
10229                }
10230            };
10231            let Mixer::Full(fa) = &layer.mixer else {
10232                panic!("gemma4 layer {il} not full-attn")
10233            };
10234            let o =
10235                self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
10236            let mut cur = e.uninit(n_embd)?;
10237            e.rms_norm(
10238                &o,
10239                layer.post_attn_norm.float_data(),
10240                &mut cur,
10241                n_embd,
10242                1,
10243                eps,
10244            )?;
10245            let next_norm = if il + 1 < n_layers {
10246                Some(self.layers[il + 1].attn_norm.float_data())
10247            } else {
10248                None
10249            };
10250            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
10251            x = xn;
10252            h_carry = hn;
10253        }
10254        let mut hn = e.uninit(n_embd)?;
10255        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
10256        let mut logits = e.matmul(&self.output, &hn, 1)?;
10257        self.gemma4_suppress(e, &mut logits, 1)?; // cap skipped (monotonic); the mask is not
10258        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
10259        e.inc_seqlen(pos_d)?;
10260        if cap_bucket_max.is_none() {
10261            cache.pos += 1;
10262        }
10263        Ok(())
10264    }
10265
10266    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
10267    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
10268    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
10269    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
10270
10271    /// Build the slot set (call OUTSIDE any capture).
10272    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
10273        let n_embd = self.cfg.n_embd as usize;
10274        let n_vocab = self.output.out_features();
10275        let n_layers = self.layers.len();
10276        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
10277        for il in 0..n_layers {
10278            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
10279            qmax = qmax.max(nh * hd);
10280            kvmax = kvmax.max(nkv * hd);
10281            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
10282                ffmax = ffmax.max(ffn_gate.out_features());
10283            }
10284        }
10285        Ok(G4DcSlots {
10286            x: e.uninit(n_embd)?,
10287            xn: e.uninit(n_embd)?,
10288            cur: e.uninit(n_embd)?,
10289            hq: e.alloc_i8_uninit(n_embd)?,
10290            hd_: e.uninit(n_embd / 32)?,
10291            q0: e.uninit(qmax)?,
10292            k0: e.uninit(kvmax)?,
10293            v0: e.uninit(kvmax)?,
10294            q: e.uninit(qmax)?,
10295            k: e.uninit(kvmax)?,
10296            v: e.uninit(kvmax)?,
10297            attn: e.uninit(qmax)?,
10298            o: e.uninit(n_embd)?,
10299            attn_out: e.uninit(n_embd)?,
10300            zsh: e.uninit(n_embd)?,
10301            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
10302            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
10303            zq: e.alloc_i8_uninit(n_embd.max(qmax))?,
10304            zd: e.uninit(n_embd.max(qmax) / 32)?,
10305            gate: e.uninit(ffmax)?,
10306            up: e.uninit(ffmax)?,
10307            act: e.uninit(ffmax)?,
10308            actq: e.alloc_i8_uninit(ffmax)?,
10309            actd: e.uninit(ffmax / 32)?,
10310            f0: e.uninit(n_embd)?,
10311            sn: e.uninit(n_embd)?,
10312            hn: e.uninit(n_embd)?,
10313            logits: e.uninit(n_vocab)?,
10314        })
10315    }
10316
10317    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
10318    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
10319    fn g4_matvec_m1_into(
10320        &self,
10321        e: &Engine,
10322        w: &crate::model::GpuTensor,
10323        aq: &CudaSlice<i8>,
10324        ad: &CudaSlice<f32>,
10325        y: &mut CudaSlice<f32>,
10326    ) -> Result<(), Box<dyn std::error::Error>> {
10327        use crate::model::GpuTensor;
10328        let (bytes, qtype, row_bytes, scale, rp) = match w {
10329            GpuTensor::Quant {
10330                bytes,
10331                qtype,
10332                row_bytes,
10333                scale,
10334                rp,
10335                ..
10336            } => (bytes, *qtype, *row_bytes, *scale, *rp),
10337            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
10338        };
10339        let (mbytes, mrp) = match w {
10340            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
10341            _ => (bytes, rp),
10342        };
10343        e.qmatvec_mmvq_into(
10344            mbytes,
10345            aq,
10346            ad,
10347            1,
10348            w.in_features(),
10349            w.out_features(),
10350            qtype,
10351            row_bytes,
10352            scale,
10353            mrp,
10354            y,
10355        )
10356    }
10357
10358    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
10359    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
10360    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
10361    #[allow(clippy::too_many_arguments)]
10362    pub fn gemma4_decode_step_dc_slotted(
10363        &self,
10364        e: &Engine,
10365        token_d: &CudaSlice<u32>,
10366        pos_d: &mut CudaSlice<i32>,
10367        embd_gpu: &CudaSlice<u8>,
10368        embd_qt: i32,
10369        embd_rb: usize,
10370        cache: &mut Cache,
10371        n_vocab: usize,
10372        cap_bucket_max: Option<(usize, usize)>,
10373        sl: &mut G4DcSlots,
10374        tok_out: &mut CudaSlice<u32>,
10375        ring: Option<(&mut CudaSlice<u32>, usize)>,
10376    ) -> Result<(), Box<dyn std::error::Error>> {
10377        let n_embd = self.cfg.n_embd as usize;
10378        let eps = self.cfg.rms_eps;
10379        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
10380        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
10381        let n_layers = self.layers.len();
10382        let mut has_carry = false;
10383        for il in 0..n_layers {
10384            if !has_carry {
10385                e.rms_norm_q8_1_into(
10386                    &sl.x,
10387                    self.layers[il].attn_norm.float_data(),
10388                    n_embd,
10389                    1,
10390                    eps,
10391                    &mut sl.hq,
10392                    &mut sl.hd_,
10393                )?;
10394            }
10395            has_carry = true;
10396            let layer = &self.layers[il];
10397            let Mixer::Full(fa) = &layer.mixer else {
10398                panic!("gemma4 layer {il} not full-attn")
10399            };
10400            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
10401            e.rms_norm(
10402                &sl.o,
10403                layer.post_attn_norm.float_data(),
10404                &mut sl.cur,
10405                n_embd,
10406                1,
10407                eps,
10408            )?;
10409            let next_norm = if il + 1 < n_layers {
10410                Some(self.layers[il + 1].attn_norm.float_data())
10411            } else {
10412                None
10413            };
10414            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
10415            std::mem::swap(&mut sl.x, &mut sl.xn);
10416        }
10417        e.rms_norm(
10418            &sl.x,
10419            self.output_norm.float_data(),
10420            &mut sl.hn,
10421            n_embd,
10422            1,
10423            eps,
10424        )?;
10425        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10426        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
10427        {
10428            let (zq, zd) = (&sl.zq, &sl.zd);
10429            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
10430            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
10431            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
10432        }
10433        self.gemma4_suppress(e, &mut sl.logits, 1)?;
10434        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
10435        if let Some((ring, base)) = ring {
10436            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
10437            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
10438            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
10439            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
10440        }
10441        e.inc_seqlen(pos_d)?;
10442        if cap_bucket_max.is_none() {
10443            cache.pos += 1;
10444        }
10445        Ok(())
10446    }
10447
10448    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
10449    #[allow(clippy::too_many_arguments)]
10450    fn gemma4_decode_attn_dc_slotted(
10451        &self,
10452        e: &Engine,
10453        fa: &crate::hybrid::FullAttnLayer,
10454        il: usize,
10455        pos_d: &CudaSlice<i32>,
10456        cache: &mut Cache,
10457        cap_bucket_max: Option<(usize, usize)>,
10458        sl: &mut G4DcSlots,
10459    ) -> Result<(), Box<dyn std::error::Error>> {
10460        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10461        let eps = self.cfg.rms_eps;
10462        let aux = self.gemma4_aux.as_ref().unwrap();
10463        let ones = aux.ones(e);
10464        #[cfg(debug_assertions)]
10465        crate::debug_assert_tensor_stream_device(
10466            ones,
10467            &e.stream(),
10468            "gemma4_decode_attn_dc_slotted.ones",
10469        );
10470        {
10471            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
10472            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
10473            if swa {
10474                if !e.matmul_q4_fused3_into(
10475                    &fa.wq, &fa.wk, &fa.wv, hq, hdq, &mut sl.q0, &mut sl.k0, &mut sl.v0,
10476                )? {
10477                    return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
10478                }
10479            } else {
10480                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
10481                    return Err("slotted step: fused2 unavailable".into());
10482                }
10483                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
10484                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
10485            }
10486        }
10487        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
10488        // kernel-for-kernel (graph stream-identity gate).
10489        let ff = if swa {
10490            None
10491        } else {
10492            Some(
10493                aux.rope_freqs(e)
10494                    .expect("gemma4 global rope needs rope_freqs.weight"),
10495            )
10496        };
10497        #[cfg(debug_assertions)]
10498        if let Some(ff) = ff {
10499            crate::debug_assert_tensor_stream_device(
10500                ff,
10501                &e.stream(),
10502                "gemma4_decode_attn_dc_slotted.rope_freqs",
10503            );
10504        }
10505        let kvl = cache.kv[il].as_mut().unwrap();
10506        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10507        if crate::Engine::qkv_append_on() {
10508            // append fold (2026-07-23): mirrors dc_into.
10509            e.rms_norm_qkv_rope_append_dc(
10510                &sl.q0,
10511                &sl.k0,
10512                &sl.v0,
10513                fa.q_norm.float_data(),
10514                fa.k_norm.float_data(),
10515                ones,
10516                &mut sl.q,
10517                &mut sl.k,
10518                &mut sl.v,
10519                hd,
10520                nh,
10521                nkv,
10522                pos_d,
10523                nh,
10524                nkv,
10525                base,
10526                1.0,
10527                ff,
10528                eps,
10529                &mut kvl.k,
10530                &mut kvl.v,
10531                &kvl.len_d,
10532                kvl.k_tok_bytes,
10533                kvl.v_tok_bytes,
10534                kv_fp8,
10535            )?;
10536        } else {
10537            e.rms_norm_qkv_rope(
10538                &sl.q0,
10539                &sl.k0,
10540                &sl.v0,
10541                fa.q_norm.float_data(),
10542                fa.k_norm.float_data(),
10543                ones,
10544                &mut sl.q,
10545                &mut sl.k,
10546                &mut sl.v,
10547                hd,
10548                nh,
10549                nkv,
10550                pos_d,
10551                nh,
10552                nkv,
10553                base,
10554                1.0,
10555                ff,
10556                eps,
10557            )?;
10558            e.append_kv_quantized_dc(
10559                &sl.k,
10560                &sl.v,
10561                &mut kvl.k,
10562                &mut kvl.v,
10563                &kvl.len_d,
10564                kvl.kv_dim_k,
10565                kvl.kv_dim_v,
10566                kvl.k_tok_bytes,
10567                kvl.v_tok_bytes,
10568                kv_fp8,
10569            )?;
10570        }
10571        e.inc_seqlen(&mut kvl.len_d)?;
10572        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
10573        let k_view = e.view_u8(&kvl.k, kvl.k.len());
10574        let v_view = e.view_u8(&kvl.v, kvl.v.len());
10575        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10576        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10577        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
10578        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
10579        // the dc_into arm branch-for-branch (stream gate).
10580        let mut fa_q8 = false;
10581        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10582            e.fa_decode_rows(
10583                &sl.q,
10584                &k_view,
10585                &v_view,
10586                &mut sl.attn,
10587                hd,
10588                nh,
10589                nkv,
10590                b_glob - 1,
10591                1,
10592                scale,
10593                kvl.k_tok_bytes,
10594                kvl.v_tok_bytes,
10595                Some((&kvl.len_d, -1)),
10596                false,
10597                false,
10598                Some((&mut sl.zq, &mut sl.zd)),
10599            )?;
10600            fa_q8 = true;
10601        } else if swa && b_swa > win && hd == 256 && rows_on {
10602            e.fa_decode_rows_w(
10603                &sl.q,
10604                &k_view,
10605                &v_view,
10606                &mut sl.attn,
10607                hd,
10608                nh,
10609                nkv,
10610                &kvl.len_d,
10611                -1,
10612                1,
10613                scale,
10614                win,
10615                kvl.k_tok_bytes,
10616                kvl.v_tok_bytes,
10617                Some((&mut sl.zq, &mut sl.zd)),
10618            )?;
10619            fa_q8 = true;
10620        } else {
10621            let b = if swa { b_swa } else { b_glob };
10622            e.fa_decode_dc(
10623                &sl.q,
10624                &k_view,
10625                &v_view,
10626                &mut sl.attn,
10627                hd,
10628                nh,
10629                nkv,
10630                &kvl.len_d,
10631                b,
10632                scale,
10633                kvl.k_tok_bytes,
10634                kvl.v_tok_bytes,
10635                swa && crate::Engine::wkv_on(),
10636            )?;
10637        }
10638        if !fa_q8 {
10639            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
10640            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
10641        }
10642        {
10643            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10644            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10645            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
10646        }
10647        Ok(())
10648    }
10649
10650    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
10651    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
10652    fn gemma4_layer_tail_slotted(
10653        &self,
10654        e: &Engine,
10655        layer: &crate::hybrid::HybridLayer,
10656        next_norm: Option<&CudaSlice<f32>>,
10657        sl: &mut G4DcSlots,
10658    ) -> Result<(), Box<dyn std::error::Error>> {
10659        let n_embd = self.cfg.n_embd as usize;
10660        let eps = self.cfg.rms_eps;
10661        let bits = layer.gemma4.as_ref().unwrap();
10662        let crate::hybrid::Ffn::Dense {
10663            ffn_gate,
10664            ffn_up,
10665            ffn_down,
10666        } = &layer.ffn
10667        else {
10668            return Err("slotted tail: dense ffn only".into());
10669        };
10670        e.add_rms_norm(
10671            &sl.cur,
10672            &sl.x,
10673            bits.ffn_norm.float_data(),
10674            &mut sl.attn_out,
10675            &mut sl.zsh,
10676            n_embd,
10677            1,
10678            eps,
10679        )?;
10680        let n_ff = ffn_gate.out_features();
10681        {
10682            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
10683            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
10684        }
10685        {
10686            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
10687            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
10688            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
10689                return Err("slotted tail: ffn fused2 unavailable".into());
10690            }
10691        }
10692        debug_assert!(e.uses_q8_1_fast(ffn_down));
10693        {
10694            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
10695            let upv = e.view(upr, n_ff);
10696            let up_all = upv.slice(0..n_ff);
10697            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
10698            e.gelu_tanh_mul_q8_1_into(
10699                gr,
10700                &up_all,
10701                &mut sl.act,
10702                n_ff,
10703                1,
10704                &mut sl.actq,
10705                &mut sl.actd,
10706            )?;
10707        }
10708        {
10709            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
10710            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
10711            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
10712        }
10713        e.rms_norm(
10714            &sl.f0,
10715            bits.post_ffw_norm.float_data(),
10716            &mut sl.sn,
10717            n_embd,
10718            1,
10719            eps,
10720        )?;
10721        match next_norm {
10722            Some(w) => {
10723                e.add_scale_rms_norm_q8_1_into(
10724                    &sl.sn,
10725                    &sl.attn_out,
10726                    bits.layer_scale,
10727                    w,
10728                    &mut sl.xn,
10729                    n_embd,
10730                    1,
10731                    eps,
10732                    &mut sl.hq,
10733                    &mut sl.hd_,
10734                )?;
10735            }
10736            None => {
10737                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
10738            }
10739        }
10740        Ok(())
10741    }
10742
10743    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
10744    #[allow(clippy::too_many_arguments)]
10745    fn gemma4_decode_attn_dc(
10746        &self,
10747        e: &Engine,
10748        fa: &crate::hybrid::FullAttnLayer,
10749        il: usize,
10750        hq: &CudaSlice<i8>,
10751        hdq: &CudaSlice<f32>,
10752        pos_d: &CudaSlice<i32>,
10753        cache: &mut Cache,
10754        cap_bucket_max: Option<(usize, usize)>,
10755    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
10756        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
10757        let eps = self.cfg.rms_eps;
10758        let aux = self.gemma4_aux.as_ref().unwrap();
10759        let ones = aux.ones(e);
10760        #[cfg(debug_assertions)]
10761        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_decode_attn_dc.ones");
10762        let (q0, k0, v0) = if swa {
10763            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
10764                Some(t3) => t3,
10765                None => {
10766                    let h0 = e.zeros(0)?;
10767                    (
10768                        e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
10769                        e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
10770                        e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?,
10771                    )
10772                }
10773            }
10774        } else {
10775            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
10776                Some(p) => p,
10777                None => {
10778                    let h0 = e.zeros(0)?;
10779                    (
10780                        e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
10781                        e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
10782                    )
10783                }
10784            };
10785            let v0 = e.clone_dtod(&k0)?;
10786            (q0, k0, v0)
10787        };
10788        let mut q = e.uninit(nh * hd)?;
10789        let mut k = e.uninit(nkv * hd)?;
10790        let mut v = e.uninit(nkv * hd)?;
10791        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
10792        let ff = if swa {
10793            None
10794        } else {
10795            Some(
10796                aux.rope_freqs(e)
10797                    .expect("gemma4 global rope needs rope_freqs.weight"),
10798            )
10799        };
10800        #[cfg(debug_assertions)]
10801        if let Some(ff) = ff {
10802            crate::debug_assert_tensor_stream_device(
10803                ff,
10804                &e.stream(),
10805                "gemma4_decode_attn_dc.rope_freqs",
10806            );
10807        }
10808        let kvl = cache.kv[il].as_mut().unwrap();
10809        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
10810        if crate::Engine::qkv_append_on() {
10811            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
10812            e.rms_norm_qkv_rope_append_dc(
10813                &q0,
10814                &k0,
10815                &v0,
10816                fa.q_norm.float_data(),
10817                fa.k_norm.float_data(),
10818                ones,
10819                &mut q,
10820                &mut k,
10821                &mut v,
10822                hd,
10823                nh,
10824                nkv,
10825                pos_d,
10826                nh,
10827                nkv,
10828                base,
10829                1.0,
10830                ff,
10831                eps,
10832                &mut kvl.k,
10833                &mut kvl.v,
10834                &kvl.len_d,
10835                kvl.k_tok_bytes,
10836                kvl.v_tok_bytes,
10837                kv_fp8,
10838            )?;
10839        } else {
10840            e.rms_norm_qkv_rope(
10841                &q0,
10842                &k0,
10843                &v0,
10844                fa.q_norm.float_data(),
10845                fa.k_norm.float_data(),
10846                ones,
10847                &mut q,
10848                &mut k,
10849                &mut v,
10850                hd,
10851                nh,
10852                nkv,
10853                pos_d,
10854                nh,
10855                nkv,
10856                base,
10857                1.0,
10858                ff,
10859                eps,
10860            )?;
10861            e.append_kv_quantized_dc(
10862                &k,
10863                &v,
10864                &mut kvl.k,
10865                &mut kvl.v,
10866                &kvl.len_d,
10867                kvl.kv_dim_k,
10868                kvl.kv_dim_v,
10869                kvl.k_tok_bytes,
10870                kvl.v_tok_bytes,
10871                kv_fp8,
10872            )?;
10873        }
10874        e.inc_seqlen(&mut kvl.len_d)?;
10875        let mut attn = e.uninit(nh * hd)?;
10876        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
10877        // rides g4_matvec_m1_into instead of matmul's internal quantize.
10878        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
10879        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
10880        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
10881        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
10882        // (gemma4_e4b_attn, +0.65% valid window).
10883        match cap_bucket_max {
10884            None => {
10885                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
10886                // decode (SWA layers attend the last `sliding_window` keys); the device
10887                // counters carry only the append slot + the graph seam.
10888                kvl.len += 1;
10889                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10890                if !swa
10891                    && hd == 512
10892                    && kvl.len >= crate::fa512_min_tkv()
10893                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10894                {
10895                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
10896                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
10897                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10898                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10899                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10900                    e.fa_decode_rows(
10901                        &q,
10902                        &kp,
10903                        &vp,
10904                        &mut attn,
10905                        hd,
10906                        nh,
10907                        nkv,
10908                        kvl.len - 1,
10909                        1,
10910                        scale,
10911                        kvl.k_tok_bytes,
10912                        kvl.v_tok_bytes,
10913                        Some((&kvl.len_d, -1)),
10914                        false,
10915                        false,
10916                        Some((&mut aq8, &mut ad8)),
10917                    )?;
10918                    fa_q8 = Some((aq8, ad8));
10919                } else if swa
10920                    && kvl.len > win
10921                    && hd == 256
10922                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
10923                {
10924                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
10925                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
10926                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
10927                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10928                    e.fa_decode_rows_w(
10929                        &q,
10930                        &kp,
10931                        &vp,
10932                        &mut attn,
10933                        hd,
10934                        nh,
10935                        nkv,
10936                        &kvl.len_d,
10937                        -1,
10938                        1,
10939                        scale,
10940                        win,
10941                        kvl.k_tok_bytes,
10942                        kvl.v_tok_bytes,
10943                        Some((&mut aq8, &mut ad8)),
10944                    )?;
10945                    fa_q8 = Some((aq8, ad8));
10946                } else {
10947                    let (off_tok, t_kv) = if swa && kvl.len > win {
10948                        (kvl.len - win, win)
10949                    } else {
10950                        (0, kvl.len)
10951                    };
10952                    let k_view = e.view_u8_range(
10953                        &kvl.k,
10954                        off_tok * kvl.k_tok_bytes,
10955                        (off_tok + t_kv) * kvl.k_tok_bytes,
10956                    );
10957                    let v_view = e.view_u8_range(
10958                        &kvl.v,
10959                        off_tok * kvl.v_tok_bytes,
10960                        (off_tok + t_kv) * kvl.v_tok_bytes,
10961                    );
10962                    e.fa_decode_kvmod(
10963                        &q,
10964                        &k_view,
10965                        &v_view,
10966                        &mut attn,
10967                        hd,
10968                        nh,
10969                        nkv,
10970                        t_kv,
10971                        scale,
10972                        kvl.k_tok_bytes,
10973                        kvl.v_tok_bytes,
10974                        swa && crate::Engine::wkv_on(),
10975                    )?;
10976                }
10977            }
10978            Some((b_swa, b_glob)) => {
10979                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
10980                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
10981                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
10982                // the RUNG max for the rows family (kernels derive per-replay splits from
10983                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
10984                let k_view = e.view_u8(&kvl.k, kvl.k.len());
10985                let v_view = e.view_u8(&kvl.v, kvl.v.len());
10986                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
10987                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
10988                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
10989                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
10990                    e.fa_decode_rows(
10991                        &q,
10992                        &k_view,
10993                        &v_view,
10994                        &mut attn,
10995                        hd,
10996                        nh,
10997                        nkv,
10998                        b_glob - 1,
10999                        1,
11000                        scale,
11001                        kvl.k_tok_bytes,
11002                        kvl.v_tok_bytes,
11003                        Some((&kvl.len_d, -1)),
11004                        false,
11005                        false,
11006                        Some((&mut aq8, &mut ad8)),
11007                    )?;
11008                    fa_q8 = Some((aq8, ad8));
11009                } else if swa && b_swa > win && hd == 256 && rows_on {
11010                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
11011                    e.fa_decode_rows_w(
11012                        &q,
11013                        &k_view,
11014                        &v_view,
11015                        &mut attn,
11016                        hd,
11017                        nh,
11018                        nkv,
11019                        &kvl.len_d,
11020                        -1,
11021                        1,
11022                        scale,
11023                        win,
11024                        kvl.k_tok_bytes,
11025                        kvl.v_tok_bytes,
11026                        Some((&mut aq8, &mut ad8)),
11027                    )?;
11028                    fa_q8 = Some((aq8, ad8));
11029                } else {
11030                    let b = if swa { b_swa } else { b_glob };
11031                    e.fa_decode_dc(
11032                        &q,
11033                        &k_view,
11034                        &v_view,
11035                        &mut attn,
11036                        hd,
11037                        nh,
11038                        nkv,
11039                        &kvl.len_d,
11040                        b,
11041                        scale,
11042                        kvl.k_tok_bytes,
11043                        kvl.v_tok_bytes,
11044                        swa && crate::Engine::wkv_on(),
11045                    )?;
11046                }
11047            }
11048        }
11049        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
11050        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
11051        if let Some((aq8, ad8)) = fa_q8 {
11052            let mut y = e.uninit(fa.wo.out_features())?;
11053            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
11054            return Ok(y);
11055        }
11056        Ok(e.matmul(&fa.wo, &attn, 1)?)
11057    }
11058
11059    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
11060    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
11061    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
11062    /// views in-graph); caller gates and falls back to the dc-eager loop.
11063    pub fn gemma4_generate_graph(
11064        &self,
11065        e: &Engine,
11066        prompt_pos: usize,
11067        first_token: u32,
11068        cache: &mut Cache,
11069        max_new: usize,
11070        eos: &[u32],
11071        mut on_token: impl FnMut(u32) -> bool,
11072    ) -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
11073        if self.is_gemma4_e4b() {
11074            return Err(
11075                "E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm"
11076                    .into(),
11077            );
11078        }
11079        use crate::decode::StopReason;
11080        let n_vocab = self.output.out_features();
11081        let n_embd = self.cfg.n_embd as usize;
11082        let embd_gpu = self
11083            .embd_gpu
11084            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11085        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11086        for kvl in cache.kv.iter_mut().flatten() {
11087            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
11088        }
11089        let mut token_d = e.stream().clone_htod(&[first_token])?;
11090        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
11091        let g4 = self.cfg.gemma4.as_ref().unwrap();
11092        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
11093        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
11094        let nkv_s = g4
11095            .head_count_kv
11096            .iter()
11097            .zip(g4.swa_pattern.iter())
11098            .find(|p| *p.1)
11099            .map(|p| *p.0 as usize)
11100            .unwrap_or(8);
11101        let nkv_g = g4
11102            .head_count_kv
11103            .iter()
11104            .zip(g4.swa_pattern.iter())
11105            .find(|p| !*p.1)
11106            .map(|p| *p.0 as usize)
11107            .unwrap_or(2);
11108        let mut graphs: std::collections::HashMap<
11109            ((bool, usize), (bool, usize), bool, bool),
11110            (
11111                cudarc::driver::CudaGraph,
11112                Vec<Box<dyn std::any::Any + Send>>,
11113            ),
11114        > = Default::default();
11115        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
11116        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
11117        let mut slots = self.g4_dc_slots(e)?;
11118        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
11119        // baked at the door entry (the modulo keeps every capture valid indefinitely).
11120        const RING: usize = 64;
11121        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
11122        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
11123        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
11124        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
11125        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
11126        const DRAIN: usize = 1;
11127        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
11128        let ring_base = prompt_pos;
11129        let mut out = Vec::with_capacity(max_new);
11130        let mut reason = StopReason::MaxNew;
11131        let mut next = first_token;
11132        let mut captures = 0usize;
11133        for _ in 0..max_new {
11134            out.push(next);
11135            if eos.contains(&next) {
11136                reason = StopReason::Eos;
11137                break;
11138            }
11139            if !on_token(next) {
11140                reason = StopReason::Callback;
11141                break;
11142            }
11143            let t_kv = cache.pos + 1;
11144            // Bucket key per ARM (graph arc step 3):
11145            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
11146            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
11147            //    the component collapses to a single marker).
11148            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
11149            //    at/above it — the kernel derives splits from len_d per replay, so buckets
11150            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
11151            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11152            let f512 = crate::fa512_min_tkv();
11153            let key_s = if t_kv > win {
11154                (true, usize::MAX)
11155            } else {
11156                e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on())
11157            };
11158            let (key_g, rung_end) = if t_kv >= f512 {
11159                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
11160                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
11161                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
11162                ((true, end), end)
11163            } else {
11164                (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv)
11165            };
11166            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
11167            if !graphs.contains_key(&key) {
11168                let bucket_max = (t_kv, rung_end);
11169                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
11170                let snap = cache.snapshot(e)?;
11171                let pos_save = e.dtoh_i32_one(&pos_d)?;
11172                let len_save: Vec<Option<i32>> = cache
11173                    .kv
11174                    .iter()
11175                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap()))
11176                    .collect();
11177                let tok_save = e.dtoh_u32_one(&token_d)?;
11178                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
11179                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
11180                // regression class, and this door's measured -8.8%. The keeper pins warmup
11181                // transients so the captured graph holds kernel nodes only.
11182                let graph = {
11183                    let tok_ref = &mut token_d;
11184                    let pos_ref = &mut pos_d;
11185                    let cache_ref = &mut *cache;
11186                    let slots_ref = &mut slots;
11187                    let ring_ref = &mut ring;
11188                    e.capture_graph_retained_flags(
11189                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
11190                        |e| {
11191                        // self-feeding: the argmax writes token_d itself.
11192                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
11193                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
11194                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
11195                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
11196                                                           cache_ref, n_vocab, Some(bucket_max),
11197                                                           sl, tok_ref, Some((rg, ring_base)))
11198                    })?
11199                };
11200                cache.rollback(e, &snap, 0)?;
11201                e.set_i32_one(&mut pos_d, pos_save)?;
11202                for (il, ls) in len_save.iter().enumerate() {
11203                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
11204                        e.set_i32_one(&mut kvl.len_d, *v)?;
11205                    }
11206                }
11207                e.set_u32_one(&mut token_d, tok_save)?;
11208                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
11209                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
11210                        eprintln!("[graph-census] {c:?}");
11211                    }
11212                }
11213                graphs.insert(key, graph);
11214                captures += 1;
11215            }
11216            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
11217            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
11218            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
11219            // the budget; capture warmups already emitted their tokens through the ring.
11220            let mut chunk = 1usize;
11221            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN")
11222                .ok()
11223                .and_then(|v| v.parse().ok())
11224                .unwrap_or(DRAIN);
11225            while chunk < drain_cap && out.len() + chunk < max_new {
11226                let t_next = cache.pos + 1 + chunk;
11227                let key_s2 = if t_next > win {
11228                    (true, usize::MAX)
11229                } else {
11230                    e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on())
11231                };
11232                let key_g2 = if t_next >= f512 {
11233                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
11234                } else {
11235                    e.fa_bucket_key(t_next, hd_g, nkv_g, false)
11236                };
11237                if (key_s2, key_g2, t_next >= f512, t_next > win) != key {
11238                    break;
11239                }
11240                chunk += 1;
11241            }
11242            let g = &graphs.get(&key).unwrap().0;
11243            for _ in 0..chunk {
11244                g.launch()?;
11245            }
11246            e.stream().synchronize()?;
11247            let ringh = e.dtoh_u32(&ring)?;
11248            for j in 0..chunk {
11249                let pos_j = cache.pos + j;
11250                let tok_j = ringh[(pos_j - ring_base) % RING];
11251                cache.pos += 0; // advanced below in one shot
11252                if j + 1 == chunk {
11253                    next = tok_j;
11254                } else {
11255                    out.push(tok_j);
11256                    if eos.contains(&tok_j) || !on_token(tok_j) {
11257                        reason = if eos.contains(&tok_j) {
11258                            StopReason::Eos
11259                        } else {
11260                            StopReason::Callback
11261                        };
11262                        // roll device/host state back to the stop point.
11263                        let keep = cache.pos + j + 1;
11264                        e.set_i32_one(&mut pos_d, keep as i32)?;
11265                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11266                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
11267                            kvl.len = keep;
11268                        }
11269                        cache.pos = keep;
11270                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11271                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11272                        }
11273                        return Ok((out, reason));
11274                    }
11275                }
11276            }
11277            cache.pos += chunk;
11278            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
11279                kvl.len += chunk;
11280            }
11281        }
11282        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
11283            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
11284        }
11285        Ok((out, reason))
11286    }
11287
11288    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
11289    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
11290    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
11291    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
11292    /// logits (host) + advances cache.pos by t.
11293    pub(crate) fn gemma4_decode_step_t(
11294        &self,
11295        e: &Engine,
11296        tokens: &[u32],
11297        pos0: usize,
11298        cache: &mut Cache,
11299    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
11300        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
11301    }
11302
11303    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
11304    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
11305    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
11306    pub(crate) fn gemma4_decode_step_t_am(
11307        &self,
11308        e: &Engine,
11309        tokens: &[u32],
11310        pos0: usize,
11311        cache: &mut Cache,
11312    ) -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11313        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11314        let t = tokens.len();
11315        let n_vocab = self.output.out_features();
11316        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
11317        for i in 0..t {
11318            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
11319        }
11320        Ok((e.dtoh_u32(&toks)?, hn))
11321    }
11322
11323    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
11324    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
11325    pub(crate) fn gemma4_decode_step_t_am_dev(
11326        &self,
11327        e: &Engine,
11328        tok_d: &CudaSlice<u32>,
11329        t: usize,
11330        pos0: usize,
11331        cache: &mut Cache,
11332    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11333        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
11334        let n_vocab = self.output.out_features();
11335        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11336        for i in 0..t {
11337            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11338        }
11339        Ok((vam, hn))
11340    }
11341
11342    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
11343    /// llama's h_nextn convention).
11344    pub(crate) fn gemma4_decode_step_t_h(
11345        &self,
11346        e: &Engine,
11347        tokens: &[u32],
11348        pos0: usize,
11349        cache: &mut Cache,
11350    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11351        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
11352        let t = tokens.len();
11353        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
11354        e.softcap(&mut ld, cap, t * self.output.out_features())?;
11355        Ok((e.dtoh(&ld)?, hn))
11356    }
11357
11358    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
11359    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
11360    pub(crate) fn verify_stream_scratch(
11361        &self,
11362        e: &Engine,
11363        cap: usize,
11364    ) -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
11365        Ok(VerifyStreamScratch {
11366            pos_d: e.htod_i32(&vec![0i32; cap])?,
11367            row_ctrs: (0..cap)
11368                .map(|_| e.htod_i32(&[0]))
11369                .collect::<Result<_, _>>()?,
11370        })
11371    }
11372
11373    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
11374    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
11375    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
11376    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
11377    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
11378    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
11379    /// sync, exactly the turnaround the burst exists to remove.
11380    pub(crate) fn gemma4_verify_t_am_stream(
11381        &self,
11382        e: &Engine,
11383        tok_d: &CudaSlice<u32>,
11384        t: usize,
11385        ctr: &CudaSlice<i32>,
11386        hint: usize,
11387        cache: &mut Cache,
11388        scr: &mut VerifyStreamScratch,
11389    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11390        let n_embd = self.cfg.n_embd as usize;
11391        let eps = self.cfg.rms_eps;
11392        assert!(t <= scr.row_ctrs.len() && t <= 64);
11393        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
11394        for i in 0..t {
11395            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
11396        }
11397        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
11398        let embd_gpu = self
11399            .embd_gpu
11400            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11401        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11402        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
11403        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11404        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11405        let n_layers = self.layers.len();
11406        for (il, layer) in self.layers.iter().enumerate() {
11407            let (hq, hdq) = match h_carry.take() {
11408                Some(p) => p,
11409                None => {
11410                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11411                }
11412            };
11413            let Mixer::Full(fa) = &layer.mixer else {
11414                panic!("gemma4 layer {il} not full-attn")
11415            };
11416            let o = self
11417                .gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache, hint, row_ctrs)?;
11418            let mut cur = e.uninit(t * n_embd)?;
11419            e.rms_norm(
11420                &o,
11421                layer.post_attn_norm.float_data(),
11422                &mut cur,
11423                n_embd,
11424                t,
11425                eps,
11426            )?;
11427            let next_norm = if il + 1 < n_layers {
11428                Some(self.layers[il + 1].attn_norm.float_data())
11429            } else {
11430                None
11431            };
11432            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
11433            x = xn;
11434            h_carry = hn;
11435            self.dflash_tap(e, cache, il, &x, t)?;
11436        }
11437        let mut hn = e.uninit(t * n_embd)?;
11438        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11439        let ld = e.matmul(&self.output, &hn, t)?;
11440        let n_vocab = self.output.out_features();
11441        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
11442        for i in 0..t {
11443            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
11444        }
11445        Ok((vam, hn))
11446    }
11447
11448    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
11449    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
11450    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
11451    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
11452    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
11453    /// kernel later if it shows in the profile).
11454    fn dflash_tap(
11455        &self,
11456        e: &Engine,
11457        cache: &mut Cache,
11458        il: usize,
11459        x: &CudaSlice<f32>,
11460        t: usize,
11461    ) -> Result<(), Box<dyn std::error::Error>> {
11462        let Some(taps) = cache.dflash_taps.as_mut() else {
11463            return Ok(());
11464        };
11465        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else {
11466            return Ok(());
11467        };
11468        let h = taps.hidden;
11469        let n_taps = taps.layer_ids.len();
11470        debug_assert_eq!(taps.t, t);
11471        let xv = e.view(x, t * h);
11472        for r in 0..t {
11473            let row = xv.slice(r * h..(r + 1) * h);
11474            e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
11475        }
11476        Ok(())
11477    }
11478
11479    fn gemma4_verify_trunk(
11480        &self,
11481        e: &Engine,
11482        tokens: &[u32],
11483        pos0: usize,
11484        cache: &mut Cache,
11485        tok_dev: Option<&CudaSlice<u32>>,
11486    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11487        let n_embd = self.cfg.n_embd as usize;
11488        let eps = self.cfg.rms_eps;
11489        let t = tokens.len();
11490        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
11491        let pos_d = e.htod_i32(&pos)?;
11492        let mut x = match tok_dev {
11493            Some(td) => {
11494                let embd_gpu = self
11495                    .embd_gpu
11496                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
11497                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
11498                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
11499            }
11500            None => e.htod(&self.embd.gather(n_embd, tokens))?,
11501        };
11502        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
11503        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
11504        let n_layers = self.layers.len();
11505        for (il, layer) in self.layers.iter().enumerate() {
11506            let (hq, hdq) = match h_carry.take() {
11507                Some(p) => p,
11508                None => {
11509                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?
11510                }
11511            };
11512            let Mixer::Full(fa) = &layer.mixer else {
11513                panic!("gemma4 layer {il} not full-attn")
11514            };
11515            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
11516            let mut cur = e.uninit(t * n_embd)?;
11517            e.rms_norm(
11518                &o,
11519                layer.post_attn_norm.float_data(),
11520                &mut cur,
11521                n_embd,
11522                t,
11523                eps,
11524            )?;
11525            let next_norm = if il + 1 < n_layers {
11526                Some(self.layers[il + 1].attn_norm.float_data())
11527            } else {
11528                None
11529            };
11530            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
11531            x = xn;
11532            h_carry = hn;
11533            self.dflash_tap(e, cache, il, &x, t)?;
11534        }
11535        let mut hn = e.uninit(t * n_embd)?;
11536        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
11537        let mut ld = e.matmul(&self.output, &hn, t)?;
11538        self.gemma4_suppress(e, &mut ld, t)?; // before the per-row argmax consumers
11539        cache.pos += t;
11540        Ok((ld, hn))
11541    }
11542
11543    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
11544    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
11545    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
11546    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
11547    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
11548    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
11549    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
11550    #[allow(clippy::too_many_arguments)]
11551    fn gemma4_verify_attn_stream(
11552        &self,
11553        e: &Engine,
11554        fa: &crate::hybrid::FullAttnLayer,
11555        il: usize,
11556        hq: &CudaSlice<i8>,
11557        hdq: &CudaSlice<f32>,
11558        pos_d: &CudaSlice<i32>,
11559        t: usize,
11560        cache: &mut Cache,
11561        hint: usize,
11562        row_ctrs: &[CudaSlice<i32>],
11563    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11564        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11565        let eps = self.cfg.rms_eps;
11566        let aux = self.gemma4_aux.as_ref().unwrap();
11567        let ones = aux.ones(e);
11568        #[cfg(debug_assertions)]
11569        crate::debug_assert_tensor_stream_device(
11570            ones,
11571            &e.stream(),
11572            "gemma4_verify_attn_stream.ones",
11573        );
11574        let h0 = e.zeros(0)?;
11575        let h = &h0;
11576        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
11577        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
11578        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11579        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11580        let fused_qkv = if f2b {
11581            if swa {
11582                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11583                    .map(|(a, b, c)| (a, b, Some(c)))
11584            } else {
11585                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11586                    .map(|(a, b)| (a, b, None))
11587            }
11588        } else {
11589            None
11590        };
11591        let (q0, k0, v0) = match fused_qkv {
11592            Some((a, b, cv)) => {
11593                let v = match cv {
11594                    Some(c) => c,
11595                    None => e.clone_dtod(&b)?,
11596                };
11597                (a, b, v)
11598            }
11599            None => {
11600                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11601                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11602                let v0 = if swa {
11603                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11604                } else {
11605                    e.clone_dtod(&k0)?
11606                };
11607                (q0, k0, v0)
11608            }
11609        };
11610        let mut q = e.uninit(t * nh * hd)?;
11611        let mut k = e.uninit(t * nkv * hd)?;
11612        let mut v = e.uninit(t * nkv * hd)?;
11613        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
11614        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
11615        let ff = if swa {
11616            None
11617        } else {
11618            Some(
11619                aux.rope_freqs(e)
11620                    .expect("gemma4 global rope needs rope_freqs.weight"),
11621            )
11622        };
11623        #[cfg(debug_assertions)]
11624        if let Some(ff) = ff {
11625            crate::debug_assert_tensor_stream_device(
11626                ff,
11627                &e.stream(),
11628                "gemma4_verify_attn_stream.rope_freqs",
11629            );
11630        }
11631        e.rms_norm_qkv_rope(
11632            &q0,
11633            &k0,
11634            &v0,
11635            fa.q_norm.float_data(),
11636            fa.k_norm.float_data(),
11637            ones,
11638            &mut q,
11639            &mut k,
11640            &mut v,
11641            hd,
11642            nh * t,
11643            nkv * t,
11644            pos_d,
11645            nh,
11646            nkv,
11647            base,
11648            1.0,
11649            ff,
11650            eps,
11651        )?;
11652        let kvl = cache.kv[il].as_mut().unwrap();
11653        // append at the DEVICE slot; the counter advances by t on-device.
11654        e.append_kv_quantized_rows_dc(
11655            &k,
11656            &v,
11657            &mut kvl.k,
11658            &mut kvl.v,
11659            &kvl.len_d,
11660            t,
11661            kvl.kv_dim_k,
11662            kvl.kv_dim_v,
11663            kvl.k_tok_bytes,
11664            kvl.v_tok_bytes,
11665            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
11666        )?;
11667        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
11668        // the sole len writer after this round's attention (base stays = old len, plus = 0).
11669        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11670        let mut attn = e.uninit(t * nh * hd)?;
11671        let k_view = e.view_u8(&kvl.k, kvl.k.len());
11672        let v_view = e.view_u8(&kvl.v, kvl.v.len());
11673        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
11674        // and a stable window regime — the same rung/regime keys as the draft graph).
11675        if swa && hint + 1 >= win {
11676            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
11677            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
11678            e.fa_decode_rows_w(
11679                &q,
11680                &k_view,
11681                &v_view,
11682                &mut attn,
11683                hd,
11684                nh,
11685                nkv,
11686                &kvl.len_d,
11687                0,
11688                t,
11689                scale,
11690                win,
11691                kvl.k_tok_bytes,
11692                kvl.v_tok_bytes,
11693                None,
11694            )?;
11695        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
11696            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
11697            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
11698            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
11699            // Burst entry gates the horizon onto one side of the crossover, so hint decides
11700            // for every row.
11701            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
11702            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
11703            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
11704            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
11705            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
11706            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
11707            // any bucket >= the live length is exact.
11708            let bucket = (hint + t + 2)
11709                .next_power_of_two()
11710                .min(crate::fa512_min_tkv().saturating_sub(1));
11711            let qv = e.view(&q, t * nh * hd);
11712            for i in 0..t {
11713                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
11714                let mut q_one = e.uninit(nh * hd)?;
11715                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
11716                let mut a_one = e.uninit(nh * hd)?;
11717                e.fa_decode_dc(
11718                    &q_one,
11719                    &k_view,
11720                    &v_view,
11721                    &mut a_one,
11722                    hd,
11723                    nh,
11724                    nkv,
11725                    &row_ctrs[i],
11726                    bucket,
11727                    scale,
11728                    kvl.k_tok_bytes,
11729                    kvl.v_tok_bytes,
11730                    false,
11731                )?;
11732                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
11733            }
11734        } else if hd == 512 {
11735            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
11736            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
11737            e.fa_decode_rows(
11738                &q,
11739                &k_view,
11740                &v_view,
11741                &mut attn,
11742                hd,
11743                nh,
11744                nkv,
11745                hint,
11746                t,
11747                scale,
11748                kvl.k_tok_bytes,
11749                kvl.v_tok_bytes,
11750                Some((&kvl.len_d, 0)),
11751                false,
11752                false,
11753                None,
11754            )?;
11755        } else {
11756            // hd256 under-window: v4 device-len rows twin.
11757            e.fa_decode_rows_dc(
11758                &q,
11759                &k_view,
11760                &v_view,
11761                &mut attn,
11762                hd,
11763                nh,
11764                nkv,
11765                &kvl.len_d,
11766                hint + t,
11767                t,
11768                scale,
11769                kvl.k_tok_bytes,
11770                kvl.v_tok_bytes,
11771                0,
11772                swa && crate::Engine::wkv_on(),
11773            )?;
11774        }
11775        Ok(e.matmul(&fa.wo, &attn, t)?)
11776    }
11777
11778    fn gemma4_verify_attn(
11779        &self,
11780        e: &Engine,
11781        fa: &crate::hybrid::FullAttnLayer,
11782        il: usize,
11783        hq: &CudaSlice<i8>,
11784        hdq: &CudaSlice<f32>,
11785        pos_d: &CudaSlice<i32>,
11786        t: usize,
11787        cache: &mut Cache,
11788    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11789        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
11790        let eps = self.cfg.rms_eps;
11791        let aux = self.gemma4_aux.as_ref().unwrap();
11792        let ones = aux.ones(e);
11793        #[cfg(debug_assertions)]
11794        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_verify_attn.ones");
11795        let n_embd = self.cfg.n_embd as usize;
11796        let _ = n_embd;
11797
11798        let h0 = e.zeros(0)?;
11799        let h = &h0;
11800        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
11801        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
11802        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11803        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
11804        let fused_qkv = if f2b {
11805            if swa {
11806                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
11807                    .map(|(a, b, c)| (a, b, Some(c)))
11808            } else {
11809                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
11810                    .map(|(a, b)| (a, b, None))
11811            }
11812        } else {
11813            None
11814        };
11815        let (q0, k0, v0) = match fused_qkv {
11816            Some((a, b, cv)) => {
11817                let v = match cv {
11818                    Some(c) => c,
11819                    None => e.clone_dtod(&b)?,
11820                };
11821                (a, b, v)
11822            }
11823            None => {
11824                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
11825                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
11826                let v0 = if swa {
11827                    e.matmul_pre(&fa.wv, hq, hdq, h, t)?
11828                } else {
11829                    e.clone_dtod(&k0)?
11830                };
11831                (q0, k0, v0)
11832            }
11833        };
11834        let mut q = e.uninit(t * nh * hd)?;
11835        let mut k = e.uninit(t * nkv * hd)?;
11836        let mut v = e.uninit(t * nkv * hd)?;
11837        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
11838        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
11839        let ff = if swa {
11840            None
11841        } else {
11842            Some(
11843                aux.rope_freqs(e)
11844                    .expect("gemma4 global rope needs rope_freqs.weight"),
11845            )
11846        };
11847        #[cfg(debug_assertions)]
11848        if let Some(ff) = ff {
11849            crate::debug_assert_tensor_stream_device(
11850                ff,
11851                &e.stream(),
11852                "gemma4_verify_attn.rope_freqs",
11853            );
11854        }
11855        e.rms_norm_qkv_rope(
11856            &q0,
11857            &k0,
11858            &v0,
11859            fa.q_norm.float_data(),
11860            fa.k_norm.float_data(),
11861            ones,
11862            &mut q,
11863            &mut k,
11864            &mut v,
11865            hd,
11866            nh * t,
11867            nkv * t,
11868            pos_d,
11869            nh,
11870            nkv,
11871            base,
11872            1.0,
11873            ff,
11874            eps,
11875        )?;
11876        let kvl = cache.kv[il].as_mut().unwrap();
11877        let base_len = kvl.len;
11878        e.append_kv_quantized_rows(
11879            &k,
11880            &v,
11881            &mut kvl.k,
11882            &mut kvl.v,
11883            base_len,
11884            t,
11885            kvl.kv_dim_k,
11886            kvl.kv_dim_v,
11887            kvl.k_tok_bytes,
11888            kvl.v_tok_bytes,
11889            (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
11890        )?;
11891        kvl.len += t;
11892        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
11893        let mut attn = e.uninit(t * nh * hd)?;
11894        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
11895        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
11896        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
11897            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
11898            // decode rides the SAME symbol at t=1 (parity law).
11899            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
11900        if rows_ok && (!swa || base_len + t <= win) {
11901            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
11902            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
11903            if hd == 512 {
11904                // device-len twin: sync the counter to the verify base (async arg-store).
11905                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11906                e.fa_decode_rows(
11907                    &q,
11908                    &k_view,
11909                    &v_view,
11910                    &mut attn,
11911                    hd,
11912                    nh,
11913                    nkv,
11914                    base_len,
11915                    t,
11916                    scale,
11917                    kvl.k_tok_bytes,
11918                    kvl.v_tok_bytes,
11919                    Some((&kvl.len_d, 0)),
11920                    false,
11921                    swa && crate::Engine::wkv_on(),
11922                    None,
11923                )?;
11924            } else {
11925                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
11926                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
11927                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
11928                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11929                e.fa_decode_rows_dc(
11930                    &q,
11931                    &k_view,
11932                    &v_view,
11933                    &mut attn,
11934                    hd,
11935                    nh,
11936                    nkv,
11937                    &kvl.len_d,
11938                    base_len + t,
11939                    t,
11940                    scale,
11941                    kvl.k_tok_bytes,
11942                    kvl.v_tok_bytes,
11943                    0,
11944                    swa && crate::Engine::wkv_on(),
11945                )?;
11946            }
11947            return Ok(e.matmul(&fa.wo, &attn, t)?);
11948        }
11949        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
11950        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
11951        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
11952        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
11953        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
11954        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
11955        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
11956        if hd == 256
11957            && swa
11958            && base_len + 1 >= win
11959            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
11960        {
11961            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
11962            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
11963            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
11964            e.fa_decode_rows_w(
11965                &q,
11966                &k_view,
11967                &v_view,
11968                &mut attn,
11969                hd,
11970                nh,
11971                nkv,
11972                &kvl.len_d,
11973                0,
11974                t,
11975                scale,
11976                win,
11977                kvl.k_tok_bytes,
11978                kvl.v_tok_bytes,
11979                None,
11980            )?;
11981            return Ok(e.matmul(&fa.wo, &attn, t)?);
11982        }
11983        for i in 0..t {
11984            let avail = base_len + i + 1;
11985            let (off_tok, t_kv) = if swa && avail > win {
11986                (avail - win, win)
11987            } else {
11988                (0, avail)
11989            };
11990            let k_view = e.view_u8_range(
11991                &kvl.k,
11992                off_tok * kvl.k_tok_bytes,
11993                (off_tok + t_kv) * kvl.k_tok_bytes,
11994            );
11995            let v_view = e.view_u8_range(
11996                &kvl.v,
11997                off_tok * kvl.v_tok_bytes,
11998                (off_tok + t_kv) * kvl.v_tok_bytes,
11999            );
12000            let qi = e.view(&q, t * nh * hd);
12001            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
12002            let mut q_one = e.uninit(nh * hd)?;
12003            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
12004            let mut a_one = e.uninit(nh * hd)?;
12005            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
12006            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
12007            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
12008            if swa
12009                && avail > win
12010                && hd == 256
12011                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12012            {
12013                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12014                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12015                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12016                e.fa_decode_rows_w(
12017                    &q_one,
12018                    &kp,
12019                    &vp,
12020                    &mut a_one,
12021                    hd,
12022                    nh,
12023                    nkv,
12024                    &kvl.len_d,
12025                    0,
12026                    1,
12027                    scale,
12028                    win,
12029                    kvl.k_tok_bytes,
12030                    kvl.v_tok_bytes,
12031                    None,
12032                )?;
12033            } else if !swa
12034                && hd == 512
12035                && avail >= crate::fa512_min_tkv()
12036                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0")
12037            {
12038                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
12039                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
12040                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
12041                e.fa_decode_rows(
12042                    &q_one,
12043                    &kp,
12044                    &vp,
12045                    &mut a_one,
12046                    hd,
12047                    nh,
12048                    nkv,
12049                    avail - 1,
12050                    1,
12051                    scale,
12052                    kvl.k_tok_bytes,
12053                    kvl.v_tok_bytes,
12054                    Some((&kvl.len_d, 0)),
12055                    false,
12056                    false,
12057                    None,
12058                )?;
12059            } else {
12060                e.fa_decode_kvmod(
12061                    &q_one,
12062                    &k_view,
12063                    &v_view,
12064                    &mut a_one,
12065                    hd,
12066                    nh,
12067                    nkv,
12068                    t_kv,
12069                    scale,
12070                    kvl.k_tok_bytes,
12071                    kvl.v_tok_bytes,
12072                    swa && crate::Engine::wkv_on(),
12073                )?;
12074            }
12075            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
12076        }
12077        Ok(e.matmul(&fa.wo, &attn, t)?)
12078    }
12079
12080    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
12081    /// h_seed = pre-output_norm hidden). Advances cache.pos.
12082    pub(crate) fn gemma4_decode_step_h(
12083        &self,
12084        e: &Engine,
12085        token: u32,
12086        cache: &mut Cache,
12087    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12088        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
12089        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
12090        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
12091        // unsplit rather than guessing a fence.
12092        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
12093            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
12094        }
12095        if crate::pp::pp_cuts(self.layers.len()).is_some() {
12096            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
12097        }
12098        let n_embd = self.cfg.n_embd as usize;
12099        let eps = self.cfg.rms_eps;
12100        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12101        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12102        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12103        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
12104        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
12105        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12106        let n_layers = self.layers.len();
12107        for (il, layer) in self.layers.iter().enumerate() {
12108            let (hq, hdq) = match h_carry.take() {
12109                Some(p) => p,
12110                None => {
12111                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?
12112                }
12113            };
12114            let Mixer::Full(fa) = &layer.mixer else {
12115                panic!("gemma4 layer {il} not full-attn")
12116            };
12117            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
12118            let mut cur = e.uninit(n_embd)?;
12119            e.rms_norm(
12120                &o,
12121                layer.post_attn_norm.float_data(),
12122                &mut cur,
12123                n_embd,
12124                1,
12125                eps,
12126            )?;
12127            let next_norm = if il + 1 < n_layers {
12128                Some(self.layers[il + 1].attn_norm.float_data())
12129            } else {
12130                None
12131            };
12132            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
12133            x = xn;
12134            h_carry = hn;
12135        }
12136        let mut hn = e.uninit(n_embd)?;
12137        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12138        let h_seed = e.clone_dtod(&x)?;
12139        let mut ld = e.matmul(&self.output, &hn, 1)?;
12140        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12141        e.softcap(&mut ld, cap, self.output.out_features())?; // R4 on device (262k host tanh ~ms/step)
12142        self.gemma4_suppress(e, &mut ld, 1)?;
12143        let logits = e.dtoh(&ld)?;
12144        cache.pos += 1;
12145        Ok((logits, h_seed))
12146    }
12147
12148    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
12149    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
12150    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
12151    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
12152    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
12153    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
12154    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
12155    fn gemma4_decode_layers(
12156        &self,
12157        e: &Engine,
12158        mut x: CudaSlice<f32>,
12159        lo: usize,
12160        hi: usize,
12161        pos_d: &CudaSlice<i32>,
12162        cache: &mut Cache,
12163    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12164        let n_embd = self.cfg.n_embd as usize;
12165        let eps = self.cfg.rms_eps;
12166        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
12167        for il in lo..hi {
12168            let layer = &self.layers[il];
12169            let (hq, hdq) = match h_carry.take() {
12170                Some(p) => p,
12171                // range head: il == lo — norm against THIS layer's attn_norm.
12172                None => {
12173                    e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?
12174                }
12175            };
12176            let Mixer::Full(fa) = &layer.mixer else {
12177                panic!("gemma4 layer {il} not full-attn")
12178            };
12179            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
12180            let mut cur = e.uninit(n_embd)?;
12181            e.rms_norm(
12182                &o,
12183                layer.post_attn_norm.float_data(),
12184                &mut cur,
12185                n_embd,
12186                1,
12187                eps,
12188            )?;
12189            let next_norm = if il + 1 < hi {
12190                Some(self.layers[il + 1].attn_norm.float_data())
12191            } else {
12192                None
12193            };
12194            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
12195            x = xn;
12196            h_carry = hn;
12197        }
12198        Ok(x)
12199    }
12200
12201    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
12202    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
12203    /// boundary handoff — same choreography as the generic arm (decode.rs), same
12204    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
12205    /// stage 1 = layers [split, n) + output_norm + softcapped head.
12206    /// Each stage uploads its own copy of the step's position scalar on its own stream.
12207    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
12208    fn gemma4_decode_step_h_pp2(
12209        &self,
12210        e: &Engine,
12211        token: u32,
12212        cache: &mut Cache,
12213        split: usize,
12214    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12215        if crate::pp::pp2_streams_off() {
12216            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
12217        }
12218        let rt = crate::pp::Pp2Rt::get(e)?;
12219        let e0 = rt.engine(0, e);
12220        let e1 = rt.engine(1, e);
12221        let n_embd = self.cfg.n_embd as usize;
12222        let eps = self.cfg.rms_eps;
12223        let pos = cache.pos as i32;
12224
12225        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
12226        let slot = {
12227            let _st0 = rt.enter(0);
12228            let pos_d = e0.htod_i32(&[pos])?;
12229            #[cfg(debug_assertions)]
12230            crate::debug_assert_tensor_stream_device(
12231                &pos_d,
12232                &e0.stream(),
12233                "gemma4_decode_step_h_pp2.stage0.pos_d",
12234            );
12235            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
12236            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12237            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
12238            rt.tx(0, &x, n_embd)?
12239        };
12240
12241        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
12242        let _st1 = rt.enter(1);
12243        let pos_d = e1.htod_i32(&[pos])?;
12244        #[cfg(debug_assertions)]
12245        crate::debug_assert_tensor_stream_device(
12246            &pos_d,
12247            &e1.stream(),
12248            "gemma4_decode_step_h_pp2.stage1.pos_d",
12249        );
12250        let x = rt.rx(0, slot, n_embd)?;
12251        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
12252
12253        let mut hn = e1.uninit(n_embd)?;
12254        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12255        let h_seed = e1.clone_dtod(&x)?;
12256        let mut ld = e1.matmul(&self.output, &hn, 1)?;
12257        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12258        e1.softcap(&mut ld, cap, self.output.out_features())?;
12259        self.gemma4_suppress(e1, &mut ld, 1)?;
12260        let logits = e1.dtoh(&ld)?;
12261        cache.pos += 1;
12262        Ok((logits, h_seed))
12263    }
12264
12265    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
12266    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
12267    fn gemma4_decode_step_h_pp2_samestream(
12268        &self,
12269        e: &Engine,
12270        token: u32,
12271        cache: &mut Cache,
12272        split: usize,
12273    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12274        let n_embd = self.cfg.n_embd as usize;
12275        let eps = self.cfg.rms_eps;
12276        let pos_d = e.htod_i32(&[cache.pos as i32])?;
12277
12278        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
12279        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
12280        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
12281        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
12282
12283        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
12284        let boundary_tx = e.clone_dtod(&x)?;
12285        let boundary_rx = e.clone_dtod(&boundary_tx)?;
12286
12287        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
12288        let x =
12289            self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
12290
12291        let mut hn = e.uninit(n_embd)?;
12292        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
12293        let h_seed = e.clone_dtod(&x)?;
12294        let mut ld = e.matmul(&self.output, &hn, 1)?;
12295        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
12296        e.softcap(&mut ld, cap, self.output.out_features())?;
12297        self.gemma4_suppress(e, &mut ld, 1)?;
12298        let logits = e.dtoh(&ld)?;
12299        cache.pos += 1;
12300        Ok((logits, h_seed))
12301    }
12302}
12303
12304// ============================ step35 (Step-3.7-Flash) ==================================
12305// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
12306// FAMILY and not a few branches inside the generic `full_attn*` chain:
12307//
12308//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
12309//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
12310//      shapes and the FA head counts would be wrong on 33 of 45 layers.
12311//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
12312//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
12313//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
12314//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
12315//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
12316//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
12317//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
12318//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
12319//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
12320//
12321// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
12322impl HybridModel {
12323    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
12324    /// synthesize a drafter or trunk layer from a neighboring class.
12325    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
12326        let geometry = self
12327            .cfg
12328            .layer_geometry(il as u32)
12329            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
12330        debug_assert_eq!(
12331            geometry.attention_gate,
12332            memra_gguf::config::AttentionGateKind::SeparateHead
12333        );
12334        geometry
12335    }
12336
12337    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
12338    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
12339    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
12340    ///
12341    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
12342    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
12343    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
12344    /// `cache`:
12345    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
12346    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
12347    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
12348    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
12349    ///     contract, lane/chunkinv-flip).
12350    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
12351    ///     q/k/v, no cache side effect.
12352    ///
12353    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
12354    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
12355    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
12356    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
12357    /// still contains must be masked per query. memra's window convention
12358    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
12359    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
12360    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
12361    ///
12362    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
12363    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
12364    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
12365    ///
12366    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
12367    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
12368    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
12369    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
12370    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
12371    /// hidden rows, and the generated text — a function of the chunk size:
12372    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
12373    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
12374    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
12375    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
12376    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
12377    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
12378    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
12379    ///   one-token change in a documented machine-config knob changed the answer.
12380    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
12381    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
12382    /// the same rows moves the logits by ~1.8.
12383    ///
12384    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
12385    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
12386    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
12387    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
12388    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
12389    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
12390    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
12391    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
12392    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
12393    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
12394    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
12395    /// those with t_kv <= win = 512.
12396    #[allow(clippy::too_many_arguments)]
12397    fn step35_attn_pre_wo(
12398        &self,
12399        e: &Engine,
12400        fa: &FullAttnLayer,
12401        mut g3: Vec<CudaSlice<f32>>,
12402        hg: Option<&CudaSlice<f32>>,
12403        gt_pre: Option<&CudaSlice<f32>>,
12404        pos_d: &CudaSlice<i32>,
12405        t: usize,
12406        cache: Option<&mut Cache>,
12407        il: usize,
12408        seq_end: usize,
12409    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12410        let geometry = self.step35_geom(il);
12411        let hd = geometry.head_dim_k as usize;
12412        let nkv = geometry.n_head_kv as usize;
12413        let nh = geometry.n_head as usize;
12414        let rbase = geometry.rope_base;
12415        let scale = geometry.attention_scale();
12416        let swa = geometry.window.is_some();
12417        let eps = self.cfg.rms_eps;
12418        let win = geometry.window.unwrap_or(0) as usize;
12419        let n_rot = geometry.n_rot as usize;
12420
12421        let v = g3.pop().unwrap();
12422        let k0 = g3.pop().unwrap();
12423        let q0 = g3.pop().unwrap();
12424
12425        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
12426        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
12427        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
12428        let mut q = e.uninit(t * nh * hd)?;
12429        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
12430        let mut k = e.uninit(t * nkv * hd)?;
12431        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
12432        let ff = if geometry.rope_factors {
12433            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12434        } else {
12435            None
12436        };
12437        #[cfg(debug_assertions)]
12438        if let Some(ff) = ff {
12439            crate::debug_assert_tensor_stream_device(
12440                ff,
12441                &e.stream(),
12442                "step35_attn_pre_wo.rope_freqs",
12443            );
12444        }
12445        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
12446
12447        let mut attn = e.uninit(t * nh * hd)?;
12448        match cache {
12449            Some(cache) => {
12450                let base_len = cache.kv[il].as_ref().unwrap().len;
12451                // Read per layer call, never in a measured default.
12452                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
12453                let legacy_calllocal = std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
12454                let off = if swa {
12455                    let raw = base_len.saturating_sub(win - 1);
12456                    if legacy_tkv || legacy_calllocal {
12457                        raw
12458                    } else {
12459                        raw & !31usize
12460                    }
12461                } else {
12462                    0
12463                };
12464                {
12465                    let kvl = cache.kv[il].as_mut().unwrap();
12466                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
12467                    let write_row = e.prepare_kv_append(kvl, off, t)?;
12468                    e.append_kv_quantized_rows(
12469                        &k,
12470                        &v,
12471                        &mut kvl.k,
12472                        &mut kvl.v,
12473                        write_row,
12474                        t,
12475                        kvl.kv_dim_k,
12476                        kvl.kv_dim_v,
12477                        kvl.k_tok_bytes,
12478                        kvl.v_tok_bytes,
12479                        crate::Engine::kv_fp8_on(),
12480                    )?;
12481                    kvl.len += t;
12482                    let new_len = kvl.len as i32;
12483                    e.set_i32_one(&mut kvl.len_d, new_len)?;
12484                }
12485                let kvl = cache.kv[il].as_ref().unwrap();
12486                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
12487                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
12488                // unaligned view offset here. Both halves are load-bearing for the canaries:
12489                // on the FA default the predicate arms agree bitwise wherever they can differ
12490                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
12491                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
12492                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
12493                // on the current FA path: its tile grid starts at the chunk/call boundary.
12494                // SWA: trim the view to the oldest key any query in this chunk can reach —
12495                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
12496                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
12497                // kernel's online-softmax recurrence groups keys into BK tiles relative to
12498                // the VIEW START — so an unaligned off regroups the same absolute keys into
12499                // different tiles at different chunk sizes = different (m,l) rounding =
12500                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
12501                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
12502                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
12503                // size; the <=31 extra leading keys are older than EVERY query's window
12504                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
12505                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
12506                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
12507                // the floor arm's bits do not move either (gated: G2f, battery 2).
12508                let t_kv = base_len + t - off;
12509                let physical = kvl.physical_rows(off, off + t_kv)?;
12510                let k_view = e.view_u8_range(
12511                    &kvl.k,
12512                    physical.start * kvl.k_tok_bytes,
12513                    physical.end * kvl.k_tok_bytes,
12514                );
12515                let v_view = e.view_u8_range(
12516                    &kvl.v,
12517                    physical.start * kvl.v_tok_bytes,
12518                    physical.end * kvl.v_tok_bytes,
12519                );
12520                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
12521                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
12522                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
12523                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
12524                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
12525                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
12526                // construction, so the invariance assertion MUST break under it (the seam whose
12527                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
12528                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
12529                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
12530                // cached (probes flip it in-process). Never on in a measured default run.
12531                let swa_naive = if legacy_tkv {
12532                    t_kv > win
12533                } else {
12534                    seq_end > win
12535                };
12536                if swa && swa_naive {
12537                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
12538                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
12539                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
12540                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
12541                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
12542                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
12543                    // identically to the unwindowed one modulo the mask, which is the point.
12544                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
12545                    // selected on `seq_end` like every arm here, so the class is uniform for
12546                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
12547                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
12548                    // the f32 floor (the previous numeric config, kept as the A/B seam).
12549                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
12550                        e.sdpa_naive_w_quantized_view(
12551                            &q,
12552                            &k_view,
12553                            &v_view,
12554                            &mut attn,
12555                            hd,
12556                            nh,
12557                            nkv,
12558                            t,
12559                            t_kv,
12560                            scale,
12561                            true,
12562                            win,
12563                            kvl.k_tok_bytes,
12564                            kvl.v_tok_bytes,
12565                        )?;
12566                    } else {
12567                        e.fa_prefill_view_ws_w_hd128(
12568                            &q,
12569                            &k_view,
12570                            &v_view,
12571                            &mut attn,
12572                            hd,
12573                            nh,
12574                            nkv,
12575                            t,
12576                            t_kv,
12577                            scale,
12578                            true,
12579                            win,
12580                            kvl.k_tok_bytes,
12581                            kvl.v_tok_bytes,
12582                        )?;
12583                    }
12584                } else if std::env::var("MEMRA_NOFA").is_ok() {
12585                    e.sdpa_naive_quantized_view(
12586                        &q,
12587                        &k_view,
12588                        &v_view,
12589                        &mut attn,
12590                        hd,
12591                        nh,
12592                        nkv,
12593                        t,
12594                        t_kv,
12595                        scale,
12596                        true,
12597                        kvl.k_tok_bytes,
12598                        kvl.v_tok_bytes,
12599                    )?;
12600                } else {
12601                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
12602                    // reach past the window, so the window mask is a no-op under causal and every
12603                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
12604                    // request either way, which is what makes the chunk size arithmetic-free.
12605                    e.fa_prefill_view_ws(
12606                        &q,
12607                        &k_view,
12608                        &v_view,
12609                        &mut attn,
12610                        hd,
12611                        nh,
12612                        nkv,
12613                        t,
12614                        t_kv,
12615                        scale,
12616                        true,
12617                        kvl.k_tok_bytes,
12618                        kvl.v_tok_bytes,
12619                        crate::Engine::kv_fp8_on(),
12620                    )?;
12621                }
12622            }
12623            None => {
12624                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
12625                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
12626                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
12627                // seq_end here too or it re-opens the same door.
12628                debug_assert_eq!(
12629                    seq_end, t,
12630                    "step35 cacheless prefill is monolithic (seq_end == t)"
12631                );
12632                if swa && seq_end > win {
12633                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
12634                } else if std::env::var("MEMRA_NOFA").is_ok() {
12635                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12636                } else {
12637                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
12638                }
12639            }
12640        }
12641
12642        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
12643        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
12644        let gw = fa
12645            .attn_gate
12646            .as_ref()
12647            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12648        let gt_owned = if gt_pre.is_none() {
12649            Some(e.matmul(
12650                gw,
12651                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
12652                t,
12653            )?)
12654        } else {
12655            None
12656        };
12657        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
12658        let mut ag = e.uninit(t * nh * hd)?;
12659        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
12660        Ok(ag)
12661    }
12662
12663    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
12664    /// `forward_last`, t2probe). Post-`wo`.
12665    pub(crate) fn step35_attn(
12666        &self,
12667        e: &Engine,
12668        fa: &FullAttnLayer,
12669        h: &CudaSlice<f32>,
12670        pos_d: &CudaSlice<i32>,
12671        t: usize,
12672        il: usize,
12673    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12674        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
12675        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
12676        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
12677        Ok(e.matmul(&fa.wo, &ag, t)?)
12678    }
12679
12680    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
12681    /// resident quantized cache, attend through the cache view). Post-`wo`.
12682    ///
12683    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
12684    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
12685    /// own extent.
12686    #[allow(clippy::too_many_arguments)]
12687    pub(crate) fn step35_attn_prime(
12688        &self,
12689        e: &Engine,
12690        fa: &FullAttnLayer,
12691        h: &CudaSlice<f32>,
12692        hx: Option<&CudaSlice<u8>>,
12693        pos_d: &CudaSlice<i32>,
12694        t: usize,
12695        cache: &mut Cache,
12696        il: usize,
12697        seq_end: usize,
12698    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12699        let g3 = match hx {
12700            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
12701            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
12702        };
12703        let ag =
12704            self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, Some(cache), il, seq_end)?;
12705        Ok(e.matmul(&fa.wo, &ag, t)?)
12706    }
12707
12708    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
12709    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
12710    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
12711    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
12712    /// requiring `attn_gate`).
12713    ///
12714    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
12715    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
12716    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
12717    #[allow(clippy::too_many_arguments)]
12718    pub(crate) fn step35_decode_attn(
12719        &self,
12720        e: &Engine,
12721        fa: &FullAttnLayer,
12722        il: usize,
12723        h: &CudaSlice<f32>,
12724        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
12725        pos_d: &CudaSlice<i32>,
12726        cache: &mut Cache,
12727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12728        let geometry = self.step35_geom(il);
12729        let hd = geometry.head_dim_k as usize;
12730        let nkv = geometry.n_head_kv as usize;
12731        let nh = geometry.n_head as usize;
12732        let rbase = geometry.rope_base;
12733        let scale = geometry.attention_scale();
12734        let swa = geometry.window.is_some();
12735        let eps = self.cfg.rms_eps;
12736        let win = geometry.window.unwrap_or(0) as usize;
12737        let n_rot = geometry.n_rot as usize;
12738        let n_embd = self.cfg.n_embd as usize;
12739        let gw = fa
12740            .attn_gate
12741            .as_ref()
12742            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
12743
12744        let (q0, k0, v0, gt) = match pre_q {
12745            Some((hq, hdq)) => {
12746                debug_assert!(
12747                    e.uses_q8_1_fast(gw),
12748                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
12749                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast"
12750                );
12751                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
12752                    Some(t3) => t3,
12753                    None => (
12754                        e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
12755                        e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
12756                        e.matmul_pre(&fa.wv, hq, hdq, h, 1)?,
12757                    ),
12758                };
12759                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
12760                (a, b, c, gt)
12761            }
12762            None => {
12763                if e.uses_q8_1_fast(&fa.wq)
12764                    && e.uses_q8_1_fast(&fa.wk)
12765                    && e.uses_q8_1_fast(&fa.wv)
12766                    && e.uses_q8_1_fast(gw)
12767                {
12768                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
12769                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
12770                        Some(t3) => t3,
12771                        None => (
12772                            e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
12773                            e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
12774                            e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
12775                        ),
12776                    };
12777                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
12778                    (a, b, c, gt)
12779                } else {
12780                    (
12781                        e.matmul(&fa.wq, h, 1)?,
12782                        e.matmul(&fa.wk, h, 1)?,
12783                        e.matmul(&fa.wv, h, 1)?,
12784                        e.matmul(gw, h, 1)?,
12785                    )
12786                }
12787            }
12788        };
12789
12790        let mut q = e.uninit(nh * hd)?;
12791        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
12792        let mut k = e.uninit(nkv * hd)?;
12793        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
12794        let ff = if swa {
12795            None
12796        } else {
12797            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
12798        };
12799        #[cfg(debug_assertions)]
12800        if let Some(ff) = ff {
12801            crate::debug_assert_tensor_stream_device(
12802                ff,
12803                &e.stream(),
12804                "step35_decode_attn.rope_freqs",
12805            );
12806        }
12807        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
12808
12809        if std::env::var("MEMRA_NOFA").is_ok() {
12810            return Err(
12811                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
12812                        cache; unset MEMRA_NOFA to use fa_decode"
12813                    .into(),
12814            );
12815        }
12816        let kvl = cache.kv[il].as_mut().unwrap();
12817        let next_len = kvl.len + 1;
12818        let (off, t_kv) = if swa && next_len > win {
12819            (next_len - win, win)
12820        } else {
12821            (0, next_len)
12822        };
12823        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
12824        e.append_kv_quantized(
12825            &k,
12826            &v0,
12827            &mut kvl.k,
12828            &mut kvl.v,
12829            write_row,
12830            kvl.kv_dim_k,
12831            kvl.kv_dim_v,
12832            kvl.k_tok_bytes,
12833            kvl.v_tok_bytes,
12834            crate::Engine::kv_fp8_on(),
12835        )?;
12836        kvl.len = next_len;
12837        let physical = kvl.physical_rows(off, off + t_kv)?;
12838        let k_view = e.view_u8_range(
12839            &kvl.k,
12840            physical.start * kvl.k_tok_bytes,
12841            physical.end * kvl.k_tok_bytes,
12842        );
12843        let v_view = e.view_u8_range(
12844            &kvl.v,
12845            physical.start * kvl.v_tok_bytes,
12846            physical.end * kvl.v_tok_bytes,
12847        );
12848        let mut attn = e.uninit(nh * hd)?;
12849        e.fa_decode_kvmod(
12850            &q,
12851            &k_view,
12852            &v_view,
12853            &mut attn,
12854            hd,
12855            nh,
12856            nkv,
12857            t_kv,
12858            scale,
12859            kvl.k_tok_bytes,
12860            kvl.v_tok_bytes,
12861            crate::Engine::kv_fp8_on(),
12862        )?;
12863
12864        let mut ag = e.uninit(nh * hd)?;
12865        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
12866        Ok(e.matmul(&fa.wo, &ag, 1)?)
12867    }
12868}
12869
12870// ===================================================================================== //
12871//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
12872//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
12873//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
12874//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
12875//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
12876//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
12877// ===================================================================================== //
12878impl HybridModel {
12879    pub fn is_gemma4_e4b(&self) -> bool {
12880        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
12881    }
12882
12883    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
12884    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
12885    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
12886    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
12887        let g = self.cfg.gemma4.as_ref().unwrap();
12888        let swa = g.swa_pattern[il];
12889        let hd = if swa {
12890            g.key_length_swa
12891        } else {
12892            g.key_length_global
12893        } as usize;
12894        let Mixer::Full(fa) = &self.layers[il].mixer else {
12895            panic!("e4b layer {il} not full-attn")
12896        };
12897        let nh = fa.wq.out_features() / hd;
12898        let nkv = fa.wk.out_features() / hd;
12899        (
12900            hd,
12901            nkv,
12902            nh,
12903            if swa {
12904                g.rope_base_swa
12905            } else {
12906                g.rope_base_global
12907            },
12908            1.0,
12909            swa,
12910        )
12911    }
12912
12913    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
12914    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
12915        self.layers[il]
12916            .gemma4
12917            .as_ref()
12918            .and_then(|b| b.e4b.as_ref())
12919            .and_then(|e4| e4.kv_share.map(|t| t as usize))
12920    }
12921
12922    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
12923    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
12924    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
12925    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
12926    fn gemma4_e4b_inp_pl(
12927        &self,
12928        e: &Engine,
12929        tokens: &[u32],
12930        x_scaled: &CudaSlice<f32>,
12931        t: usize,
12932    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12933        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
12934        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
12935    }
12936
12937    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
12938    fn gemma4_e4b_inp_pl_dev(
12939        &self,
12940        e: &Engine,
12941        tok_d: &CudaSlice<u32>,
12942        x_scaled: &CudaSlice<f32>,
12943        t: usize,
12944    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12945        let aux = self.gemma4_aux.as_ref().unwrap();
12946        let m = aux.e4b.as_ref().unwrap();
12947        let n_embd = self.cfg.n_embd as usize;
12948        let n_layer = self.layers.len();
12949        let width = m.n_epl * n_layer;
12950        let tbl = m.tok_tbl_gpu.get_or_init(|| {
12951            e.upload_u8(&m.tok_embd_bytes)
12952                .expect("e4b per-layer token table upload")
12953        });
12954        let mut a =
12955            e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt, m.tok_embd_row_bytes)?;
12956        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
12957        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
12958        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
12959        let mut pn = e.uninit(t * width)?;
12960        e.rms_norm(
12961            &p,
12962            m.proj_norm.float_data(),
12963            &mut pn,
12964            m.n_epl,
12965            t * n_layer,
12966            self.cfg.rms_eps,
12967        )?;
12968        let mut out = e.uninit(t * width)?;
12969        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
12970        Ok(out)
12971    }
12972
12973    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
12974    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
12975    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
12976    /// already holds this forward's rows — the target runs earlier in the stack).
12977    #[allow(clippy::too_many_arguments)]
12978    fn gemma4_e4b_attn(
12979        &self,
12980        e: &Engine,
12981        il: usize,
12982        hq: &CudaSlice<i8>,
12983        hdq: &CudaSlice<f32>,
12984        pos_d: &CudaSlice<i32>,
12985        t: usize,
12986        cache: &mut Cache,
12987        dc_bucket: Option<usize>,
12988    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12989        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
12990        let eps = self.cfg.rms_eps;
12991        let aux = self.gemma4_aux.as_ref().unwrap();
12992        let ones = aux.ones(e);
12993        #[cfg(debug_assertions)]
12994        crate::debug_assert_tensor_stream_device(ones, &e.stream(), "gemma4_e4b_attn.ones");
12995        let Mixer::Full(fa) = &self.layers[il].mixer else {
12996            unreachable!()
12997        };
12998        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
12999        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
13000        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
13001        let h0 = e.zeros(0)?;
13002        let h = &h0;
13003
13004        let ff = if swa {
13005            None
13006        } else {
13007            Some(
13008                aux.rope_freqs(e)
13009                    .expect("e4b global rope needs rope_freqs.weight"),
13010            )
13011        };
13012        #[cfg(debug_assertions)]
13013        if let Some(ff) = ff {
13014            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "gemma4_e4b_attn.rope_freqs");
13015        }
13016        let share = self.gemma4_e4b_kv_target(il);
13017        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
13018        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
13019        let mut q;
13020        if let Some(_tgt) = share {
13021            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
13022            q = e.uninit(t * nh * hd)?;
13023            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
13024            // empty; q0 stands in for the unused k/v pointers).
13025            let mut kdummy = e.uninit(1)?;
13026            let mut vdummy = e.uninit(1)?;
13027            e.rms_norm_qkv_rope(
13028                &q0,
13029                &q0,
13030                &q0,
13031                fa.q_norm.float_data(),
13032                fa.q_norm.float_data(),
13033                ones,
13034                &mut q,
13035                &mut kdummy,
13036                &mut vdummy,
13037                hd,
13038                nh * t,
13039                0,
13040                pos_d,
13041                nh,
13042                1,
13043                base,
13044                1.0,
13045                ff,
13046                eps,
13047            )?;
13048        } else {
13049            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
13050            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
13051            // q|k|v rows — the cat norm+rope twin consumes it directly.
13052            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
13053            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
13054            q = e.uninit(t * nh * hd)?;
13055            let mut k = e.uninit(t * nkv * hd)?;
13056            let mut v = e.uninit(t * nkv * hd)?;
13057            if t == 1 && cat.is_some() {
13058                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
13059                e.rms_norm_qkv_rope_cat(
13060                    &qkv0,
13061                    fa.q_norm.float_data(),
13062                    fa.k_norm.float_data(),
13063                    ones,
13064                    &mut q,
13065                    &mut k,
13066                    &mut v,
13067                    hd,
13068                    nh,
13069                    nkv,
13070                    pos_d,
13071                    nh,
13072                    nkv,
13073                    base,
13074                    1.0,
13075                    ff,
13076                    eps,
13077                )?;
13078            } else {
13079                let (q0, k0, v0) = match if t == 1 {
13080                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
13081                } else {
13082                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
13083                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
13084                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13085                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
13086                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
13087                    } else {
13088                        None
13089                    }
13090                } {
13091                    Some(triple) => triple,
13092                    None => (
13093                        e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
13094                        e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
13095                        e.matmul_pre(&fa.wv, hq, hdq, h, t)?,
13096                    ), // E4B: real v (K != V)
13097                };
13098                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
13099                // the normed rows; V ones-rms, never roped).
13100                e.rms_norm_qkv_rope(
13101                    &q0,
13102                    &k0,
13103                    &v0,
13104                    fa.q_norm.float_data(),
13105                    fa.k_norm.float_data(),
13106                    ones,
13107                    &mut q,
13108                    &mut k,
13109                    &mut v,
13110                    hd,
13111                    nh * t,
13112                    nkv * t,
13113                    pos_d,
13114                    nh,
13115                    nkv,
13116                    base,
13117                    1.0,
13118                    ff,
13119                    eps,
13120                )?;
13121            }
13122            let kvl = cache.kv[il].as_mut().unwrap();
13123            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
13124            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
13125            // degenerate tok-0 stream, 2026-07-12).
13126            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13127            if dc_bucket.is_some() {
13128                // DC arm (graph serving): append at the len_d slot, advance the counter
13129                // in-stream — replay-correct, no host len in the launch args. Host mirrors
13130                // are NOT touched here (the replay loop owns them; a bump at capture-record
13131                // time would double-count the capture iteration).
13132                debug_assert!(t == 1);
13133                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
13134                e.append_kv_quantized_row_dc_inc(
13135                    &k,
13136                    &v,
13137                    &mut kvl.k,
13138                    &mut kvl.v,
13139                    &mut kvl.len_d,
13140                    kvl.kv_dim_k,
13141                    kvl.kv_dim_v,
13142                    kvl.k_tok_bytes,
13143                    kvl.v_tok_bytes,
13144                    cls,
13145                )?;
13146            } else {
13147                e.append_kv_quantized_rows(
13148                    &k,
13149                    &v,
13150                    &mut kvl.k,
13151                    &mut kvl.v,
13152                    kvl.len,
13153                    t,
13154                    kvl.kv_dim_k,
13155                    kvl.kv_dim_v,
13156                    kvl.k_tok_bytes,
13157                    kvl.v_tok_bytes,
13158                    cls,
13159                )?;
13160                kvl.len += t;
13161            }
13162            kv_f32 = Some((k, v));
13163        }
13164        // attention: per-row causal fa over the (own or target) quantized cache. The cache
13165        // already contains this forward's rows in both arms; row i attends [.., base+i].
13166        let kvl_idx = share.unwrap_or(il);
13167        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
13168        let base_len = kvl.len - t; // pre-append length (target appended this forward too)
13169        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
13170        let mut attn = e.uninit(t * nh * hd)?;
13171        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
13172        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
13173        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
13174        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
13175        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
13176        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
13177        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
13178        //     rows (the T=K verify kernel; the target appended this forward's rows already).
13179        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
13180        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
13181        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
13182        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
13183            if let Some((kf, vf)) = &kv_f32 {
13184                if hd == 256 && t <= win {
13185                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13186                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13187                }
13188                if hd == 256 && swa && t > win {
13189                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13190                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13191                }
13192                if hd == 512 && !swa {
13193                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13194                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13195                }
13196            } else if share.is_some() {
13197                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13198                let k_view = e.view_u8(&kvl.k, kvl.k.len());
13199                let v_view = e.view_u8(&kvl.v, kvl.v.len());
13200                if hd == 256 && (!swa || t <= win) {
13201                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
13202                    e.fa_prefill_view(
13203                        &q,
13204                        &k_view,
13205                        &v_view,
13206                        &mut attn,
13207                        hd,
13208                        nh,
13209                        nkv,
13210                        t,
13211                        t,
13212                        scale,
13213                        true,
13214                        kvl.k_tok_bytes,
13215                        kvl.v_tok_bytes,
13216                        g,
13217                    )?;
13218                    return Ok(e.matmul(&fa.wo, &attn, t)?);
13219                }
13220                // remaining shared classes (swa above the window; hd512 globals): dequant
13221                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
13222                let kv_dim = nkv * hd;
13223                let mut kf = e.uninit(t * kv_dim)?;
13224                let mut vf = e.uninit(t * kv_dim)?;
13225                e.fa_dequant_kv_view_f32(
13226                    &k_view,
13227                    &v_view,
13228                    &mut kf,
13229                    &mut vf,
13230                    kv_dim,
13231                    kv_dim,
13232                    t,
13233                    kvl.k_tok_bytes,
13234                    kvl.v_tok_bytes,
13235                    g,
13236                )?;
13237                if hd == 512 {
13238                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
13239                } else {
13240                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
13241                }
13242                return Ok(e.matmul(&fa.wo, &attn, t)?);
13243            }
13244        }
13245        if let Some(bucket) = dc_bucket {
13246            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
13247            // fa_decode_dc over the live counter. len_d already advanced past this token
13248            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
13249            // counter (advanced when the target ran earlier in the stack).
13250            assert!(t == 1);
13251            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
13252            // and under the window every live t_kv sits below it — cap the capture bucket
13253            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
13254            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
13255            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
13256            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
13257                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
13258            } else {
13259                bucket
13260            };
13261            let k_view = e.view_u8(&kvl.k, kvl.k.len());
13262            let v_view = e.view_u8(&kvl.v, kvl.v.len());
13263            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
13264            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
13265            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
13266            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
13267            // captured into the dc graph like any other launch. Extending the cascade to
13268            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
13269            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
13270            // MEMRA_WPF=0 rollback seam.
13271            if crate::Engine::wpf_level() >= 1 {
13272                e.prefetch_weight_l2(&fa.wo)?;
13273            }
13274            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
13275            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
13276            if e.uses_q8_1_fast(&fa.wo) {
13277                let mut oq = e.alloc_i8_uninit(nh * hd)?;
13278                let mut od = e.zeros(nh * hd / 32)?;
13279                e.fa_decode_dc_q8(
13280                    &q,
13281                    &k_view,
13282                    &v_view,
13283                    &mut attn,
13284                    hd,
13285                    nh,
13286                    nkv,
13287                    &kvl.len_d,
13288                    bucket,
13289                    scale,
13290                    kvl.k_tok_bytes,
13291                    kvl.v_tok_bytes,
13292                    g,
13293                    Some((&mut oq, &mut od)),
13294                )?;
13295                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
13296            }
13297            e.fa_decode_dc(
13298                &q,
13299                &k_view,
13300                &v_view,
13301                &mut attn,
13302                hd,
13303                nh,
13304                nkv,
13305                &kvl.len_d,
13306                bucket,
13307                scale,
13308                kvl.k_tok_bytes,
13309                kvl.v_tok_bytes,
13310                g,
13311            )?;
13312            return Ok(e.matmul(&fa.wo, &attn, t)?);
13313        }
13314        for i in 0..t {
13315            let avail = base_len + i + 1;
13316            let (off_tok, t_kv) = if swa && avail > win {
13317                (avail - win, win)
13318            } else {
13319                (0, avail)
13320            };
13321            let k_view = e.view_u8_range(
13322                &kvl.k,
13323                off_tok * kvl.k_tok_bytes,
13324                (off_tok + t_kv) * kvl.k_tok_bytes,
13325            );
13326            let v_view = e.view_u8_range(
13327                &kvl.v,
13328                off_tok * kvl.v_tok_bytes,
13329                (off_tok + t_kv) * kvl.v_tok_bytes,
13330            );
13331            let qv = e.view(&q, t * nh * hd);
13332            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
13333            let mut q_one = e.uninit(nh * hd)?;
13334            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
13335            let mut a_one = e.uninit(nh * hd)?;
13336            // read class MUST match the append class (globals are e4m3 under gkv): the
13337            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
13338            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
13339            e.fa_decode_kvmod(
13340                &q_one,
13341                &k_view,
13342                &v_view,
13343                &mut a_one,
13344                hd,
13345                nh,
13346                nkv,
13347                t_kv,
13348                scale,
13349                kvl.k_tok_bytes,
13350                kvl.v_tok_bytes,
13351                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
13352            )?;
13353            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
13354        }
13355        Ok(e.matmul(&fa.wo, &attn, t)?)
13356    }
13357
13358    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
13359    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
13360    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
13361    /// layer; does NOT advance cache.pos (caller owns pos).
13362    fn gemma4_e4b_trunk(
13363        &self,
13364        e: &Engine,
13365        tokens: &[u32],
13366        pos0: usize,
13367        cache: &mut Cache,
13368        head_last: bool,
13369    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13370        let n_embd = self.cfg.n_embd as usize;
13371        let t = tokens.len();
13372        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13373        let pos_d = e.htod_i32(&pos)?;
13374        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
13375        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13376        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
13377        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
13378    }
13379
13380    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
13381    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
13382    /// eager chain by construction: SAME functions, not twins).
13383    fn gemma4_e4b_trunk_core(
13384        &self,
13385        e: &Engine,
13386        x_in: CudaSlice<f32>,
13387        inp_pl: CudaSlice<f32>,
13388        pos_d: &CudaSlice<i32>,
13389        t: usize,
13390        cache: &mut Cache,
13391        dc_bucket: Option<usize>,
13392        cap_logits: bool,
13393        head_last: bool,
13394    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13395        let n_embd = self.cfg.n_embd as usize;
13396        let eps = self.cfg.rms_eps;
13397        let n_layer = self.layers.len();
13398        let mut x = x_in;
13399        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
13400        let n_epl = aux_e4b.n_epl;
13401
13402        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
13403        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
13404        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
13405        // head rides matmul_pre too. First layer's pair comes from a standalone fused
13406        // norm+quant.
13407        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
13408        for il in 0..n_layer {
13409            let layer = &self.layers[il];
13410            let (hq, hdq) = match h_carry.take() {
13411                Some(p) => p,
13412                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
13413            };
13414            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
13415            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
13416            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
13417            let bits = layer.gemma4.as_ref().unwrap();
13418            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
13419            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
13420            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
13421            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
13422            // the fused single-phase reduction is NOT FP-order-identical to the unfused
13423            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
13424            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
13425            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
13426            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
13427            // gate dropped, decode AND verify ride the same fused chain — parity by
13428            // construction, VERIFY-GATE 0.000e0.
13429            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
13430            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
13431                e,
13432                layer,
13433                &o,
13434                &x,
13435                t,
13436                Some(layer.post_attn_norm.float_data()),
13437                fuse_exit,
13438            )?;
13439            let mut resid = e.uninit(t * n_embd)?;
13440            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
13441            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
13442            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
13443            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
13444            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
13445            let g = if fuse_exit {
13446                // sn here = RAW f0 (post_ffw deferred).
13447                let (rq, rd) = e.rms_pre_add_q8_1(
13448                    &sn,
13449                    bits.post_ffw_norm.float_data(),
13450                    &attn_out,
13451                    &mut resid,
13452                    n_embd,
13453                    t,
13454                    self.cfg.rms_eps,
13455                )?;
13456                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
13457            } else {
13458                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
13459                e.matmul(&e4b.inp_gate, &resid, t)?
13460            };
13461            let mut act = e.uninit(t * n_epl)?;
13462            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
13463                let ipv = e.view(&inp_pl, n_epl * n_layer);
13464                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
13465                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
13466                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
13467            } else {
13468                let mut inp_this = e.uninit(t * n_epl)?;
13469                e.copy_rows_strided(
13470                    &inp_pl,
13471                    &mut inp_this,
13472                    n_epl,
13473                    t,
13474                    n_epl * n_layer,
13475                    il * n_epl,
13476                )?;
13477                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
13478                e.matmul(&e4b.proj, &act, t)?
13479            };
13480            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
13481            // ONE launch (glue-fusion lane; last layer emits through output_norm).
13482            let next_norm = if il + 1 < n_layer {
13483                self.layers[il + 1].attn_norm.float_data()
13484            } else {
13485                self.output_norm.float_data()
13486            };
13487            let mut xn = e.uninit(t * n_embd)?;
13488            let pair = e.rms_pre_add_scale_rms_norm_q8_1(
13489                &y,
13490                e4b.post_norm.float_data(),
13491                &resid,
13492                bits.layer_scale,
13493                next_norm,
13494                &mut xn,
13495                n_embd,
13496                t,
13497                eps,
13498            )?;
13499            h_carry = Some(pair);
13500            x = xn;
13501        }
13502        // the head consumes the last layer's fused (output_norm) emit. head_last callers
13503        // (prime, last_only forward) need only the final row's logits — the all-T head is
13504        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
13505        let (oq, odq) = h_carry.take().unwrap();
13506        let h0 = e.zeros(0)?;
13507        let hm = if head_last { 1 } else { t };
13508        let (hq, hd) = if head_last && t > 1 {
13509            let mut q1 = e.uninit_i8(n_embd)?;
13510            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
13511            let nb = n_embd / 32;
13512            let mut d1 = e.uninit(nb)?;
13513            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
13514            (q1, d1)
13515        } else {
13516            (oq, odq)
13517        };
13518        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
13519        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
13520        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
13521        // Logit-returning callers (host logits / spec prime) keep the capped emit.
13522        if cap_logits {
13523            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
13524            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
13525        }
13526        self.gemma4_suppress(e, &mut ld, hm)?; // mask both capped and argmax-only consumers
13527        Ok((ld, x))
13528    }
13529
13530    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
13531    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
13532    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
13533    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
13534    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
13535    /// covers exactly the layers that appended).
13536    pub fn gemma4_e4b_decode_step_t_am_dev(
13537        &self,
13538        e: &Engine,
13539        tok_d: &CudaSlice<u32>,
13540        t: usize,
13541        pos0: usize,
13542        cache: &mut Cache,
13543    ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13544        let n_embd = self.cfg.n_embd as usize;
13545        let eps = self.cfg.rms_eps;
13546        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
13547        let pos_d = e.htod_i32(&pos)?;
13548        let embd_gpu = self
13549            .embd_gpu
13550            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
13551        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
13552        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
13553        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
13554        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
13555        let (ld, xp) =
13556            self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, false)?;
13557        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
13558        // emit is already capped, matching the eager chain bit-for-bit).
13559        let n_vocab = self.output.out_features();
13560        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
13561        for i in 0..t {
13562            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
13563        }
13564        let mut hn = e.uninit(t * n_embd)?;
13565        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13566        cache.pos += t;
13567        Ok((vam, hn))
13568    }
13569
13570    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
13571    /// prime path — mirror of `gemma4_decode_step_t_h`).
13572    pub(crate) fn gemma4_e4b_decode_step_t_h(
13573        &self,
13574        e: &Engine,
13575        tokens: &[u32],
13576        pos0: usize,
13577        cache: &mut Cache,
13578    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13579        let n_embd = self.cfg.n_embd as usize;
13580        let eps = self.cfg.rms_eps;
13581        let t = tokens.len();
13582        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
13583        let mut hn = e.uninit(t * n_embd)?;
13584        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
13585        cache.pos += t;
13586        Ok((e.dtoh(&ld)?, hn))
13587    }
13588
13589    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
13590    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
13591    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
13592    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
13593    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
13594    pub fn gemma4_e4b_decode_step_dcg(
13595        &self,
13596        e: &Engine,
13597        token_d: &mut CudaSlice<u32>,
13598        pos_d: &mut CudaSlice<i32>,
13599        embd_gpu: &CudaSlice<u8>,
13600        embd_qt: i32,
13601        embd_rb: usize,
13602        cache: &mut Cache,
13603        n_vocab: usize,
13604        bucket: usize,
13605    ) -> Result<(), Box<dyn std::error::Error>> {
13606        let n_embd = self.cfg.n_embd as usize;
13607        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13608        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13609        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13610        let (ld, _x) =
13611            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket), false, false)?;
13612        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
13613        e.inc_seqlen(pos_d)?;
13614        Ok(())
13615    }
13616
13617    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
13618    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
13619    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
13620    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
13621    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
13622    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
13623    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
13624    #[allow(clippy::too_many_arguments)]
13625    pub fn gemma4_e4b_decode_step_dc(
13626        &self,
13627        e: &Engine,
13628        token_d: &CudaSlice<u32>,
13629        pos_d: &mut CudaSlice<i32>,
13630        embd_gpu: &CudaSlice<u8>,
13631        embd_qt: i32,
13632        embd_rb: usize,
13633        cache: &mut Cache,
13634        n_vocab: usize,
13635    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
13636        let n_embd = self.cfg.n_embd as usize;
13637        let eps = self.cfg.rms_eps;
13638        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
13639        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
13640        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
13641        let (ld, _x) =
13642            self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false, false)?;
13643        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
13644        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
13645        e.inc_seqlen(pos_d)?;
13646        cache.pos += 1;
13647        let _ = eps;
13648        Ok(tok_out)
13649    }
13650
13651    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
13652    /// pre-output_norm hidden). Advances cache.pos.
13653    pub(crate) fn gemma4_e4b_decode_step_h(
13654        &self,
13655        e: &Engine,
13656        token: u32,
13657        cache: &mut Cache,
13658    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13659        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
13660        let logits = e.dtoh(&ld)?;
13661        cache.pos += 1;
13662        Ok((logits, x))
13663    }
13664
13665    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
13666    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
13667    /// fast; the prefill fa arms come later.
13668    pub(crate) fn gemma4_e4b_prime(
13669        &self,
13670        e: &Engine,
13671        tokens: &[u32],
13672        cache: &mut Cache,
13673    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13674        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
13675        // process-kill as gemma4_prime — refuse per-request.
13676        if cache.pos != 0 {
13677            return Err(
13678                "e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
13679                        call or decode tokenwise"
13680                    .into(),
13681            );
13682        }
13683        let n_embd = self.cfg.n_embd as usize;
13684        let t = tokens.len();
13685        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
13686        cache.pos += t;
13687        let last = e.dtoh(&ld)?; // head_last: ld is already the final row only
13688        let xv = e.view(&x, t * n_embd);
13689        let row = xv.slice((t - 1) * n_embd..t * n_embd);
13690        let mut h_seed = e.uninit(n_embd)?;
13691        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
13692        Ok((last, h_seed, x))
13693    }
13694
13695    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
13696    pub(crate) fn gemma4_e4b_forward(
13697        &self,
13698        e: &Engine,
13699        tokens: &[u32],
13700        last_only: bool,
13701    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
13702        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
13703        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
13704        Ok(e.dtoh(&ld)?) // head_last already reduced to the final row when last_only
13705    }
13706}
13707
13708#[cfg(test)]
13709mod prime_chunk_schedule_tests {
13710    use super::{
13711        PRIME_MIN_T, PRIME_PIPE_MIN_CHUNK, dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges,
13712        fixed_prime_chunk_ranges_for_ring,
13713    };
13714
13715    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
13716        ranges.iter().map(|(start, end)| end - start).collect()
13717    }
13718
13719    fn auto_chunk(t: usize) -> usize {
13720        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
13721    }
13722
13723    #[test]
13724    fn fixed_schedule_retains_measured_geometry() {
13725        assert_eq!(
13726            sizes(&fixed_prime_chunk_ranges(461, 128)),
13727            vec![128, 128, 128, 77]
13728        );
13729        assert_eq!(
13730            sizes(&fixed_prime_chunk_ranges(1833, 230)),
13731            vec![230, 230, 230, 230, 230, 230, 230, 223]
13732        );
13733        assert_eq!(sizes(&fixed_prime_chunk_ranges(4096, 512)), vec![512; 8]);
13734        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
13735        assert_eq!(capped, vec![4096, 4088, 16]);
13736        assert!(capped.iter().all(|&rows| rows <= 4096));
13737        assert_eq!(
13738            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
13739            vec![4100],
13740            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
13741        );
13742    }
13743
13744    #[test]
13745    fn dynamic_schedule_matches_registered_shapes() {
13746        let cases = [
13747            (461, vec![64, 141, 132, 124]),
13748            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
13749            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
13750        ];
13751        for (t, expected) in cases {
13752            let chunk = auto_chunk(t);
13753            let fixed = fixed_prime_chunk_ranges(t, chunk);
13754            assert_eq!(
13755                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
13756                expected
13757            );
13758        }
13759    }
13760
13761    #[test]
13762    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
13763        for t in 256..=8192 {
13764            let chunk = auto_chunk(t);
13765            let fixed = fixed_prime_chunk_ranges(t, chunk);
13766            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
13767            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
13768            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
13769            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
13770            for pair in dynamic.windows(2) {
13771                assert_eq!(pair[0].1, pair[1].0, "T={t}");
13772            }
13773            assert!(
13774                dynamic
13775                    .iter()
13776                    .all(|(start, end)| end - start >= PRIME_MIN_T),
13777                "T={t} sizes={:?}",
13778                sizes(&dynamic)
13779            );
13780            if dynamic.len() >= 3 {
13781                let chunk_sizes = sizes(&dynamic);
13782                assert!(
13783                    chunk_sizes[0] < chunk_sizes[1],
13784                    "T={t} sizes={chunk_sizes:?}"
13785                );
13786                assert!(
13787                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
13788                    "T={t} sizes={chunk_sizes:?}"
13789                );
13790            }
13791        }
13792    }
13793}
13794
13795#[cfg(test)]
13796mod page_prefetch_tests {
13797    use super::{
13798        grouped_worker_prefetch_position, page_prefetch_positions,
13799        page_prefetch_window_from_values, worker_prefetch_positions,
13800    };
13801
13802    #[test]
13803    fn page_prefetch_window_keeps_existing_opt_in_default() {
13804        assert_eq!(page_prefetch_window_from_values(false, None), 0);
13805        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
13806        assert_eq!(page_prefetch_window_from_values(true, None), 1);
13807        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
13808        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
13809        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
13810    }
13811
13812    #[test]
13813    fn rolling_page_prefetch_advises_each_future_expert_once() {
13814        let advised: Vec<_> = (0..7)
13815            .flat_map(|position| page_prefetch_positions(position, 7, 3))
13816            .collect();
13817        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
13818
13819        let one_ahead: Vec<_> = (0..4)
13820            .flat_map(|position| page_prefetch_positions(position, 4, 1))
13821            .collect();
13822        assert_eq!(one_ahead, vec![1, 2, 3]);
13823        assert!(page_prefetch_positions(0, 4, 0).is_empty());
13824    }
13825
13826    #[test]
13827    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
13828        assert_eq!(grouped_worker_prefetch_position(0, None), None);
13829        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
13830            .chain(
13831                (0..4).filter_map(|position| grouped_worker_prefetch_position(4, Some(position))),
13832            )
13833            .collect();
13834        assert_eq!(positions, vec![0, 1, 2, 3]);
13835        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
13836    }
13837
13838    #[test]
13839    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
13840        let queued: Vec<_> = (0..8)
13841            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
13842            .collect();
13843        assert_eq!(queued, (0..8).collect::<Vec<_>>());
13844
13845        let one_at_a_time: Vec<_> = (0..4)
13846            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
13847            .collect();
13848        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
13849        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
13850    }
13851}
13852
13853pub struct G4DcSlots {
13854    x: CudaSlice<f32>,
13855    xn: CudaSlice<f32>,
13856    cur: CudaSlice<f32>,
13857    hq: CudaSlice<i8>,
13858    hd_: CudaSlice<f32>,
13859    q0: CudaSlice<f32>,
13860    k0: CudaSlice<f32>,
13861    v0: CudaSlice<f32>,
13862    q: CudaSlice<f32>,
13863    k: CudaSlice<f32>,
13864    v: CudaSlice<f32>,
13865    attn: CudaSlice<f32>,
13866    o: CudaSlice<f32>,
13867    attn_out: CudaSlice<f32>,
13868    zsh: CudaSlice<f32>,
13869    zq: CudaSlice<i8>,
13870    zd: CudaSlice<f32>,
13871    gate: CudaSlice<f32>,
13872    up: CudaSlice<f32>,
13873    act: CudaSlice<f32>,
13874    actq: CudaSlice<i8>,
13875    actd: CudaSlice<f32>,
13876    f0: CudaSlice<f32>,
13877    sn: CudaSlice<f32>,
13878    hn: CudaSlice<f32>,
13879    logits: CudaSlice<f32>,
13880}