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 cudarc::driver::CudaSlice;
6use memra_gguf::config::ModelConfig;
7use crate::Engine;
8use crate::cache::Cache;
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
124/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
125pub(crate) struct AttnPre {
126    pub q: cudarc::driver::CudaSlice<f32>,
127    pub k: cudarc::driver::CudaSlice<f32>,
128    pub v: cudarc::driver::CudaSlice<f32>,
129    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
130}
131
132/// task #18: one sequence's GDN prep outputs (the scan inputs).
133pub(crate) struct GdnPrep {
134    pub hk: usize,
135    pub q_l2: cudarc::driver::CudaSlice<f32>,
136    pub k_l2: cudarc::driver::CudaSlice<f32>,
137    pub v_g: cudarc::driver::CudaSlice<f32>,
138    pub beta: cudarc::driver::CudaSlice<f32>,
139    pub g_log: cudarc::driver::CudaSlice<f32>,
140    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
141    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
142}
143
144/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
145pub(crate) struct VerifyStreamScratch {
146    pub pos_d: CudaSlice<i32>,
147    pub row_ctrs: Vec<CudaSlice<i32>>,
148}
149use crate::hybrid::{HybridModel, Mixer, FullAttnLayer, LinearAttnLayer, MoeWeights};
150
151struct MoeInputTraceWriter {
152    dir: std::path::PathBuf,
153    index: std::fs::File,
154    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
155}
156
157static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<
158    std::sync::Mutex<Option<MoeInputTraceWriter>>,
159> = std::sync::OnceLock::new();
160
161/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
162/// per-expert launch chain). See `moe_gdec_token`.
163fn gdec_enabled() -> bool {
164    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
165    *E.get_or_init(|| std::env::var("MEMRA_MOE_GDEC").map(|v| v != "0").unwrap_or(true))
166}
167
168/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
169/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
170/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
171/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
172/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
173/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
174/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
175/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
176/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
177/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
178fn moe_slab_enabled() -> bool {
179    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
180}
181
182/// Expert-grouped dispatch remains opt-in after the local 5090 transfer gate rejected the
183/// default flip. `=0` selects the established path, while any other explicit value enables the
184/// grouped research arm for the current call.
185fn moe_grouped_enabled(_cfg: &ModelConfig, _prefill: bool) -> bool {
186    std::env::var("MEMRA_MOE_GROUPED")
187        .map(|value| value != "0")
188        .unwrap_or(false)
189}
190
191/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
192/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
193fn moe_prefetch_enabled() -> bool {
194    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
196        || crate::spill_pread::worker_enabled())
197}
198
199/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
200/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
201/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
202/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
203fn moe_page_prefetch_window() -> usize {
204    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
205    *W.get_or_init(|| page_prefetch_window_from_values(
206        std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
207        std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
208    ))
209}
210
211fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
212    if !enabled {
213        return 0;
214    }
215    raw_window
216        .and_then(|value| value.parse().ok())
217        .unwrap_or(1)
218}
219
220/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
221/// full window; each later position adds one expert at the far edge. Thus widening the window does
222/// not repeatedly issue `MADV_WILLNEED` for the same range.
223fn page_prefetch_positions(
224    position: usize,
225    len: usize,
226    window: usize,
227) -> std::ops::Range<usize> {
228    if window == 0 || position >= len {
229        return len..len;
230    }
231    let (start, count) = if position == 0 {
232        (1, window)
233    } else {
234        (position.saturating_add(window), 1)
235    };
236    let start = start.min(len);
237    start..start.saturating_add(count).min(len)
238}
239
240/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
241/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
242fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
243    let position = current.map_or(0, |position| position.saturating_add(1));
244    (position < order_len).then_some(position)
245}
246
247/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
248/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
249/// window. Position zero primes the current expert too: its three independent reads can run in
250/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
251fn worker_prefetch_window() -> usize {
252    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
253    *WINDOW.get_or_init(|| {
254        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
255        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
256            .ok()
257            .and_then(|value| value.parse::<usize>().ok())
258            .unwrap_or(automatic.max(1))
259    })
260}
261
262/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
263/// this includes the current expert when the window is seeded so all three current projections
264/// enter the CPU pool together.
265fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
266    if window == 0 || position >= len {
267        return len..len;
268    }
269    let (start, count) = if position == 0 {
270        (0, window)
271    } else {
272        (position.saturating_add(window).saturating_sub(1), 1)
273    };
274    let start = start.min(len);
275    start..start.saturating_add(count).min(len)
276}
277
278/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
279/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
280/// expert weight pointers come from the per-layer device table. Requires the fused router (the
281/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
282fn moe_dev_enabled() -> bool {
283    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
284    *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
285        && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
286}
287
288/// Device sigmoid top-k is the default for Step-3.7 / M3 / Hy3 / GLM-DSA. `MEMRA_SIG_ROUTER=0` restores
289/// the full-logit DtoH plus `moe_route_sigmoid_host` oracle without changing expert dispatch.
290fn sigmoid_router_enabled() -> bool {
291    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
292    *E.get_or_init(|| std::env::var("MEMRA_SIG_ROUTER").map(|v| v != "0").unwrap_or(true))
293}
294
295/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
296/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
297/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
298/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
299fn moe_q8_enabled() -> bool {
300    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
301    *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
302}
303
304/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
305/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
306fn expert_dp4a_supported(qt: i32) -> bool {
307    qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
308        || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
309}
310
311fn q8_expert_supported(qt: i32) -> bool {
312    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
313    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
314    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
315    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
316    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
317    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
318    let kq = *KQ.get_or_init(|| {
319        std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
320    });
321    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
322    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
323    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
324    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
325    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
326    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
327    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
328    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
329        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
330}
331
332/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
333/// k-quant tensors must fall to the _em dot path instead.
334fn q8_expert_dec_supported(qt: i32) -> bool {
335    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
336}
337
338/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
339/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
340/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
341/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
342/// q35 layers, which is why that cell measured FLAT.
343fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
344    match qt {
345        crate::QT_Q4_0 => in_f % 32 == 0,
346        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
347        | crate::QT_Q6_K => in_f % 256 == 0,
348        _ => false,
349    }
350}
351
352/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
353/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
354fn moe_prewarm_enabled() -> bool {
355    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
356    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
357}
358
359/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
360/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
361/// can then vote for and exercise those experts on GPU before the cache is frozen.
362fn cpu_expert_profile_admit_enabled() -> bool {
363    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
364    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
365}
366
367/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
368/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
369/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
370pub const PRIME_MIN_T: usize = 16;
371const PRIME_PIPE_MICROBATCHES: usize = 8;
372const PRIME_PIPE_MIN_CHUNK: usize = 128;
373const PRIME_PIPE_EDGE_MIN_CHUNK: usize = 64;
374const PRIME_PIPE_LINEAR_WORK: usize = 8;
375
376fn prime_pp2_auto_geometry(n_layers: usize) -> bool {
377    crate::pp::prime_pp_on()
378        && !crate::pp::pp2_streams_off()
379        && crate::pp::pp_cuts(n_layers).is_some_and(|cuts| cuts.len() == 3)
380}
381
382/// Effective internal prime chunk. An explicit MEMRA_PRIME_CHUNK is authoritative.
383/// Naked PP-2 primes use the measured pipeline geometry: up to eight microchunks, never
384/// below 128 tokens, while the legacy 4096-token cap remains the long-context bound.
385pub fn prime_chunk_tokens(t: usize, n_layers: usize) -> usize {
386    if let Ok(value) = std::env::var("MEMRA_PRIME_CHUNK") {
387        let parsed = value.parse::<usize>().unwrap_or(crate::cache::PRIME_CHUNK_MAX_TOKENS);
388        return if crate::cache::swa_ring_on() {
389            if parsed == 0 {
390                crate::cache::PRIME_CHUNK_MAX_TOKENS
391            } else {
392                parsed.min(crate::cache::PRIME_CHUNK_MAX_TOKENS)
393            }
394        } else {
395            parsed
396        };
397    }
398    let chunk = crate::cache::PRIME_CHUNK_MAX_TOKENS;
399    if prime_pp2_auto_geometry(n_layers) && t >= 2 * PRIME_PIPE_MIN_CHUNK {
400        chunk.min(
401            t.div_ceil(PRIME_PIPE_MICROBATCHES)
402                .max(PRIME_PIPE_MIN_CHUNK),
403        )
404    } else {
405        chunk
406    }
407}
408
409fn fixed_prime_chunk_ranges(t: usize, chunk: usize) -> Vec<(usize, usize)> {
410    fixed_prime_chunk_ranges_for_ring(t, chunk, crate::cache::swa_ring_on())
411}
412
413fn fixed_prime_chunk_ranges_for_ring(t: usize, chunk: usize, ring_on: bool) -> Vec<(usize, usize)> {
414    if chunk == 0 || t <= chunk {
415        return vec![(0, t)];
416    }
417    let mut ranges = Vec::with_capacity(t.div_ceil(chunk));
418    let mut start = 0usize;
419    while start < t {
420        let mut end = (start + chunk).min(t);
421        if t - end > 0 && t - end < PRIME_MIN_T {
422            if ring_on {
423                let shifted = t - PRIME_MIN_T;
424                end = if shifted > start { shifted } else { t };
425            } else {
426                end = t;
427            }
428        }
429        ranges.push((start, end));
430        start = end;
431    }
432    ranges
433}
434
435fn prime_chunk_work(prefix: usize, total: usize) -> u128 {
436    let prefix = prefix as u128;
437    prefix * (prefix + (PRIME_PIPE_LINEAR_WORK as u128) * (total as u128))
438}
439
440fn dynamic_prime_chunk_ranges(
441    t: usize,
442    fixed_chunk: usize,
443    fixed: &[(usize, usize)],
444) -> Vec<(usize, usize)> {
445    let n = fixed.len();
446    if n < 3 {
447        return fixed.to_vec();
448    }
449
450    let max_first = t - (n - 1) * PRIME_MIN_T;
451    let first = fixed_chunk
452        .div_ceil(2)
453        .max(PRIME_PIPE_EDGE_MIN_CHUNK)
454        .min(max_first);
455    let mut ranges = Vec::with_capacity(n);
456    ranges.push((0, first));
457
458    let first_work = prime_chunk_work(first, t);
459    let work_span = prime_chunk_work(t, t) - first_work;
460    let denominator = (n - 1) as u128;
461    let mut previous = first;
462    for boundary in 1..n - 1 {
463        let target = first_work * denominator + work_span * (boundary as u128);
464        let remaining = n - 1 - boundary;
465        let mut low = previous + PRIME_MIN_T;
466        let mut high = t - remaining * PRIME_MIN_T;
467        while low < high {
468            let mid = low + (high - low) / 2;
469            if prime_chunk_work(mid, t) * denominator >= target {
470                high = mid;
471            } else {
472                low = mid + 1;
473            }
474        }
475        ranges.push((previous, low));
476        previous = low;
477    }
478    ranges.push((previous, t));
479    ranges
480}
481
482/// Internal prime ranges. The naked PP-2 pipeline defaults to a short-fill,
483/// equal-modeled-time schedule; MEMRA_PRIME_CHUNK_SCHED=fixed restores the measured
484/// equal-token ranges. An explicit MEMRA_PRIME_CHUNK always retains fixed semantics.
485pub fn prime_chunk_ranges(t: usize, n_layers: usize) -> Vec<(usize, usize)> {
486    let explicit_chunk = std::env::var_os("MEMRA_PRIME_CHUNK").is_some();
487    let chunk = prime_chunk_tokens(t, n_layers);
488    let fixed = fixed_prime_chunk_ranges(t, chunk);
489    let dynamic = match std::env::var("MEMRA_PRIME_CHUNK_SCHED") {
490        Ok(value) => value == "dynamic",
491        Err(_) => true,
492    };
493    if explicit_chunk || !dynamic || !prime_pp2_auto_geometry(n_layers) {
494        fixed
495    } else {
496        dynamic_prime_chunk_ranges(t, chunk, &fixed)
497    }
498}
499
500impl HybridModel {
501    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
502    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
503    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
504    /// (it forces a dtoh + host hash per layer).
505    fn prime_trace_path() -> Option<&'static str> {
506        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
507        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
508            .as_deref()
509    }
510
511    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
512    pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
513        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
514        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
515        let cfg = &self.cfg;
516        let n_embd = cfg.n_embd as usize;
517        let t = tokens.len();
518        let eps = cfg.rms_eps;
519        let pos: Vec<i32> = (0..t as i32).collect();
520        let pos_d = e.htod_i32(&pos)?;
521
522        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
523
524        for (il, layer) in self.layers.iter().enumerate() {
525            // attn_norm
526            let mut h = e.uninit(t * n_embd)?;
527            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
528
529            let mixed = match &layer.mixer {
530                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
531                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
532                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
533            };
534
535            // residual 1
536            let mut x1 = e.uninit(t * n_embd)?;
537            e.add(&x, &mixed, &mut x1, t * n_embd)?;
538
539            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
540            let mut z = e.uninit(t * n_embd)?;
541            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
542            let ffn_out = match &layer.ffn {
543                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
544                    let n_ff = ffn_gate.out_features();
545                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
546                    let up = g2.pop().unwrap();
547                    let gate = g2.pop().unwrap();
548                    let mut act = e.uninit(t * n_ff)?;
549                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
550                    // both the dense MLP and the shared expert, and its limit is
551                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
552                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
553                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
554                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
555                    e.matmul(ffn_down, &act, t)?
556                }
557                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
558            };
559            let mut x2 = e.uninit(t * n_embd)?;
560            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
561            x = x2;
562        }
563
564        let mut hn = e.uninit(t * n_embd)?;
565        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
566        let logits = e.matmul(&self.output, &hn, t)?;
567        Ok(e.dtoh(&logits)?)
568    }
569
570    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
571    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
572    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
573    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
574    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
575    pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
576        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
577        let cfg = &self.cfg;
578        let n_embd = cfg.n_embd as usize;
579        let t = tokens.len();
580        let eps = cfg.rms_eps;
581        let pos: Vec<i32> = (0..t as i32).collect();
582        let pos_d = e.htod_i32(&pos)?;
583
584        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
585        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
586        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
587        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
588        for (il, layer) in self.layers.iter().enumerate() {
589            let mut h = e.uninit(t * n_embd)?;
590            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
591            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
592            let mixed = match &layer.mixer {
593                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
594                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
595                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
596            };
597            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
598            let mut x1 = e.uninit(t * n_embd)?;
599            e.add(&x, &mixed, &mut x1, t * n_embd)?;
600            let mut z = e.uninit(t * n_embd)?;
601            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
602            let ffn_out = match &layer.ffn {
603                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
604                    let n_ff = ffn_gate.out_features();
605                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
606                    let up = g2.pop().unwrap();
607                    let gate = g2.pop().unwrap();
608                    let mut act = e.uninit(t * n_ff)?;
609                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
610                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
611                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
612                    e.matmul(ffn_down, &act, t)?
613                }
614                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
615            };
616            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
617            let mut x2 = e.uninit(t * n_embd)?;
618            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
619            x = x2;
620        }
621        // norm over all T, then slice the LAST row and run lm_head on that single row.
622        let mut hn = e.uninit(t * n_embd)?;
623        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
624        let last = e.view(&hn, t * n_embd);            // [T, n_embd]
625        let last_row = last.slice((t - 1) * n_embd..t * n_embd);  // [1, n_embd]
626        let mut hlast = e.uninit(n_embd)?;
627        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
628        let logits = e.matmul(&self.output, &hlast, 1)?;   // [1, n_vocab] — lm_head on ONE row
629        Ok(e.dtoh(&logits)?)
630    }
631
632    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
633    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
634    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
635    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
636    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
637    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
638    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
639    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
640    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
641    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
642    ///       argmax gate is the accuracy authority, exactly as for forward_last);
643    ///   (c) `cache.pos`/KV len/len_d advance by T.
644    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
645    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
646    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
647    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
648    ///
649    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
650    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
651    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
652    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
653    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
654    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
655    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
656    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
657    /// differently under load — research/tick-seg-20260807, receipt in
658    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
659    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
660    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
661    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
662    /// caller that SPLITS one request across calls passes the remainder.
663    pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize)
664                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
665        let n_embd = self.cfg.n_embd as usize;
666        let t = tokens.len();
667        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
668        // session cache — every chunk (including the first) takes the continuation arm
669        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
670        assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
671        assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
672
673        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
674        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
675        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
676        // each chunk runs the full layer stack with transients sized to the chunk, appending its
677        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
678        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
679        // exactly the state carry it was built for). Full-attn chunks after the first attend to
680        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
681        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
682        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
683        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
684        if self.is_gemma4_e4b() {
685            return self.gemma4_e4b_prime(e, tokens, cache);
686        }
687        if self.cfg.gemma4.is_some() {
688            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
689            return self.gemma4_prime(e, tokens, cache);
690        }
691        let ranges = prime_chunk_ranges(t, self.layers.len());
692        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
693        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
694        // the prefill's ARITHMETIC, so two rigs with different values produced different
695        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
696        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
697        // (VERDICT.md) — and it is NOT what docs originally said:
698        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
699        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
700        //     output head), so growing a chunk cannot move an existing row's value.
701        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
702        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
703        //     not describe our leak.
704        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
705        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
706        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
707        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
708        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
709        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
710        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
711        // the source — every row is in one numeric class, so the chunk size no longer steers
712        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
713        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
714        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
715        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
716        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
717        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
718        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
719        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
720        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
721        // across calls, the request still ends at the same absolute position, whatever the tick
722        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
723        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
724        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
725        // default. Read per call, not cached (the probe flips it in-process between arms). Never
726        // on in a measured default run.
727        let legacy_calllocal =
728            std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
729        let seq_end = if legacy_calllocal {
730            cache.pos + t
731        } else {
732            cache.pos + t + queued_after
733        };
734        if ranges.len() == 1 {
735            return self.prime_chunk(e, tokens, cache, seq_end);
736        }
737        // PIPELINED PP-2 PRIME (lane/cx-pipeline-prime, 2026-08-08): overlap stage 0 of
738        // chunk N+1 with stage 1 of chunk N. The serial split stays reachable through
739        // MEMRA_PRIME_PIPE=0 and is the exactness oracle. N>2 keeps the serial walker;
740        // this lane owns the balanced two-stage schedule only.
741        if crate::pp::prime_pipe_on()
742            && crate::pp::prime_pp_on()
743            && !crate::pp::pp2_streams_off()
744        {
745            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()).filter(|f| f.len() == 3) {
746                if crate::pp::pp_multi_stream_same_device() {
747                    return Err(
748                        "prime chunk pipeline refused with 2 stage streams on one device — \
749                         that concurrent-stream placement remains quarantined by the deferred \
750                         pp flake record. Use one device per stage or MEMRA_PRIME_PIPE=0 for \
751                         the serial split."
752                            .into(),
753                    );
754                }
755                return self.prime_cache_pp2_pipelined(
756                    e, tokens, cache, seq_end, &ranges, &fence,
757                );
758            }
759        }
760        let mut hiddens = e.uninit(t * n_embd)?;
761        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
762        for &(start, end) in &ranges {
763            let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
764            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
765            last = Some((l, hs));
766        }
767        let (logits, h_seed) = last.unwrap();
768        Ok((logits, h_seed, hiddens))
769    }
770
771    /// PP-2 chunk scheduler: stage 1 of chunk N and stage 0 of chunk N+1 are both queued
772    /// before N's epilogue D2H drains the last-stage stream. Arithmetic is unchanged:
773    /// every chunk still runs the same two `prime_layers` ranges, boundary copy, output
774    /// norm, lm head, and caller hidden-stack copy as the serial split.
775    fn prime_cache_pp2_pipelined(
776        &self,
777        e: &Engine,
778        tokens: &[u32],
779        cache: &mut Cache,
780        seq_end: usize,
781        ranges: &[(usize, usize)],
782        fence: &[usize],
783    ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
784        debug_assert_eq!(fence.len(), 3);
785        debug_assert!(ranges.len() >= 2);
786        let rt = crate::pp::PpNRt::get(e)?;
787        assert_eq!(rt.n_stages(), 2, "prime pipeline requires exactly two PP stages");
788        let n_embd = self.cfg.n_embd as usize;
789        let t = tokens.len();
790        let initial_base = cache.pos;
791        let caller_stream = e.stream();
792
793        // #87 reverse publication before any new stage allocation, then prewarm both
794        // boundary slots while the stage streams are otherwise empty. Lazy-growing slot B
795        // after stage 1(N) is queued would synchronize that stream and erase the first
796        // overlap on a two-chunk prompt.
797        rt.fence_stages_behind(&caller_stream)?;
798        let max_payload = ranges
799            .iter()
800            .map(|(s, e)| (e - s) * n_embd)
801            .max()
802            .unwrap();
803        rt.prepare_overlap_slots(0, max_payload)?;
804
805        let mut hiddens = e.uninit(t * n_embd)?;
806        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
807        let mut stage_caches = PrimeCacheStages::new(cache, fence[1]);
808        let (cache0, cache1) = stage_caches.parts();
809        let (first_start, first_end) = ranges[0];
810        let mut slot = self.prime_pp2_stage0_enqueue(
811            e,
812            rt,
813            &tokens[first_start..first_end],
814            cache0,
815            seq_end,
816            fence,
817            initial_base + first_start,
818            true,
819        )?;
820        cache0.pos = initial_base + first_end;
821
822        for (i, &(start, end)) in ranges.iter().enumerate() {
823            let base = initial_base + start;
824            debug_assert_eq!(
825                cache1.pos, base,
826                "stage 1 must drain chunks in original position order"
827            );
828            let (out, next_slot) = if let Some(&(next_start, next_end)) = ranges.get(i + 1) {
829                let next_base = initial_base + next_start;
830                debug_assert_eq!(
831                    cache0.pos, next_base,
832                    "stage 0 must issue chunks in original position order"
833                );
834                let cache0_stage = &mut *cache0;
835                // Step's MoE router readback synchronizes once per layer. Two CUDA streams
836                // on one host thread therefore serialize even if the calls are ordered as
837                // a pipeline. Drive the disjoint stage caches from two scoped host threads:
838                // stage 1 consumes slot N while stage 0 produces slot N+1.
839                std::thread::scope(
840                    |scope| -> Result<_, Box<dyn std::error::Error>> {
841                        let stage0 = scope.spawn(move || -> Result<usize, String> {
842                            let next = self
843                                .prime_pp2_stage0_enqueue(
844                                    e,
845                                    rt,
846                                    &tokens[next_start..next_end],
847                                    cache0_stage,
848                                    seq_end,
849                                    fence,
850                                    next_base,
851                                    true,
852                                )
853                                .map_err(|err| err.to_string())?;
854                            cache0_stage.pos = initial_base + next_end;
855                            Ok(next)
856                        });
857                        let x = self.prime_pp2_stage1_enqueue(
858                            e,
859                            rt,
860                            slot,
861                            end - start,
862                            cache1,
863                            seq_end,
864                            fence,
865                            base,
866                            true,
867                        )?;
868                        let out = {
869                            rt.bind_stage(1)?;
870                            let _st1 = rt.enter(1);
871                            let e1 = rt.engine(1, e);
872                            self.prime_chunk_epilogue(e1, x, end - start, cache1)?
873                        };
874                        let next = stage0
875                            .join()
876                            .map_err(|_| "pipeprime stage-0 host walker panicked")?
877                            .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
878                        Ok((out, Some(next)))
879                    },
880                )?
881            } else {
882                let x = self.prime_pp2_stage1_enqueue(
883                    e,
884                    rt,
885                    slot,
886                    end - start,
887                    cache1,
888                    seq_end,
889                    fence,
890                    base,
891                    true,
892                )?;
893                let out = {
894                    rt.bind_stage(1)?;
895                    let _st1 = rt.enter(1);
896                    let e1 = rt.engine(1, e);
897                    self.prime_chunk_epilogue(e1, x, end - start, cache1)?
898                };
899                (out, None)
900            };
901
902            rt.publish_to(1, &caller_stream)?;
903            e.copy_into(
904                &mut hiddens,
905                start * n_embd,
906                &out.2,
907                (end - start) * n_embd,
908            )?;
909            last = Some((out.0, out.1));
910            crate::pp::PRIME_SPLIT_CHUNKS
911                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
912
913            if let Some(next) = next_slot {
914                // The caller copy above reads a stage-1 allocation. Before stage 1 of the
915                // next chunk can allocate/reuse blocks, mirror #87's body-entry fence.
916                // Stage 0(N+1) is already queued before this wait is appended, so its
917                // overlap with stage 1(N) is preserved.
918                rt.fence_stages_behind(&caller_stream)?;
919                slot = next;
920            }
921        }
922
923        debug_assert_eq!(cache0.pos, initial_base + t);
924        debug_assert_eq!(cache1.pos, initial_base + t);
925        let (logits, h_seed) = last.unwrap();
926        Ok((logits, h_seed, hiddens))
927    }
928
929    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
930    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
931    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
932    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
933    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
934    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
935    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
936        if Engine::gdn_db_on()
937            && Engine::gdn_chunked_enabled() && t >= 16
938            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
939            && num_k * 2 == num_v
940        {
941            num_k
942        } else {
943            num_v
944        }
945    }
946
947    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
948    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
949    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
950    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
951    fn f16out_on(e: &Engine, t: usize) -> bool {
952        crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
953            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
954    }
955
956    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
957    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
958    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
959    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
960    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
961    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
962    /// see one entry, byte-identical behavior.
963    pub fn prime_slabs_get(
964        &self,
965        e: &Engine,
966        t: usize,
967        n_embd: usize,
968        n_ff_max: usize,
969    ) -> Result<std::sync::Arc<std::sync::Mutex<PrimeSlabs>>, Box<dyn std::error::Error>> {
970        let mut slabs = self.prime_slabs.lock().unwrap();
971        let dev = e.ctx().ordinal();
972        let need_new = match slabs.get(&dev) {
973            None => true,
974            Some(sl) => sl.lock().unwrap().t_cap < t,
975        };
976        if need_new {
977            slabs.insert(dev, std::sync::Arc::new(std::sync::Mutex::new(PrimeSlabs {
978                t_cap: t,
979                h: e.uninit(t * n_embd)?,
980                x1: e.uninit(t * n_embd)?,
981                z: e.uninit(t * n_embd)?,
982                act: e.uninit(t * n_ff_max)?,
983                xa: e.uninit(t * n_embd)?,
984                xb: e.uninit(t * n_embd)?,
985                h16: e.alloc_u8_uninit(t * n_embd * 2)?,
986                z16: e.alloc_u8_uninit(t * n_embd * 2)?,
987                gate: e.uninit(t * n_ff_max)?,
988                up: e.uninit(t * n_ff_max)?,
989                ffn_out: e.uninit(t * n_embd)?,
990                seg_glue: Vec::new(),
991                mixed: e.uninit(t * n_embd)?,
992                seg_mid: Vec::new(),
993                seg_t: 0,
994            })));
995        }
996        Ok(slabs.get(&dev).expect("prime slab inserted").clone())
997    }
998
999    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
1000    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
1001    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
1002    fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize)
1003                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1004        if crate::pp::pp_host_bounce_active()
1005            && (self.cfg.gemma4.is_some() || !crate::pp::prime_pp_on())
1006        {
1007            return Err(
1008                "prime_chunk: refused with MEMRA_PP_HOST_BOUNCE=1 because this configuration \
1009                 has no active prime stage split and would peer-read remote weights; keep \
1010                 MEMRA_PRIME_PP enabled and use a PP-prime-supported model"
1011                    .into(),
1012            );
1013        }
1014        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
1015        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
1016        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
1017        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
1018        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
1019        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
1020        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
1021        // loader is off and there is nothing remote to split for.
1022        if self.cfg.gemma4.is_none()
1023            && !crate::pp::pp2_streams_off()
1024            && crate::pp::prime_pp_on()
1025        {
1026            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1027                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
1028            }
1029        }
1030        if crate::pp::pp_host_bounce_active() {
1031            return Err(
1032                "prime_chunk: MEMRA_PP_HOST_BOUNCE=1 found no valid prime stage split; \
1033                 refusing an unsplit remote-weight walk"
1034                    .into(),
1035            );
1036        }
1037        let t = tokens.len();
1038        let base = cache.pos;
1039        debug_assert!(seq_end >= base + t, "prime_chunk: seq_end must cover this chunk");
1040        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1041        let pos_d = e.htod_i32(&pos)?;
1042
1043        let x_embed = self.embed(e, tokens)?;   // [T, n_embd]
1044        let x = self.prime_layers(
1045            e, x_embed, 0, self.layers.len(), &pos_d, t, base, cache, seq_end,
1046        )?;
1047        self.prime_chunk_epilogue(e, x, t, cache)
1048    }
1049
1050    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
1051    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
1052    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
1053    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
1054    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
1055    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
1056    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
1057    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
1058    ///     the plain add (materialize) and the next stage hoists its own first norm — the
1059    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
1060    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
1061    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
1062    ///     each stage walks through its own resident transients;
1063    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
1064    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
1065    #[allow(clippy::too_many_arguments)]
1066    fn prime_layers(&self, e: &Engine, x_in: CudaSlice<f32>, lo: usize, hi: usize,
1067                    pos_d: &CudaSlice<i32>, t: usize, base: usize, cache: &mut Cache,
1068                    seq_end: usize)
1069                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1070        let cfg = &self.cfg;
1071        let n_embd = cfg.n_embd as usize;
1072        let eps = cfg.rms_eps;
1073        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
1074        // standalone convert launches). Only when the f16 lane serves and T reaches the
1075        // GEMM tier; bit-identical either way.
1076        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1077        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
1078        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
1079        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
1080        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
1081        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
1082        let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
1083            crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
1084            _ => n_embd,
1085        }).max().unwrap_or(n_embd).max(n_embd);
1086        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
1087        let slab = if use_slabs {
1088            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
1089        } else {
1090            None
1091        };
1092        let mut slab_guard = slab.as_ref().map(|sl| sl.lock().unwrap());
1093        let mut x_own;   // fallback storage when slabs are off
1094        type SlabRefs<'a> = (&'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<u8>, &'a mut CudaSlice<u8>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>, &'a mut CudaSlice<f32>);
1095        let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
1096        let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
1097        let mut x_own2;
1098        match slab_guard.as_mut() {
1099            Some(g) => {
1100                let slabs = &mut **g;
1101                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
1102                let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
1103                x_cur = xa;
1104                x_nxt = xb;
1105                seg = Some((seg_glue, seg_mid, mixed, seg_t));
1106                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
1107            }
1108            None => {
1109                x_own = x_in;
1110                x_own2 = e.uninit(t * n_embd)?;
1111                x_cur = &mut x_own;
1112                x_nxt = &mut x_own2;
1113                sl = None;
1114            }
1115        }
1116        let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
1117        let mut alloc_h16; let mut alloc_z16;
1118        let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
1119        let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1120        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
1121        let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
1122        match sl {
1123            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
1124                h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
1125                sl_gate = g; sl_up = u; sl_fo = fo;
1126            }
1127            None => {
1128                alloc_h = e.uninit(t * n_embd)?;
1129                alloc_x1 = e.uninit(t * n_embd)?;
1130                alloc_z = e.uninit(t * n_embd)?;
1131                alloc_act = e.uninit(t * n_ff_max)?;
1132                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1133                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1134                alloc_gate = e.uninit(t * n_ff_max)?;
1135                alloc_up = e.uninit(t * n_ff_max)?;
1136                alloc_fo = e.uninit(t * n_embd)?;
1137                h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
1138                h16 = &mut alloc_h16; z16 = &mut alloc_z16;
1139                sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
1140            }
1141        }
1142        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
1143        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
1144        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
1145        // first prime at this t (capture does not execute -> launch right after).
1146        let n_layers = self.layers.len();
1147        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
1148        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
1149        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
1150        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
1151        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
1152        // machinery stays (byte-identical) as their foundation.
1153        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
1154        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
1155        // step35 rides its own mixer through the normal per-layer arm below.
1156        let use_seg = f16fuse && seg.is_some() && self.cfg.step35.is_none()
1157            && lo == 0 && hi == n_layers
1158            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
1159        if let Some((sg, sm, _, st)) = seg.as_mut() {
1160            if **st != t {
1161                sg.clear();
1162                sg.extend((0..n_layers).map(|_| None));
1163                sm.clear();
1164                sm.extend((0..n_layers).map(|_| None));
1165                **st = t;
1166            }
1167        }
1168        {
1169            let layer_lo = &self.layers[lo];
1170            if f16fuse {
1171                e.rms_norm_f16out(x_cur, layer_lo.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
1172            } else {
1173                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
1174            }
1175        }
1176        for il in lo..hi {
1177            let layer = &self.layers[il];
1178            let hx16 = if f16fuse { Some(&*h16) } else { None };
1179            if use_seg {
1180                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
1181                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
1182                let (pre, pre16, w_out) = match &layer.mixer {
1183                    Mixer::Full(fa) => {
1184                        let g3 = match hx16 {
1185                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1186                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1187                        };
1188                        let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
1189                        (pre, pre16, &fa.wo)
1190                    }
1191                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1192                    Mixer::Linear(la) => {
1193                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1194                        let g4 = match hx16 {
1195                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1196                            None => e.matmul_group(&ws, h, t)?,
1197                        };
1198                        let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
1199                        (pre, pre16, &la.ssm_out)
1200                    }
1201                };
1202                {
1203                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
1204                    let pre_n = pre.len() / t;
1205                    let xh_pre = match pre16 {
1206                        Some(x) => x,
1207                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
1208                    };
1209                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
1210                        let y = e.matmul(w_out, &pre, t)?;
1211                        e.copy_into(mslab, 0, &y, t * n_embd)?;
1212                    }
1213                    if sm[il].is_none() {
1214                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1215                        let w_post = layer.post_attn_norm.float_data();
1216                        e.stream().synchronize()?;
1217                        e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1218                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1219                            e.add(x_cur, mslab, x1, t * n_embd)?;
1220                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
1221                            Ok(())
1222                        })();
1223                        let g = e.stream().end_capture(
1224                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1225                        r?;
1226                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
1227                    }
1228                    sm[il].as_ref().unwrap().launch()?;
1229                }
1230            } else {
1231                let mixed = match &layer.mixer {
1232                    Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il,
1233                                                            seq_end)?,
1234                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
1235                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1236                };
1237                if f16fuse {
1238                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
1239                    // bit-identical) — the standalone add pass disappears.
1240                    e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
1241                                          x1, z, z16, n_embd, t, eps)?;
1242                } else {
1243                    e.add(x_cur, &mixed, x1, t * n_embd)?;
1244                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
1245                }
1246            }
1247            let zx16 = if f16fuse { Some(&*z16) } else { None };
1248            match &layer.ffn {
1249                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1250                    let n_ff = ffn_gate.out_features();
1251                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
1252                    // the allocating group + copy when a mirror is missing.
1253                    let mut into_ok = false;
1254                    if let Some(xh) = zx16 {
1255                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
1256                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
1257                    }
1258                    if !into_ok {
1259                        let mut g2 = match zx16 {
1260                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
1261                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
1262                        };
1263                        let up_y = g2.pop().unwrap();
1264                        let gate_y = g2.pop().unwrap();
1265                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
1266                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
1267                    }
1268                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
1269                    // operand in-epilogue; non-silu activations keep the standalone convert.
1270                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
1271                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
1272                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1273                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none()
1274                        && d_lim.is_none() {
1275                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
1276                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
1277                        Some(a16)
1278                    } else {
1279                        Self::ffn_act_lim(e, &self.cfg, sl_gate, sl_up, 1.0, 1.0, d_lim,
1280                                          act, t * n_ff)?;
1281                        None
1282                    };
1283                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
1284                    let xh_act = match act16 {
1285                        Some(x) => x,
1286                        None => e.f16_act(act, t * n_ff, n_ff)?,
1287                    };
1288                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
1289                        let y = e.matmul(ffn_down, &*act, t)?;
1290                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1291                    }
1292                }
1293                crate::hybrid::Ffn::Moe(m) => {
1294                    let y = self.moe_ffn_il_prefill(e, m, z, t, il as u16)?;
1295                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
1296                }
1297            }
1298            if use_seg && il + 1 < hi {
1299                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
1300                let w_next = self.layers[il + 1].attn_norm.float_data();
1301                let (sg, _, _, _) = seg.as_mut().unwrap();
1302                if sg[il].is_none() {
1303                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
1304                    e.stream().synchronize()?;
1305                    e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
1306                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
1307                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1308                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
1309                        Ok(())
1310                    })();
1311                    let g = e.stream().end_capture(
1312                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
1313                    r?;
1314                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
1315                }
1316                sg[il].as_ref().unwrap().launch()?;
1317            } else {
1318                if il + 1 < hi {
1319                    let w_next = self.layers[il + 1].attn_norm.float_data();
1320                    if f16fuse {
1321                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
1322                    } else {
1323                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1324                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
1325                    }
1326                } else {
1327                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
1328                }
1329            }
1330            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
1331            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
1332            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
1333            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
1334            // unset (the default) costs one OnceLock read per layer.
1335            if let Some(path) = Self::prime_trace_path() {
1336                let row = (base + t - 1) as usize;
1337                let host = e.dtoh(x_nxt)?;
1338                let last = &host[(t - 1) * n_embd..t * n_embd];
1339                use std::io::Write as _;
1340                let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
1341                let mut h64: u64 = 0xcbf29ce484222325;
1342                for v in last {
1343                    h64 ^= v.to_bits() as u64;
1344                    h64 = h64.wrapping_mul(0x100000001b3);
1345                }
1346                writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
1347                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
1348                         last[0], last[1], last[2])?;
1349            }
1350            std::mem::swap(&mut x_cur, &mut x_nxt);
1351        }
1352        // hidden-stack return: clone the final x out of the slab
1353        let mut x = e.uninit(t * n_embd)?;
1354        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
1355        drop(slab_guard);
1356        Ok(x)
1357    }
1358
1359    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
1360    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
1361    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
1362    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
1363    fn prime_chunk_epilogue(&self, e: &Engine, x: CudaSlice<f32>, t: usize, cache: &mut Cache)
1364                            -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1365        let n_embd = self.cfg.n_embd as usize;
1366        let eps = self.cfg.rms_eps;
1367        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
1368        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
1369        // the post-norm copy happens after hn exists).
1370        let mut h_seed = e.uninit(n_embd)?;
1371        if !crate::spec::spec_hpost() {
1372            e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1373        }
1374        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
1375        let mut hn = e.uninit(t * n_embd)?;
1376        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1377        if crate::spec::spec_hpost() {
1378            e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
1379        }
1380        let last = e.view(&hn, t * n_embd);
1381        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
1382        let mut hlast = e.uninit(n_embd)?;
1383        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
1384        let logits = e.matmul(&self.output, &hlast, 1)?;
1385        cache.pos += t;
1386        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
1387        // post-norm stack hn (MEMRA_SPEC_HPOST).
1388        Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
1389    }
1390
1391    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
1392    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
1393    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
1394    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
1395    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
1396    /// prefill kernels. Structure mirrors the verify split exactly:
1397    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
1398    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
1399    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
1400    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
1401    ///                  there via the sharded loader) → `publish_to`
1402    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
1403    /// round's stage-freed buffers must not be reused under the caller's queued reads);
1404    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
1405    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
1406    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
1407    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
1408    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
1409    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
1410    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
1411    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
1412    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
1413    /// and its liveness counter is bumped here — the gate goes green with this function.
1414    fn prime_chunk_ppn(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize,
1415                       fence: &[usize])
1416                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
1417        let rt = crate::pp::PpNRt::get(e)?;
1418        let n_st = fence.len() - 1;
1419        assert_eq!(
1420            rt.n_stages(), n_st,
1421            "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
1422        );
1423        let n_embd = self.cfg.n_embd as usize;
1424        let t = tokens.len();
1425        let base = cache.pos;
1426        debug_assert!(seq_end >= base + t, "prime_chunk_ppn: seq_end must cover this chunk");
1427        let payload = t * n_embd;
1428        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1429        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1430        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1431        let caller_stream = e.stream();
1432        rt.fence_stages_behind(&caller_stream)?;
1433
1434        if n_st == 2 {
1435            let slot = self.prime_pp2_stage0_enqueue(
1436                e, rt, tokens, cache, seq_end, fence, base, false,
1437            )?;
1438            let x = self.prime_pp2_stage1_enqueue(
1439                e, rt, slot, t, cache, seq_end, fence, base, false,
1440            )?;
1441            let out = {
1442                rt.bind_stage(1)?;
1443                let _st1 = rt.enter(1);
1444                let e1 = rt.engine(1, e);
1445                self.prime_chunk_epilogue(e1, x, t, cache)?
1446            };
1447            rt.publish_to(1, &caller_stream)?;
1448            crate::pp::PRIME_SPLIT_CHUNKS
1449                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1450            return Ok(out);
1451        }
1452
1453        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1454
1455        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1456        let mut slot = {
1457            let _st0 = rt.enter(0);
1458            let e0 = rt.engine(0, e);
1459            let pos_d = e0.htod_i32(&pos)?;
1460            let x = self.embed(e0, tokens)?;
1461            let x = self.prime_layers(
1462                e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1463            )?;
1464            rt.tx(0, &x, payload)?
1465            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1466        };
1467
1468        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1469        for s in 1..n_st - 1 {
1470            let _st = rt.enter(s);
1471            let es = rt.engine(s, e);
1472            let pos_d = es.htod_i32(&pos)?;
1473            let x = rt.rx(s - 1, slot, payload)?;
1474            let x = self.prime_layers(
1475                es, x, fence[s], fence[s + 1], &pos_d, t, base, cache, seq_end,
1476            )?;
1477            slot = rt.tx(s, &x, payload)?;
1478        }
1479
1480        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1481        let _stl = rt.enter(n_st - 1);
1482        let el = rt.engine(n_st - 1, e);
1483        let pos_d = el.htod_i32(&pos)?;
1484        let x = rt.rx(n_st - 2, slot, payload)?;
1485        let x = self.prime_layers(
1486            el, x, fence[n_st - 1], fence[n_st], &pos_d, t, base, cache, seq_end,
1487        )?;
1488        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1489        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1490        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1491        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1492        // stage stream host-side, but the law is stated in events, not in a dtoh side
1493        // effect a later deferred form would remove.
1494        rt.publish_to(n_st - 1, &caller_stream)?;
1495        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1496        Ok(out)
1497    }
1498
1499    fn prime_pp2_stage0_enqueue(
1500        &self,
1501        e: &Engine,
1502        rt: &crate::pp::PpNRt,
1503        tokens: &[u32],
1504        cache: &mut Cache,
1505        seq_end: usize,
1506        fence: &[usize],
1507        base: usize,
1508        pipelined: bool,
1509    ) -> Result<usize, Box<dyn std::error::Error>> {
1510        let t = tokens.len();
1511        let n_embd = self.cfg.n_embd as usize;
1512        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1513        rt.bind_stage(0)?;
1514        let _st0 = rt.enter(0);
1515        let e0 = rt.engine(0, e);
1516        let pos_d = e0.htod_i32(&pos)?;
1517        let x = self.embed(e0, tokens)?;
1518        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1519        let x = self.prime_layers(
1520            e0, x, fence[0], fence[1], &pos_d, t, base, cache, seq_end,
1521        )?;
1522        if pipelined {
1523            rt.tx_pipelined(0, &x, t * n_embd)
1524        } else {
1525            rt.tx(0, &x, t * n_embd)
1526        }
1527    }
1528
1529    fn prime_pp2_stage1_enqueue(
1530        &self,
1531        e: &Engine,
1532        rt: &crate::pp::PpNRt,
1533        slot: usize,
1534        t: usize,
1535        cache: &mut Cache,
1536        seq_end: usize,
1537        fence: &[usize],
1538        base: usize,
1539        pipelined: bool,
1540    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1541        let n_embd = self.cfg.n_embd as usize;
1542        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1543        rt.bind_stage(1)?;
1544        let _st1 = rt.enter(1);
1545        let e1 = rt.engine(1, e);
1546        let pos_d = e1.htod_i32(&pos)?;
1547        let x = rt.rx(0, slot, t * n_embd)?;
1548        let _overlap = pipelined.then(crate::pp::enter_prime_pipe_stage);
1549        self.prime_layers(
1550            e1, x, fence[1], fence[2], &pos_d, t, base, cache, seq_end,
1551        )
1552    }
1553
1554    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
1555    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
1556    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
1557    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
1558    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
1559    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
1560    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
1561    /// bookkeeping still runs on the host per call — the real replay path moves the write
1562    /// slot to the len_d device counter (increment 3).
1563    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
1564    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
1565    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
1566    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
1567    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
1568    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
1569    pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
1570                                t: usize, cache: &mut Cache,
1571                                len_d: &CudaSlice<i32>,
1572                                logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
1573                                -> Result<(), Box<dyn std::error::Error>> {
1574        let cfg = &self.cfg;
1575        let n_embd = cfg.n_embd as usize;
1576        let eps = cfg.rms_eps;
1577        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1578        let mut x = e.uninit(t * n_embd)?;
1579        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1580        for (il, layer) in self.layers.iter().enumerate() {
1581            let mut h = e.uninit(t * n_embd)?;
1582            let mut hx16: Option<CudaSlice<u8>> = None;
1583            if f16fuse {
1584                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1585                e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
1586                hx16 = Some(b16);
1587            } else {
1588                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1589            }
1590            let mixed = match &layer.mixer {
1591                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
1592                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
1593                // come from the caller (see step35_attn_pre_wo's doc note).
1594                Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache,
1595                                                        il, t)?,
1596                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1597                Mixer::Linear(la) => {
1598                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1599                    let g4 = match hx16.as_ref() {
1600                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1601                        None => e.matmul_group(&ws, &h, t)?,
1602                    };
1603                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1604                }
1605            };
1606            let mut x1 = e.uninit(t * n_embd)?;
1607            e.add(&x, &mixed, &mut x1, t * n_embd)?;
1608            let mut z = e.uninit(t * n_embd)?;
1609            let mut zx16: Option<CudaSlice<u8>> = None;
1610            if f16fuse {
1611                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1612                e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
1613                zx16 = Some(b16);
1614            } else {
1615                e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
1616            }
1617            let ffn_out = match &layer.ffn {
1618                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1619                    let n_ff = ffn_gate.out_features();
1620                    let mut g2 = match &zx16 {
1621                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1622                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1623                    };
1624                    let up = g2.pop().unwrap();
1625                    let gate = g2.pop().unwrap();
1626                    let mut act = e.uninit(t * n_ff)?;
1627                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
1628                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
1629                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
1630                    e.matmul(ffn_down, &act, t)?
1631                }
1632                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_prefill(e, m, &z, t, il as u16)?,
1633            };
1634            let mut x2 = e.uninit(t * n_embd)?;
1635            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1636            x = x2;
1637        }
1638        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
1639        if !crate::spec::spec_hpost() {
1640            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1641        }
1642        let mut hn = e.uninit(t * n_embd)?;
1643        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1644        if crate::spec::spec_hpost() {
1645            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1646        }
1647        let mut hlast = e.uninit(n_embd)?;
1648        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1649        let logits = e.matmul(&self.output, &hlast, 1)?;
1650        let nv = logits.len();
1651        e.copy_into(logits_out, 0, &logits, nv)?;
1652        Ok(())
1653    }
1654
1655    fn step35_prime_batch_on() -> bool {
1656        std::env::var("MEMRA_STEP35_PRIME_BATCH").as_deref() != Ok("0")
1657    }
1658
1659    /// Step35 cross-request prime range: weight-streaming work runs once at `m=sum(T)`;
1660    /// sequence-scoped attention/KV work stays on each request's own cache and positions.
1661    #[allow(clippy::too_many_arguments)]
1662    fn step35_prime_batch_layers(
1663        &self,
1664        e: &Engine,
1665        mut x: CudaSlice<f32>,
1666        lo: usize,
1667        hi: usize,
1668        ts: &[usize],
1669        offs: &[usize],
1670        pos_ds: &[CudaSlice<i32>],
1671        caches: &mut [&mut Cache],
1672    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1673        let cfg = &self.cfg;
1674        let n_embd = cfg.n_embd as usize;
1675        let eps = cfg.rms_eps;
1676        let b = ts.len();
1677        let total: usize = ts.iter().sum();
1678        let f16fuse = crate::f16_ffi::pp_f16_enabled() && total >= 16;
1679
1680        let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
1681                     -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1682            let mut out = Vec::with_capacity(b);
1683            for s in 0..b {
1684                let mut ys = e.uninit(ts[s] * dim)?;
1685                e.copy_view_into(
1686                    &mut ys,
1687                    0,
1688                    &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim),
1689                    ts[s] * dim,
1690                )?;
1691                out.push(ys);
1692            }
1693            Ok(out)
1694        };
1695
1696        for il in lo..hi {
1697            let layer = &self.layers[il];
1698            let Mixer::Full(fa) = &layer.mixer else {
1699                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
1700            };
1701
1702            let mut h = e.uninit(total * n_embd)?;
1703            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1704            if f16fuse {
1705                e.rms_norm_f16out(
1706                    &x,
1707                    layer.attn_norm.float_data(),
1708                    &mut h,
1709                    &mut hx16,
1710                    n_embd,
1711                    total,
1712                    eps,
1713                )?;
1714            } else {
1715                e.rms_norm(
1716                    &x,
1717                    layer.attn_norm.float_data(),
1718                    &mut h,
1719                    n_embd,
1720                    total,
1721                    eps,
1722                )?;
1723            }
1724
1725            // Q/K/V + gate projections and wo are batched weight streams. The step35 core
1726            // remains per sequence, so its partial RoPE, SWA view, KV append, and gate
1727            // application stay verbatim.
1728            let gate_w = fa
1729                .attn_gate
1730                .as_ref()
1731                .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
1732            let mut g4 = if f16fuse {
1733                e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, &hx16, total)?
1734            } else {
1735                e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv, gate_w], &h, total)?
1736            };
1737            let gate = g4.pop().unwrap();
1738            let mut parts: Vec<Vec<CudaSlice<f32>>> =
1739                (0..b).map(|_| Vec::with_capacity(3)).collect();
1740            for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g4) {
1741                for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1742                    parts[s].push(ys);
1743                }
1744            }
1745            let gates = split(e, &gate, gate_w.out_features())?;
1746            let geometry = self.step35_geom(il);
1747            let hd = geometry.head_dim_k as usize;
1748            let nh = geometry.n_head as usize;
1749            let mut ag_cat = e.uninit(total * nh * hd)?;
1750            for (s, (g3s, gate)) in parts.into_iter().zip(gates).enumerate() {
1751                let ag = self.step35_attn_pre_wo(
1752                    e,
1753                    fa,
1754                    g3s,
1755                    None,
1756                    Some(&gate),
1757                    &pos_ds[s],
1758                    ts[s],
1759                    Some(&mut *caches[s]),
1760                    il,
1761                    ts[s],
1762                )?;
1763                e.copy_into(
1764                    &mut ag_cat,
1765                    offs[s] * nh * hd,
1766                    &ag,
1767                    ts[s] * nh * hd,
1768                )?;
1769            }
1770            let mixed = e.matmul(&fa.wo, &ag_cat, total)?;
1771
1772            let mut x1 = e.uninit(total * n_embd)?;
1773            let mut z = e.uninit(total * n_embd)?;
1774            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1775            if f16fuse {
1776                e.add_rms_norm_f16out(
1777                    &x,
1778                    &mixed,
1779                    layer.post_attn_norm.float_data(),
1780                    &mut x1,
1781                    &mut z,
1782                    &mut zx16,
1783                    n_embd,
1784                    total,
1785                    eps,
1786                )?;
1787            } else {
1788                e.add(&x, &mixed, &mut x1, total * n_embd)?;
1789                e.rms_norm(
1790                    &x1,
1791                    layer.post_attn_norm.float_data(),
1792                    &mut z,
1793                    n_embd,
1794                    total,
1795                    eps,
1796                )?;
1797            }
1798
1799            let ffn_out = match &layer.ffn {
1800                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1801                    let n_ff = ffn_gate.out_features();
1802                    let mut g2 = if f16fuse {
1803                        e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?
1804                    } else {
1805                        e.matmul_group(&[ffn_gate, ffn_up], &z, total)?
1806                    };
1807                    let up = g2.pop().unwrap();
1808                    let gate = g2.pop().unwrap();
1809                    let mut act = e.uninit(total * n_ff)?;
1810                    let d_lim = cfg.clamp_shexp_at(il as u32);
1811                    if Self::f16out_on(e, total) && cfg.m3.is_none() && d_lim.is_none() {
1812                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1813                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1814                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1815                            Some(y) => y,
1816                            None => e.matmul(ffn_down, &act, total)?,
1817                        }
1818                    } else {
1819                        Self::ffn_act_lim(
1820                            e,
1821                            cfg,
1822                            &gate,
1823                            &up,
1824                            1.0,
1825                            1.0,
1826                            d_lim,
1827                            &mut act,
1828                            total * n_ff,
1829                        )?;
1830                        e.matmul(ffn_down, &act, total)?
1831                    }
1832                }
1833                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1834            };
1835            let mut x2 = e.uninit(total * n_embd)?;
1836            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1837            x = x2;
1838        }
1839        Ok(x)
1840    }
1841
1842    fn step35_prime_batch_epilogue(
1843        &self,
1844        e: &Engine,
1845        x: CudaSlice<f32>,
1846        ts: &[usize],
1847        offs: &[usize],
1848        caches: &mut [&mut Cache],
1849    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1850        let n_embd = self.cfg.n_embd as usize;
1851        let total: usize = ts.iter().sum();
1852        let mut hn = e.uninit(total * n_embd)?;
1853        e.rms_norm(
1854            &x,
1855            self.output_norm.float_data(),
1856            &mut hn,
1857            n_embd,
1858            total,
1859            self.cfg.rms_eps,
1860        )?;
1861
1862        let hidden_src = if crate::spec::spec_hpost() { &hn } else { &x };
1863        let mut out = Vec::with_capacity(ts.len());
1864        for s in 0..ts.len() {
1865            let mut hidden = e.uninit(ts[s] * n_embd)?;
1866            e.copy_view_into(
1867                &mut hidden,
1868                0,
1869                &hidden_src.slice(offs[s] * n_embd..(offs[s] + ts[s]) * n_embd),
1870                ts[s] * n_embd,
1871            )?;
1872            let last0 = (offs[s] + ts[s] - 1) * n_embd;
1873            let mut h_seed = e.uninit(n_embd)?;
1874            e.copy_view_into(
1875                &mut h_seed,
1876                0,
1877                &hidden_src.slice(last0..last0 + n_embd),
1878                n_embd,
1879            )?;
1880            // Exactness-first: the serial reference runs the output head at m=1.
1881            let mut hlast = e.uninit(n_embd)?;
1882            e.copy_view_into(
1883                &mut hlast,
1884                0,
1885                &hn.slice(last0..last0 + n_embd),
1886                n_embd,
1887            )?;
1888            let logits = e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?;
1889            caches[s].pos += ts[s];
1890            out.push((logits, h_seed, hidden));
1891        }
1892        Ok(out)
1893    }
1894
1895    fn step35_prime_cache_batch(
1896        &self,
1897        e: &Engine,
1898        prompts: &[&[u32]],
1899        caches: &mut [&mut Cache],
1900    ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1901        if crate::pp::pp_host_bounce_active()
1902            && (!crate::pp::prime_pp_on()
1903                || crate::pp::pp_cuts(self.layers.len()).is_none())
1904        {
1905            return Err(
1906                "step35_prime_cache_batch: MEMRA_PP_HOST_BOUNCE=1 requires a valid prime \
1907                 stage split; refusing an unsplit remote-weight walk"
1908                    .into(),
1909            );
1910        }
1911        if !Self::step35_prime_batch_on() {
1912            return Err("step35 batched prime is disabled (MEMRA_STEP35_PRIME_BATCH=0)".into());
1913        }
1914        if caches.iter().any(|c| c.pos != 0) {
1915            return Err(
1916                "step35 batched prime currently supports complete fresh prompts only; \
1917                 continuation/tick chunks require per-request queued_after"
1918                    .into(),
1919            );
1920        }
1921
1922        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
1923        for &t in &ts {
1924            assert!(t >= PRIME_MIN_T, "step35 batched prime needs T >= {PRIME_MIN_T}");
1925        }
1926        for (s, c) in caches.iter().enumerate() {
1927            assert!(ts[s] <= c.max_ctx, "step35 batched prime exceeds cache max_ctx");
1928        }
1929        let offs: Vec<usize> = ts
1930            .iter()
1931            .scan(0usize, |a, &t| {
1932                let o = *a;
1933                *a += t;
1934                Some(o)
1935            })
1936            .collect();
1937        let total: usize = ts.iter().sum();
1938        let payload = total * self.cfg.n_embd as usize;
1939        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
1940        let positions: Vec<Vec<i32>> = ts
1941            .iter()
1942            .map(|&t| (0..t as i32).collect())
1943            .collect();
1944        let upload_positions = |e: &Engine|
1945                                -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
1946            positions
1947                .iter()
1948                .map(|p| e.htod_i32(p))
1949                .collect::<Result<_, _>>()
1950        };
1951
1952        static ONCE: std::sync::Once = std::sync::Once::new();
1953        ONCE.call_once(|| {
1954            eprintln!(
1955                "[step35-prime-batch] first concat prime: B={} tokens={total}",
1956                prompts.len()
1957            );
1958        });
1959
1960        let out = if !crate::pp::pp2_streams_off() && crate::pp::prime_pp_on() {
1961            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
1962                let rt = crate::pp::PpNRt::get(e)?;
1963                let n_st = fence.len() - 1;
1964                assert_eq!(rt.n_stages(), n_st, "step35 prime batch stage count mismatch");
1965                let caller_stream = e.stream();
1966                rt.fence_stages_behind(&caller_stream)?;
1967
1968                let mut slot = {
1969                    let _st0 = rt.enter(0);
1970                    let e0 = rt.engine(0, e);
1971                    let pos_ds = upload_positions(e0)?;
1972                    let x = self.embed(e0, &cat_tokens)?;
1973                    let x = self.step35_prime_batch_layers(
1974                        e0,
1975                        x,
1976                        fence[0],
1977                        fence[1],
1978                        &ts,
1979                        &offs,
1980                        &pos_ds,
1981                        caches,
1982                    )?;
1983                    rt.tx(0, &x, payload)?
1984                };
1985                for s in 1..n_st - 1 {
1986                    let _st = rt.enter(s);
1987                    let es = rt.engine(s, e);
1988                    let pos_ds = upload_positions(es)?;
1989                    let x = rt.rx(s - 1, slot, payload)?;
1990                    let x = self.step35_prime_batch_layers(
1991                        es,
1992                        x,
1993                        fence[s],
1994                        fence[s + 1],
1995                        &ts,
1996                        &offs,
1997                        &pos_ds,
1998                        caches,
1999                    )?;
2000                    slot = rt.tx(s, &x, payload)?;
2001                }
2002
2003                let _stl = rt.enter(n_st - 1);
2004                let el = rt.engine(n_st - 1, e);
2005                let pos_ds = upload_positions(el)?;
2006                let x = rt.rx(n_st - 2, slot, payload)?;
2007                let x = self.step35_prime_batch_layers(
2008                    el,
2009                    x,
2010                    fence[n_st - 1],
2011                    fence[n_st],
2012                    &ts,
2013                    &offs,
2014                    &pos_ds,
2015                    caches,
2016                )?;
2017                let out = self.step35_prime_batch_epilogue(el, x, &ts, &offs, caches)?;
2018                rt.publish_to(n_st - 1, &caller_stream)?;
2019                crate::pp::STEP35_PRIME_BATCH_SPLITS.fetch_add(
2020                    1,
2021                    std::sync::atomic::Ordering::Relaxed,
2022                );
2023                out
2024            } else {
2025                let pos_ds = upload_positions(e)?;
2026                let x = self.embed(e, &cat_tokens)?;
2027                let x = self.step35_prime_batch_layers(
2028                    e,
2029                    x,
2030                    0,
2031                    self.layers.len(),
2032                    &ts,
2033                    &offs,
2034                    &pos_ds,
2035                    caches,
2036                )?;
2037                self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2038            }
2039        } else {
2040            let pos_ds = upload_positions(e)?;
2041            let x = self.embed(e, &cat_tokens)?;
2042            let x = self.step35_prime_batch_layers(
2043                e,
2044                x,
2045                0,
2046                self.layers.len(),
2047                &ts,
2048                &offs,
2049                &pos_ds,
2050                caches,
2051            )?;
2052            self.step35_prime_batch_epilogue(e, x, &ts, &offs, caches)?
2053        };
2054        crate::pp::STEP35_PRIME_BATCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2055        Ok(out)
2056    }
2057
2058    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
2059    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
2060    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
2061    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
2062    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
2063    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
2064    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
2065    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
2066    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
2067    /// over the quantized past; Linear: the stateful pad_view twin — the same state
2068    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
2069    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
2070    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
2071    /// back to single-chunk serving).
2072    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
2073    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
2074    pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
2075                             -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
2076        let cfg = &self.cfg;
2077        let n_embd = cfg.n_embd as usize;
2078        let eps = cfg.rms_eps;
2079        let b = prompts.len();
2080        assert!(b >= 1 && b == caches.len());
2081        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
2082        let carried = pos0s.iter().any(|&p| p > 0);
2083        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
2084        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
2085        // generic concat attn core below (uniform geometry, no per-layer swa window, no
2086        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
2087        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
2088        if cfg.gemma4.is_some() {
2089            return Err("prime_cache_batch: gemma4 has no batched prime core (per-layer \
2090                        swa/global geometry, softcapped head) — use gemma4_prime per sequence".into());
2091        }
2092        // Step35 has a dedicated concat walk: the generic core below cannot express its
2093        // per-layer geometry/SWA/head gate, and under PP the dedicated path is stage-scoped.
2094        if cfg.step35.is_some() {
2095            return self.step35_prime_cache_batch(e, prompts, caches);
2096        }
2097        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
2098        for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
2099        for (s, c) in caches.iter().enumerate() {
2100            assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
2101        }
2102        let total: usize = ts.iter().sum();
2103        let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
2104        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
2105        let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
2106            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
2107            .collect::<Result<_, _>>()?;
2108        // split a concat [total, dim] buffer into per-seq copies
2109        let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
2110                     -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
2111            let mut out = Vec::with_capacity(b);
2112            for s in 0..b {
2113                let mut ys = e.uninit(ts[s] * dim)?;
2114                e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
2115                out.push(ys);
2116            }
2117            Ok(out)
2118        };
2119
2120        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
2121        let mut x = self.embed(e, &cat_tokens)?;   // [total, n_embd]
2122        for (il, layer) in self.layers.iter().enumerate() {
2123            let mut h = e.uninit(total * n_embd)?;
2124            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2125            e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
2126            // mixer: projection GROUP on the concat (m = total), stateful core per seq
2127            let mut mixed = e.uninit(total * n_embd)?;
2128            match &layer.mixer {
2129                Mixer::Full(fa) => {
2130                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
2131                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
2132                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
2133                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
2134                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
2135                    // back to the per-seq dispatch.
2136                    let geometry = self.cfg.full_attention_geometry_at(il as u32);
2137                    let (n_head, n_head_kv, head_dim) = (
2138                        geometry.n_head as usize,
2139                        geometry.n_head_kv as usize,
2140                        geometry.head_dim_k as usize,
2141                    );
2142                    let fa_scale = geometry.attention_scale();
2143                    let use_favl = !carried
2144                        && (2..=8).contains(&b)
2145                        && (head_dim == 256 || head_dim == 128)
2146                        && geometry.attention_gate
2147                            == memra_gguf::config::AttentionGateKind::FusedQ
2148                        && std::env::var("MEMRA_NOFA").is_err()
2149                        && std::env::var("MEMRA_FA_FLOOR").is_err()
2150                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
2151                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
2152                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
2153                    if use_favl {
2154                        let (qf_w, kf_w, vf_w) =
2155                            (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
2156                        struct APre {
2157                            q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
2158                            qn: CudaSlice<f32>, kn: CudaSlice<f32>,
2159                        }
2160                        let mut aps = Vec::with_capacity(b);
2161                        for &t in ts.iter().take(b) {
2162                            aps.push(APre {
2163                                q: e.uninit(t * n_head * head_dim)?,
2164                                gate: Some(e.uninit(t * n_head * head_dim)?),
2165                                qn: e.uninit(t * n_head * head_dim)?,
2166                                kn: e.uninit(t * n_head_kv * head_dim)?,
2167                            });
2168                        }
2169                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
2170                            let kvl = caches[0].kv[il].as_ref().unwrap();
2171                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
2172                        };
2173                        let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
2174                            let (o, t) = (offs[s], ts[s]);
2175                            let kvl = caches[s].kv[il].as_ref().unwrap();
2176                            assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
2177                                    "prime_cache_batch attn vl: fresh + capacity");
2178                            crate::AttnPreVl {
2179                                qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
2180                                kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
2181                                vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
2182                                q: e.addr_f32(&aps[s].q),
2183                                gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
2184                                qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
2185                                kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
2186                                t: t as i32, pad: 0,
2187                            }
2188                        }).collect();
2189                        e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
2190                                       head_dim, geometry.n_rot as usize, n_head, n_head_kv,
2191                                       self.cfg.rms_eps, geometry.rope_base, 1.0,
2192                                       kv_dim_k, kv_dim_v, ktb, vtb)?;
2193                        for s in 0..b {
2194                            let kvl = caches[s].kv[il].as_mut().unwrap();
2195                            kvl.len += ts[s];
2196                            let new_len = kvl.len as i32;
2197                            e.set_i32_one(&mut kvl.len_d, new_len)?;
2198                        }
2199                        let mut attns = Vec::with_capacity(b);
2200                        let mut mirrors = Vec::with_capacity(b);
2201                        for &t in ts.iter().take(b) {
2202                            attns.push(e.uninit(t * n_head * head_dim)?);
2203                            let n = t * n_head_kv * head_dim;
2204                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
2205                        }
2206                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
2207                        // promoted single-seq config is on; else the mma favl.
2208                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
2209                            Ok("0") => false,
2210                            Ok("1") => true,
2211                            _ => cfg!(memra_hopper_mma),
2212                        };
2213                        if fa3_on {
2214                            let mut q16s = Vec::with_capacity(b);
2215                            let mut v16s = Vec::with_capacity(b);
2216                            for s in 0..b {
2217                                let t = ts[s];
2218                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
2219                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
2220                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2221                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
2222                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
2223                                e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
2224                                                &mut v16, t * n_head_kv * head_dim)?;
2225                                q16s.push(q16);
2226                                v16s.push((k16, v16));
2227                            }
2228                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
2229                            let mut kp = qp;
2230                            let mut vp = qp;
2231                            let mut op = [core::ptr::null_mut::<f32>(); 8];
2232                            let mut tsv = [0i32; 8];
2233                            for s in 0..b {
2234                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
2235                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
2236                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
2237                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
2238                                tsv[s] = ts[s] as i32;
2239                            }
2240                            let rc = unsafe {
2241                                crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
2242                                                  tsv.as_ptr(), b as i32, n_head as i32,
2243                                                  n_head_kv as i32, head_dim as i32, fa_scale,
2244                                                  e.stream().cu_stream() as *mut core::ffi::c_void)
2245                            };
2246                            if rc != 0 {
2247                                return Err(format!("memra_fa3_vl rc={rc}").into());
2248                            }
2249                        } else {
2250                            let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
2251                                q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
2252                                v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
2253                                kf: e.addr_f32(&aps[s].kn),
2254                                vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
2255                                t: ts[s] as i32, pad: 0,
2256                            }).collect();
2257                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
2258                        }
2259                        for (s, attn) in attns.into_iter().enumerate() {
2260                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
2261                                e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
2262                            let mut done = false;
2263                            if let Some(xh) = &ag16 {
2264                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2265                            }
2266                            if !done {
2267                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2268                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2269                            }
2270                        }
2271                    } else {
2272                        let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
2273                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
2274                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
2275                                parts[s].push(ys);
2276                            }
2277                        }
2278                        for (s, g3s) in parts.into_iter().enumerate() {
2279                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
2280                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
2281                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
2282                            let mut done = false;
2283                            if let Some(xh) = &ag16 {
2284                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
2285                            }
2286                            if !done {
2287                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
2288                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
2289                            }
2290                        }
2291                    }
2292                }
2293                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
2294                Mixer::Linear(la) => {
2295                    // task #16: NO split copies (cores read row-offset views of the concat
2296                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
2297                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
2298                    // varlen K5 launch for all sequences.
2299                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2300                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
2301                    let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
2302                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
2303                        let (o, t) = (offs[s], ts[s]);
2304                        let mut done = false;
2305                        if let Some(xh) = &gn16 {
2306                            done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
2307                        }
2308                        if !done {
2309                            let m = e.matmul(&la.ssm_out, &gn, t)?;
2310                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
2311                        }
2312                    }
2313                }
2314            }
2315            let mut x1 = e.uninit(total * n_embd)?;
2316            let mut z = e.uninit(total * n_embd)?;
2317            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
2318            e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
2319                                  &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
2320            let ffn_out = match &layer.ffn {
2321                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
2322                    let n_ff = ffn_gate.out_features();
2323                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
2324                    let up = g2.pop().unwrap();
2325                    let gate = g2.pop().unwrap();
2326                    let mut act = e.uninit(total * n_ff)?;
2327                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
2328                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
2329                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
2330                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
2331                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
2332                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
2333                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
2334                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
2335                            Some(y) => y,
2336                            None => e.matmul(ffn_down, &act, total)?,
2337                        }
2338                    } else {
2339                        Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, d_lim,
2340                                          &mut act, total * n_ff)?;
2341                        e.matmul(ffn_down, &act, total)?
2342                    }
2343                }
2344                crate::hybrid::Ffn::Moe(m) => {
2345                    self.moe_ffn_il_prefill(e, m, &z, total, il as u16)?
2346                }
2347            };
2348            let mut x2 = e.uninit(total * n_embd)?;
2349            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
2350            x = x2;
2351        }
2352        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
2353        let mut hn = e.uninit(total * n_embd)?;
2354        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
2355        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
2356        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
2357        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
2358        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
2359        // argmax battery arbitrates, same as every other prefill GEMM change.
2360        let mut hcat = e.uninit(b * n_embd)?;
2361        for s in 0..b {
2362            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2363            e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
2364        }
2365        let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
2366        let logits_host: Option<Vec<f32>> = match &logits_cat {
2367            Some(lc) => Some(e.dtoh(lc)?),
2368            None => None,
2369        };
2370        let n_vocab = self.output.out_features();
2371        let mut hidden_all = if crate::spec::spec_hpost() {
2372            split(e, &hn, n_embd)?
2373        } else {
2374            split(e, &x, n_embd)?
2375        };
2376        let mut out = Vec::with_capacity(b);
2377        for s in 0..b {
2378            let last0 = (offs[s] + ts[s] - 1) * n_embd;
2379            let mut h_seed = e.uninit(n_embd)?;
2380            if !crate::spec::spec_hpost() {
2381                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
2382            } else {
2383                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2384            }
2385            let logits = match &logits_host {
2386                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
2387                None => {
2388                    let mut hlast = e.uninit(n_embd)?;
2389                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
2390                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
2391                }
2392            };
2393            caches[s].pos += ts[s];
2394            out.push((logits, h_seed, hidden_all.remove(0)));
2395        }
2396        Ok(out)
2397    }
2398
2399    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
2400    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
2401    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
2402    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
2403    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
2404    ///
2405    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
2406    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
2407    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
2408    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
2409    #[allow(clippy::too_many_arguments)]
2410    fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
2411                       hx: Option<&CudaSlice<u8>>,
2412                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize,
2413                       seq_end: usize)
2414                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2415        if self.cfg.step35.is_some() {
2416            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
2417        }
2418        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
2419        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
2420        // this single-seq path composes proj+core identically (byte-for-byte the old body).
2421        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
2422        let g3 = match hx {
2423            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
2424            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
2425        };
2426        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
2427    }
2428
2429    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
2430    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
2431    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
2432    fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2433                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2434                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2435        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
2436        if let Some(xh) = &ag16 {
2437            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
2438                return Ok(y);
2439            }
2440        }
2441        Ok(e.matmul(&fa.wo, &attn_g, t)?)
2442    }
2443
2444    fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
2445                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2446                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2447        let cfg = &self.cfg;
2448        let geometry = cfg.full_attention_geometry_at(il as u32);
2449        let n_head = geometry.n_head as usize;
2450        let n_head_kv = geometry.n_head_kv as usize;
2451        let head_dim = geometry.head_dim_k as usize;
2452        let scale = geometry.attention_scale();
2453        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
2454        let AttnPre { q, k, v, gate } = pre;
2455        let mut attn = e.uninit(t * n_head * head_dim)?;
2456        self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
2457                                         head_dim, n_head, n_head_kv, scale)?;
2458        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
2459    }
2460
2461    /// task #18 (attn side): projections tail through KV append — everything before the
2462    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
2463    /// present BEFORE this chunk's append (base_len; 0 == fresh).
2464    #[allow(clippy::type_complexity)]
2465    fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
2466                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
2467                            -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
2468        let cfg = &self.cfg;
2469        let geometry = cfg.full_attention_geometry_at(il as u32);
2470        let n_head = geometry.n_head as usize;
2471        let n_head_kv = geometry.n_head_kv as usize;
2472        let head_dim = geometry.head_dim_k as usize;
2473        let eps = cfg.rms_eps;
2474
2475        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
2476        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
2477        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
2478        let gated = geometry.attention_gate
2479            == memra_gguf::config::AttentionGateKind::FusedQ;
2480        let v = g3.pop().unwrap();
2481        let mut k = g3.pop().unwrap();
2482        let qf = g3.pop().unwrap();
2483        let (mut q, gate) = if gated {
2484            let mut q = e.uninit(t * n_head * head_dim)?;
2485            let mut gate = e.uninit(t * n_head * head_dim)?;
2486            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2487            (q, Some(gate))
2488        } else {
2489            (qf, None)
2490        };
2491
2492        let mut qn = e.uninit(t * n_head * head_dim)?;
2493        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2494        q = qn;
2495        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2496        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2497        k = kn;
2498        let rope_dims = geometry.n_rot as usize;
2499        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
2500        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
2501
2502        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
2503        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
2504        {
2505            let kvl = cache.kv[il].as_mut().unwrap();
2506            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
2507            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
2508                                       kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
2509                                       crate::Engine::kv_fp8_on())?;
2510            kvl.len += t;
2511            let new_len = kvl.len as i32;
2512            e.set_i32_one(&mut kvl.len_d, new_len)?;
2513        }
2514
2515        let base_len = {
2516            let kvl = cache.kv[il].as_ref().unwrap();
2517            kvl.len - t   // KV rows present BEFORE this chunk's append above
2518        };
2519        Ok((AttnPre { q, k, v, gate }, base_len))
2520    }
2521
2522    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
2523    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
2524    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
2525    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
2526    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
2527    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
2528    #[allow(clippy::too_many_arguments)]
2529    fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
2530                            v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
2531                            t: usize, cache: &mut Cache, il: usize,
2532                            head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
2533                            -> Result<(), Box<dyn std::error::Error>> {
2534        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
2535        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
2536        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
2537        // attend through the quantized cache exactly like every later chunk (quantize-then-
2538        // attend). One numeric class for every row => the chunk size cannot decide where a
2539        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
2540        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
2541        // pin-the-boundary approach).
2542        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
2543        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
2544        // with the fix unconditional, only re-introducing the class edge can prove the gate
2545        // still detects the mechanism. Never on in a measured default run.
2546        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
2547            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2548                e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2549            } else {
2550                e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2551            }
2552            return Ok(());
2553        }
2554        let kvl = cache.kv[il].as_ref().unwrap();
2555        let t_kv = base_len + t;
2556        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
2557        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
2558        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
2559        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
2560        // same numeric class, so the uniform contract holds on the fallback too.
2561        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2562            e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
2563                                        n_head_kv, t, t_kv, scale, true,
2564                                        kvl.k_tok_bytes, kvl.v_tok_bytes)?;
2565            return Ok(());
2566        }
2567        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
2568        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
2569        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
2570        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
2571        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
2572        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
2573        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
2574        let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
2575        if deqw {
2576            e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2577                                 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2578                                 crate::Engine::kv_fp8_on())?;
2579        } else {
2580            e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
2581                              t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
2582                              crate::Engine::kv_fp8_on())?;
2583        }
2584        Ok(())
2585    }
2586
2587    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
2588    /// (bit-identical composition) and hands wo its fp16 operand directly.
2589    fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
2590                            gate: &Option<CudaSlice<f32>>, t: usize,
2591                            n_head: usize, head_dim: usize)
2592                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2593        let (attn_g, ag16) = match gate {
2594            Some(gate) => {
2595                let n = t * n_head * head_dim;
2596                let mut ag = e.uninit(n)?;
2597                if Self::f16out_on(e, t) {
2598                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
2599                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
2600                    (ag, Some(a16))
2601                } else {
2602                    let mut gsig = e.uninit(n)?;
2603                    e.sigmoid(gate, &mut gsig, n)?;
2604                    e.mul(&attn, &gsig, &mut ag, n)?;
2605                    (ag, None)
2606                }
2607            }
2608            None => (attn, None),
2609        };
2610        Ok((attn_g, ag16))
2611    }
2612
2613    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
2614    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
2615    /// carried THROUGH the cache like the spec verify does: carried-ring conv
2616    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
2617    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
2618    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
2619    fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
2620                         hx: Option<&CudaSlice<u8>>, t: usize,
2621                         cache: &mut Cache, il: usize)
2622                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2623        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
2624        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
2625        let g4 = match hx {
2626            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
2627            None => e.matmul_group(&ws, h, t)?,
2628        };
2629        self.linear_attn_prime_core(e, la, g4, t, cache, il)
2630    }
2631
2632    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
2633    fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2634                              t: usize, cache: &mut Cache, il: usize)
2635                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2636        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
2637    }
2638
2639    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
2640    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
2641    /// conv ring writes back from the true tail. None = classic path, byte-identical.
2642    #[allow(clippy::too_many_arguments)]
2643    fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
2644                              t: usize, cache: &mut Cache, il: usize,
2645                              pad_len: Option<&CudaSlice<i32>>)
2646                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2647        // shim over the view twin (task #16): full-range views of the owned buffers.
2648        let ssm = self.cfg.ssm.as_ref().unwrap();
2649        let d_state = ssm.state_size as usize;
2650        let num_k = ssm.group_count as usize;
2651        let num_v = ssm.time_step_rank as usize;
2652        let key_dim = d_state * num_k;
2653        let value_dim = d_state * num_v;
2654        let conv_dim = key_dim * 2 + value_dim;
2655        let alpha = g4.pop().unwrap();                   // [T, num_v]
2656        let beta_raw = g4.pop().unwrap();                // [T, num_v]
2657        let z = g4.pop().unwrap();                       // [T, value_dim]
2658        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
2659        self.linear_attn_prime_core_pad_view(
2660            e, la,
2661            &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
2662            &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
2663            t, cache, il, pad_len)
2664    }
2665
2666    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
2667    /// shared verbatim by the per-seq scan path and the varlen batched path.
2668    #[allow(clippy::too_many_arguments)]
2669    fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
2670                            qkv_mixed: &cudarc::driver::CudaView<f32>,
2671                            beta_raw: &cudarc::driver::CudaView<f32>,
2672                            alpha: &cudarc::driver::CudaView<f32>,
2673                            t: usize, cache: &mut Cache, il: usize,
2674                            pad_len: Option<&CudaSlice<i32>>)
2675                            -> Result<GdnPrep, Box<dyn std::error::Error>> {
2676        let cfg = &self.cfg;
2677        let ssm = cfg.ssm.as_ref().unwrap();
2678        let d_state = ssm.state_size as usize;       // 128
2679        let num_k = ssm.group_count as usize;        // 16
2680        let num_v = ssm.time_step_rank as usize;     // 32
2681        let d_conv = ssm.conv_kernel as usize;       // 4
2682        let key_dim = d_state * num_k;               // 2048
2683        let value_dim = d_state * num_v;             // 4096
2684        let conv_dim = key_dim * 2 + value_dim;      // 8192
2685        let eps = cfg.rms_eps;
2686        debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
2687
2688        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
2689        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
2690        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
2691        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
2692        let rl = cache.recur[il].as_mut().unwrap();
2693        let hk = Self::gdn_hk(e, t, num_v, num_k);
2694        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
2695        let hk = if conv_fuse { hk } else { num_v };   // de-broadcast rides the fused conv
2696        let mut q_g = e.uninit(d_state * hk * t)?;
2697        let mut k_g = e.uninit(d_state * hk * t)?;
2698        let mut v_g = e.uninit(d_state * num_v * t)?;
2699        if conv_fuse {
2700            e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
2701                                  &mut q_g, &mut k_g, &mut v_g,
2702                                  conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
2703        } else {
2704            let mut conv_out = e.uninit(conv_dim * t)?;      // [conv_dim, T] channel-major, SiLU
2705            e.ssm_conv1d_tm_state_pad_v(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
2706                                  &mut conv_out, conv_dim, t, d_conv, pad_len)?;
2707            e.qkv_to_gdn_repack(&conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t)?;
2708        }
2709        let mut q_l2 = e.uninit(d_state * hk * t)?;
2710        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
2711        // Emitted only where a consumer exists (the wgmma config) — on other arches the
2712        // alloc + epilogue stores would be pure waste.
2713        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
2714            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2715            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
2716            Some(qb)
2717        } else {
2718            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
2719            None
2720        };
2721        let mut k_l2 = e.uninit(d_state * hk * t)?;
2722        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
2723        let kb16 = if Engine::l2_v2_on(d_state) {
2724            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
2725            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
2726            Some(kb)
2727        } else {
2728            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
2729            None
2730        };
2731        let mut beta = e.uninit(t * num_v)?;
2732        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
2733        let mut g_log = e.uninit(t * num_v)?;
2734        e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
2735        if let Some(len_d) = pad_len {
2736            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
2737        }
2738        Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
2739    }
2740
2741    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
2742    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
2743    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
2744    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
2745    #[allow(clippy::too_many_arguments)]
2746    fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
2747                                    g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
2748                                    caches: &mut [&mut Cache], il: usize)
2749                                    -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
2750        let ssm = self.cfg.ssm.as_ref().unwrap();
2751        let d_state = ssm.state_size as usize;
2752        let num_k = ssm.group_count as usize;
2753        let num_v = ssm.time_step_rank as usize;
2754        let key_dim = d_state * num_k;
2755        let value_dim = d_state * num_v;
2756        let conv_dim = key_dim * 2 + value_dim;
2757        let eps = self.cfg.rms_eps;
2758        let scale = 1.0 / (d_state as f32).sqrt();
2759        let b = ts.len();
2760        let c = Engine::gdn_chunk_size();
2761        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
2762        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
2763        let carried = caches.iter().any(|c| c.pos > 0);
2764        let use_vl = !carried
2765            && (2..=8).contains(&b)
2766            && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
2767            && e.gdn_mma_enabled(c)
2768            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
2769        if !use_vl {
2770            return (0..b).map(|s| {
2771                let (o, t) = (offs[s], ts[s]);
2772                self.linear_attn_prime_core_pad_view(
2773                    e, la,
2774                    &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
2775                    &g4[1].slice(o * value_dim..(o + t) * value_dim),
2776                    &g4[2].slice(o * num_v..(o + t) * num_v),
2777                    &g4[3].slice(o * num_v..(o + t) * num_v),
2778                    t, caches[s], il, None)
2779            }).collect();
2780        }
2781        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
2782        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
2783        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
2784        struct SeqBufs {
2785            conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
2786            q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
2787            gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
2788        }
2789        let d_conv = ssm.conv_kernel as usize;
2790        let f16o = Self::f16out_on(e, 16);
2791        let hk = Self::gdn_hk(e, 16, num_v, num_k);   // vl path is always chunked+mma
2792        let mut sb = Vec::with_capacity(b);
2793        let mut pres = Vec::with_capacity(b);
2794        for &t in ts.iter().take(b) {
2795            sb.push(SeqBufs {
2796                conv_out: e.uninit(conv_dim * t)?,
2797                q_g: e.uninit(d_state * hk * t)?,
2798                k_g: e.uninit(d_state * hk * t)?,
2799                v_g: e.uninit(d_state * num_v * t)?,
2800                q_l2: e.uninit(d_state * hk * t)?,
2801                k_l2: e.uninit(d_state * hk * t)?,
2802                beta: e.uninit(t * num_v)?,
2803                g_log: e.uninit(t * num_v)?,
2804                gn: e.uninit(d_state * num_v * t)?,
2805                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
2806            });
2807            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
2808        }
2809        let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
2810            let (o, t) = (offs[s], ts[s]);
2811            let rl = caches[s].recur[il].as_ref().unwrap();
2812            crate::GdnPrepVl {
2813                qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
2814                conv_state: e.addr_f32(&rl.conv_state),
2815                conv_out: e.addr_f32(&sb[s].conv_out),
2816                q_g: e.addr_f32(&sb[s].q_g), k_g: e.addr_f32(&sb[s].k_g), v_g: e.addr_f32(&sb[s].v_g),
2817                q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
2818                beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
2819                alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
2820                beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
2821                o: e.addr_f32(&pres[s].o),
2822                z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
2823                gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
2824                kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
2825                qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
2826                t: t as i32, pad: 0,
2827            }
2828        }).collect();
2829        let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
2830            let rl = caches[s].recur[il].as_ref().unwrap();
2831            crate::GdnSeqVl {
2832                kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
2833                beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
2834                wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
2835                ssnap: e.addr_u8(&pres[s].ssnap16),
2836                state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
2837                q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
2838                o: e.addr_f32(&pres[s].o),
2839                k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
2840                g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
2841                w: e.addr_f32(&pres[s].w),
2842                t: ts[s] as i32, nc: pres[s].nc as i32,
2843            }
2844        }).collect();
2845        e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
2846                       la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
2847        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
2848        // both standalone mirror launches vanish on the default config.
2849        if !Engine::l2_v2_on(d_state) {
2850            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
2851        }
2852        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
2853        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
2854            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
2855            if !Engine::l2_v2_on(d_state) {
2856                for s in 0..b {
2857                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
2858                }
2859            }
2860            let mut wa = [crate::GdnWVl::default(); 8];
2861            for s in 0..b {
2862                wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
2863            }
2864            Some(crate::GdnWVl8(wa))
2865        } else { None };
2866        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
2867        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
2868        if f16o {
2869            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
2870        }
2871        // per-seq state swap (+ non-f16out tail fallback)
2872        let mut out = Vec::with_capacity(b);
2873        for (s, bufs) in sb.into_iter().enumerate() {
2874            let rl = caches[s].recur[il].as_mut().unwrap();
2875            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2876            let (o, t) = (offs[s], ts[s]);
2877            let SeqBufs { mut gn, gn16, .. } = bufs;
2878            if f16o {
2879                out.push((gn, Some(gn16)));
2880            } else {
2881                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
2882                e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
2883                                   d_state, num_v * t, eps)?;
2884                out.push((gn, None));
2885            }
2886        }
2887        Ok(out)
2888    }
2889
2890    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
2891    /// views of the CONCAT projection outputs directly (no per-seq split copies).
2892    /// Same kernels, same values, byte-identical to the Vec shim above.
2893    #[allow(clippy::too_many_arguments)]
2894    fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
2895                              qkv_mixed: &cudarc::driver::CudaView<f32>,
2896                              z: &cudarc::driver::CudaView<f32>,
2897                              beta_raw: &cudarc::driver::CudaView<f32>,
2898                              alpha: &cudarc::driver::CudaView<f32>,
2899                              t: usize, cache: &mut Cache, il: usize,
2900                              pad_len: Option<&CudaSlice<i32>>)
2901                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
2902        let cfg = &self.cfg;
2903        let ssm = cfg.ssm.as_ref().unwrap();
2904        let d_state = ssm.state_size as usize;       // 128
2905        let num_v = ssm.time_step_rank as usize;     // 32
2906        let eps = cfg.rms_eps;
2907        let scale = 1.0 / (d_state as f32).sqrt();
2908
2909        let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
2910
2911        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
2912        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
2913        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
2914        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
2915        // verify keep the sequential kernel).
2916        let mut o = e.uninit(d_state * num_v * t)?;
2917        let rl = cache.recur[il].as_mut().unwrap();
2918        {
2919            let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
2920            e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
2921                               prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
2922                               prep.hk)?;
2923        }
2924        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2925
2926        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
2927        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
2928        let mut gn = e.uninit(d_state * num_v * t)?;
2929        let gn16 = if Self::f16out_on(e, t) {
2930            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
2931            e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
2932                                      d_state, num_v * t, eps)?;
2933            Some(g16)
2934        } else {
2935            e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
2936            None
2937        };
2938        Ok((gn, gn16))
2939    }
2940
2941    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
2942    #[allow(clippy::too_many_arguments)]
2943    fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
2944                              t: usize, cache: &mut Cache, il: usize,
2945                              pad_len: Option<&CudaSlice<i32>>)
2946                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2947        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
2948        if let Some(xh) = &gn16 {
2949            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
2950                return Ok(y);
2951            }
2952        }
2953        Ok(e.matmul(&la.ssm_out, &gn, t)?)
2954    }
2955
2956    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
2957    ///
2958    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
2959    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
2960    pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize)
2961                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2962        if self.cfg.step35.is_some() {
2963            return self.step35_attn(e, fa, h, pos_d, t, il);
2964        }
2965        let cfg = &self.cfg;
2966        let _n_embd = cfg.n_embd as usize;
2967        let geometry = cfg.full_attention_geometry_at(il as u32);
2968        let n_head = geometry.n_head as usize;
2969        let n_head_kv = geometry.n_head_kv as usize;
2970        let head_dim = geometry.head_dim_k as usize;
2971        let eps = cfg.rms_eps;
2972        let scale = geometry.attention_scale();
2973
2974        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
2975        // gate — wq out = n_head*head_dim, no split (see prime-path note).
2976        let gated = geometry.attention_gate
2977            == memra_gguf::config::AttentionGateKind::FusedQ;
2978        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
2979        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
2980        let v = g3.pop().unwrap();
2981        let mut k = g3.pop().unwrap();
2982        let qf = g3.pop().unwrap();
2983        let (mut q, gate) = if gated {
2984            let mut q = e.uninit(t * n_head * head_dim)?;
2985            let mut gate = e.uninit(t * n_head * head_dim)?;
2986            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2987            (q, Some(gate))
2988        } else {
2989            (qf, None)
2990        };
2991
2992        // QK-norm (per head_dim row), then partial RoPE.
2993        let mut qn = e.uninit(t * n_head * head_dim)?;
2994        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2995        q = qn;
2996        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2997        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2998        k = kn;
2999        let rope_dims = geometry.n_rot as usize;
3000        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, geometry.rope_base, 1.0)?;
3001        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, geometry.rope_base, 1.0)?;
3002
3003        // SDPA
3004        let mut attn = e.uninit(t * n_head * head_dim)?;
3005        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
3006        // falls back to naive sdpa.
3007        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
3008            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
3009            e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
3010        } else {
3011            e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
3012        }
3013
3014        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
3015        let attn_g = match &gate {
3016            Some(gate) => {
3017                let mut gsig = e.uninit(t * n_head * head_dim)?;
3018                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
3019                let mut ag = e.uninit(t * n_head * head_dim)?;
3020                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
3021                ag
3022            }
3023            None => attn,
3024        };
3025
3026        // o projection
3027        let o = e.matmul(&fa.wo, &attn_g, t)?;
3028        Ok(o)
3029    }
3030
3031    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
3032    pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
3033                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3034        let cfg = &self.cfg;
3035        let _n_embd = cfg.n_embd as usize;
3036        let ssm = cfg.ssm.as_ref().unwrap();
3037        let d_state = ssm.state_size as usize;       // 128
3038        let num_k = ssm.group_count as usize;        // 16
3039        let num_v = ssm.time_step_rank as usize;     // 32
3040        let d_conv = ssm.conv_kernel as usize;       // 4
3041        let head_k = d_state; let head_v = d_state;
3042        let key_dim = head_k * num_k;                // 2048
3043        let value_dim = head_v * num_v;              // 4096
3044        let conv_dim = key_dim * 2 + value_dim;      // 8192
3045        let eps = cfg.rms_eps;
3046        let scale = 1.0 / (d_state as f32).sqrt();
3047
3048        // projections
3049        // grouped: one f16 activation convert feeds all four projections (matmul_group)
3050        let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
3051        let alpha = g4.pop().unwrap();                   // [T, num_v]
3052        let beta_raw = g4.pop().unwrap();                // [T, num_v]
3053        let z = g4.pop().unwrap();                       // [T, value_dim]
3054        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
3055
3056        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
3057        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
3058        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
3059        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
3060        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
3061        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
3062        let _ = (head_k, head_v);
3063        let mut q_g = e.uninit(d_state * num_v * t)?;
3064        let mut k_g = e.uninit(d_state * num_v * t)?;
3065        let mut v_g = e.uninit(d_state * num_v * t)?;
3066        e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
3067                         conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
3068        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
3069        let mut q_l2 = e.uninit(d_state * num_v * t)?;
3070        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
3071        let mut k_l2 = e.uninit(d_state * num_v * t)?;
3072        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
3073        let v_gd = v_g;
3074
3075        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
3076        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
3077        let mut beta = e.uninit(t * num_v)?;
3078        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
3079        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
3080        let mut g_log = e.uninit(t * num_v)?;
3081        e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
3082
3083        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
3084        let state_in = e.zeros(d_state * d_state * num_v)?;  // zero state (prefill)
3085        let mut state_out = e.zeros(d_state * d_state * num_v)?;
3086        let mut o = e.uninit(d_state * num_v * t)?;
3087        e.gdn_scan_prefill(&q_l2, &k_l2, &v_gd, &g_log, &beta, None, None, &state_in, &mut state_out, &mut o, num_v, t, scale, num_v)?;
3088
3089        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
3090        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
3091        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
3092        // o rows are (t*num_v+vh) too. Good.
3093        let mut gn = e.uninit(d_state * num_v * t)?;
3094        e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
3095
3096        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
3097        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
3098        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
3099        let out = e.matmul(&la.ssm_out, &gn, t)?;
3100        Ok(out)
3101    }
3102}
3103
3104impl HybridModel {
3105    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
3106    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
3107    ///
3108    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
3109    /// different 860160-byte block than the same expert of layer 7).
3110    ///
3111    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
3112    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
3113    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
3114    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
3115    pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
3116               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3117        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), false)
3118    }
3119
3120    /// Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers
3121    /// keep `moe_ffn_il` and therefore retain their existing dispatch class.
3122    pub fn moe_ffn_il_prefill(
3123        &self,
3124        e: &Engine,
3125        m: &MoeWeights,
3126        z: &CudaSlice<f32>,
3127        t: usize,
3128        il: u16,
3129    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3130        Self::moe_ffn_inner(e, m, z, None, t, &self.cfg, il, self.max_moe_block(), true)
3131    }
3132
3133    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
3134    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
3135    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
3136    pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3137                          zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
3138               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3139        Self::moe_ffn_inner(
3140            e, m, z, zq8, t, &self.cfg, il, self.max_moe_block(), false,
3141        )
3142    }
3143
3144    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
3145    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
3146    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
3147    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
3148    ///
3149    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
3150    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
3151    pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3152                          cfg: &ModelConfig, il: u16, max_block: usize)
3153               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3154        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block, false)
3155    }
3156
3157    #[allow(clippy::too_many_arguments)]
3158    pub(crate) fn moe_ffn_inner(
3159        e: &Engine,
3160        m: &MoeWeights,
3161        z: &CudaSlice<f32>,
3162        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3163        t: usize,
3164        cfg: &ModelConfig,
3165        il: u16,
3166        max_block: usize,
3167        prefill: bool,
3168    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3169        let worker_io = crate::spill_pread::worker_enabled();
3170        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
3171        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
3172            e.with_moe_cache(max_block, |cache, _| {
3173                cache.begin_forward_epoch(il, t);
3174                if worker_io {
3175                    cache.begin_worker_scope();
3176                }
3177                Ok(())
3178            })?;
3179        }
3180        if Self::sigmoid_resident_dev_eligible(e, m, cfg) {
3181            let moe = cfg.moe.as_ref().unwrap();
3182            let n_expert = moe.expert_count as usize;
3183            let n_used = moe.expert_used_count as usize;
3184            let sigmoid = cfg.sigmoid_router().unwrap();
3185            let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
3186            Self::trace_sigmoid_router_logits(
3187                e, il, t, n_expert, n_used, &logits, m, sigmoid,
3188            )?;
3189            return Self::moe_ffn_sigmoid_dev(
3190                e, m, z, zq8, &logits, t, cfg, il, sigmoid,
3191            );
3192        }
3193        // Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED explicitly opts the
3194        // current caller into this research arm; the naked default stays on the established path.
3195        if t > 1 && moe_grouped_enabled(cfg, prefill) {
3196            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
3197            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
3198            // The grouped q8 path uses the same row-wise quantize/dot/activation/FMA programs as
3199            // sequential dispatch; mixed or q8-disabled layouts fall back to the matching f32
3200            // path. A mismatch is therefore a correctness failure, not an accepted numeric class.
3201            if std::env::var("MEMRA_MOE_GATE").is_ok() {
3202                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
3203                let g_host = e.dtoh(&grouped_out)?;
3204                let s_host = e.dtoh(&seq_out)?;
3205                let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
3206                let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
3207                if g_bytes == s_bytes {
3208                    println!("moe-gate il={il} t={t} BYTE-IDENTICAL");
3209                } else {
3210                    let diffs = g_host.iter().zip(s_host.iter()).enumerate()
3211                        .filter(|(_, (a, b))| a != b).count();
3212                    let maxdiff = g_host.iter().zip(s_host.iter())
3213                        .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
3214                    panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
3215                }
3216            }
3217            return Ok(grouped_out);
3218        }
3219        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
3220    }
3221
3222    fn sigmoid_resident_dev_eligible(
3223        e: &Engine,
3224        m: &MoeWeights,
3225        cfg: &ModelConfig,
3226    ) -> bool {
3227        let Some(moe) = cfg.moe.as_ref() else { return false };
3228        // Cached once per process: this predicate runs per MoE layer per decode step, and five
3229        // uncached env reads were ~300+ environ scans/step on the exact path this arm de-hosts.
3230        static OBSERVATION_MODE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3231        let observation_mode = *OBSERVATION_MODE.get_or_init(|| {
3232            std::env::var("MEMRA_MOE_STATS").is_ok()
3233                || std::env::var("MEMRA_MOE_TRACE").is_ok()
3234                || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
3235                || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok()
3236                || std::env::var("MEMRA_MOE_GATE").is_ok()
3237        });
3238        cfg.step35.is_some()
3239            && sigmoid_router_enabled()
3240            && moe_dev_enabled()
3241            && moe_slab_enabled()
3242            && !observation_mode
3243            && moe.expert_used_count <= 8
3244            && m.has_uniform_expert_layout()
3245            && m.gate_exps.macros.is_none()
3246            && m.up_exps.macros.is_none()
3247            && m.down_exps.macros.is_none()
3248            && !m.has_macros
3249            && moe_q8_enabled()
3250            && q8_expert_supported(m.gate_exps.qtype)
3251            && q8_expert_supported(m.up_exps.qtype)
3252            && q8_expert_supported(m.down_exps.qtype)
3253            && m.dev_exps
3254                .as_ref()
3255                .is_some_and(|dev| dev.dev == e.ctx().ordinal())
3256    }
3257
3258    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
3259    pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
3260                          cfg: &ModelConfig, il: u16, max_block: usize)
3261               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3262        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
3263    }
3264
3265    /// One router-logit selector shared by sequential and grouped dispatch. Real prefill must use
3266    /// the row-wise GEMV when exact routing is enabled: cuBLASLt's reduction changes with `m`,
3267    /// which makes expert selection depend on the caller's chunk or concat-batch shape.
3268    fn moe_router_logits(
3269        e: &Engine,
3270        m: &MoeWeights,
3271        z: &CudaSlice<f32>,
3272        t: usize,
3273        cfg: &ModelConfig,
3274    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3275        if t < PRIME_MIN_T {
3276            // Decode and speculative verify use one fixed per-row reduction program.
3277            if crate::router_kernel_on() {
3278                e.router_gemv(
3279                    m.gate_inp.float_data(),
3280                    z,
3281                    cfg.n_embd as usize,
3282                    m.gate_exps.n_expert,
3283                    t,
3284                )
3285            } else {
3286                e.matmul_decode_exact(&m.gate_inp, z, t)
3287            }
3288        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
3289            e.router_gemv(
3290                m.gate_inp.float_data(),
3291                z,
3292                cfg.n_embd as usize,
3293                m.gate_exps.n_expert,
3294                t,
3295            )
3296        } else {
3297            e.matmul(&m.gate_inp, z, t)
3298        }
3299    }
3300
3301    /// Append the host-visible router selection for one layer/forward when calibration tracing is
3302    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
3303    /// trace is independent of the dispatch optimization selected for the forward.
3304    fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
3305                        -> Result<(), Box<dyn std::error::Error>> {
3306        use std::io::Write as _;
3307        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
3308            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3309            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
3310            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
3311        }
3312        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
3313            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
3314            let pairs: Vec<String> = sel_all.iter().zip(weights)
3315                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
3316                .collect();
3317            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
3318        }
3319        Ok(())
3320    }
3321
3322    #[allow(clippy::too_many_arguments)]
3323    fn trace_sigmoid_router_logits(
3324        e: &Engine,
3325        il: u16,
3326        t: usize,
3327        n_expert: usize,
3328        n_used: usize,
3329        logits: &CudaSlice<f32>,
3330        m: &MoeWeights,
3331        (scaling_factor, route_norm): (f32, bool),
3332    ) -> Result<(), Box<dyn std::error::Error>> {
3333        if !crate::sigrouter_contract::served_logit_trace_enabled() || t != 1 {
3334            return Ok(());
3335        }
3336        let logits = e.dtoh(logits)?;
3337        let active: Vec<u8> = m
3338            .active_experts
3339            .as_ref()
3340            .map(|mask| mask.iter().map(|&enabled| u8::from(enabled)).collect())
3341            .unwrap_or_else(|| vec![1; n_expert]);
3342        let bias = m.exp_probs_b.clone().unwrap_or_else(|| vec![0.0; n_expert]);
3343        crate::sigrouter_contract::capture_served_logits(
3344            il as u32,
3345            t,
3346            n_expert,
3347            n_used,
3348            scaling_factor,
3349            route_norm,
3350            &active,
3351            &bias,
3352            &logits,
3353        )?;
3354        Ok(())
3355    }
3356
3357    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
3358    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
3359    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
3360    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
3361    fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
3362                       -> Result<(), Box<dyn std::error::Error>> {
3363        use std::io::Write as _;
3364        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
3365        let host = e.dtoh(z)?;
3366        if host.len() != t * n_embd {
3367            return Err(format!(
3368                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
3369                host.len(), t, n_embd
3370            ).into());
3371        }
3372        let bytes = unsafe {
3373            std::slice::from_raw_parts(
3374                host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
3375            )
3376        };
3377        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
3378        let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
3379        if state.is_none() {
3380            let dir = std::path::PathBuf::from(&dir);
3381            std::fs::create_dir_all(&dir)?;
3382            let index = std::fs::OpenOptions::new().create(true).append(true)
3383                .open(dir.join("index.jsonl"))?;
3384            *state = Some(MoeInputTraceWriter {
3385                dir,
3386                index,
3387                payloads: std::collections::HashMap::new(),
3388            });
3389        }
3390        let writer = state.as_mut().unwrap();
3391        if writer.dir != std::path::Path::new(&dir) {
3392            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
3393        }
3394        let file_name = format!("layer-{il:03}.f32");
3395        if !writer.payloads.contains_key(&il) {
3396            let payload = std::fs::OpenOptions::new().create(true).append(true)
3397                .open(writer.dir.join(&file_name))?;
3398            let offset = payload.metadata()?.len();
3399            writer.payloads.insert(il, (payload, offset));
3400        }
3401        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
3402        let row_offset = *offset;
3403        payload.write_all(bytes)?;
3404        *offset += bytes.len() as u64;
3405        writeln!(
3406            writer.index,
3407            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
3408             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
3409             \"payload_bytes\":{}}}",
3410            bytes.len()
3411        )?;
3412        Ok(())
3413    }
3414
3415    #[allow(clippy::too_many_arguments)]
3416    pub(crate) fn moe_ffn_sequential_zq8(
3417        e: &Engine,
3418        m: &MoeWeights,
3419        z: &CudaSlice<f32>,
3420        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
3421        t: usize,
3422        cfg: &ModelConfig,
3423        il: u16,
3424        max_block: usize,
3425    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3426        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3427        let moe = cfg.moe.as_ref().unwrap();
3428        let n_embd = cfg.n_embd as usize;          // 2048 (gate/up in_f, down out_f)
3429        let n_expert = moe.expert_count as usize;  // 256
3430        let n_used = moe.expert_used_count as usize; // 8
3431        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
3432
3433        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
3434        debug_assert_eq!(m.gate_exps.in_f, n_embd);
3435        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
3436        debug_assert_eq!(m.down_exps.in_f, n_ff_exp);  // down is TRANSPOSED: in=512
3437        debug_assert_eq!(m.down_exps.out_f, n_embd);   //                     out=2048
3438        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
3439
3440        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
3441        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
3442        let lim_exp = cfg.clamp_exp_at(il as u32);
3443        let lim_shexp = cfg.clamp_shexp_at(il as u32);
3444        let use_cache = Engine::moe_cache_enabled();
3445        let uniform_experts = m.has_uniform_expert_layout();
3446        let moe_q8 = uniform_experts && moe_q8_enabled()
3447            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3448            && q8_expert_supported(m.down_exps.qtype);
3449        // Experimental secondary backend: complete experts already resident in the SLRU stay on
3450        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
3451        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
3452        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
3453        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
3454        // commands and CI have no llama.cpp or OpenMP dependency.
3455        let cpu_expert_requested = crate::cpu_experts::configured();
3456        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
3457            return Err(std::io::Error::other(
3458                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
3459            )
3460            .into());
3461        }
3462        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
3463        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
3464        // Those backends are each deterministic but are different numeric configurations, so a
3465        // later prefill eviction can change greedy output. Freeze after the first real prefill;
3466        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
3467        // staging below and cannot change backend assignment.
3468        let freeze_cpu_residency = cpu_expert_requested
3469            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
3470        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
3471            .ok()
3472            .and_then(|value| value.parse::<usize>().ok())
3473            .is_some_and(|tokens| tokens > 0);
3474        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
3475            e.freeze_moe_cache();
3476        }
3477        let cache_frozen = use_cache && e.moe_cache_frozen();
3478        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
3479
3480        // 1. ROUTER: one selector is shared with expert-grouped prefill so changing dispatch
3481        // cannot change logits, selected expert ids, or routing weights.
3482        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
3483        if let Some(sig) = cfg.sigmoid_router() {
3484            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
3485        }
3486
3487        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
3488        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
3489        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
3490        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
3491        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
3492        // per-token host stall that dominated the 35B decode wall after stages 1+2.
3493        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
3494        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
3495        // only difference is where sel/w/pointers are READ from (device instead of params).
3496        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
3497        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
3498        // Any non-resident layer falls through to host routing + the gdec/sequential path.
3499        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
3500        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
3501        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
3502        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
3503        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
3504        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
3505        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
3506        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
3507        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
3508        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
3509        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
3510        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
3511        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
3512        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
3513        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
3514        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
3515        // now rides the dev loop below (same kernels per token as decode); pairs serves real
3516        // prefill (t >= 16, where spec never verifies).
3517        // sigmoid-router archs must NOT enter the pairs/dev arms: those dispatch arms still route
3518        // via the fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the
3519        // M3 gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Sigmoid routing has its own device
3520        // top-k kernel; MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle instead.
3521        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
3522        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
3523        // ride the macro-aware sequential/staged paths below or every expert output is off by
3524        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
3525        let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3526            && m.down_exps.macros.is_none();
3527        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
3528        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
3529        // so it cannot even see the per-layer limit.
3530        if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
3531            && !cfg.swiglu_clamped_at(il as u32)
3532            && no_exp_macros
3533            && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
3534            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3535            && q8_expert_supported(m.down_exps.qtype)
3536            && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
3537            && std::env::var("MEMRA_MOE_STATS").is_err() {
3538            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
3539        }
3540
3541        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
3542        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
3543        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
3544        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk). Sigmoid
3545        // routing has a separate device top-k kernel, but this dev arm cannot consume its sel/w
3546        // contract, so sigmoid-router arches must NOT enter it: with MOE_CACHE=1 M3 silently
3547        // routed softmax = wrong experts (gate MISMATCH 74602 vs 92, caught 2026-07-07).
3548        // MEMRA_SIG_ROUTER=0 selects the host sigmoid oracle for the supported sigmoid path.
3549        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
3550        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
3551        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
3552        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
3553        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
3554        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
3555        // Keyed off sigmoid_router() so arch #4 is denied by construction.
3556        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
3557        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
3558        let dev_ok = uniform_experts && cfg.sigmoid_router().is_none()
3559            && cfg.m3.is_none() && cfg.hy3.is_none()
3560            && !cfg.swiglu_clamped_at(il as u32);
3561        // Observation modes must route through the host-visible selection below. Otherwise a fully
3562        // resident layer returns through device dispatch before its trace/stats row is recorded,
3563        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
3564        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
3565            || std::env::var("MEMRA_MOE_TRACE").is_ok()
3566            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
3567            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
3568        if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
3569            && !observe_routes {
3570            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3571        }
3572        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
3573            && !observe_routes {
3574            let row_ok = e.with_moe_cache(max_block, |c, eng| {
3575                if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
3576                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
3577            })?;
3578            if row_ok {
3579                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
3580            }
3581        }
3582
3583        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
3584        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
3585            if cpu_hybrid {
3586                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
3587                    e,
3588                    &logits,
3589                    z,
3590                    t,
3591                    n_expert,
3592                    n_used,
3593                    m.exp_probs_b.as_deref(),
3594                    sig,
3595                    m.active_experts.as_deref(),
3596                )?;
3597                (sel, w, Some(input))
3598            } else {
3599                let (sel, w) = Self::moe_route_sigmoid_cfg(
3600                    e, &logits, t, n_expert, n_used, m, sig,
3601                )?;
3602                (sel, w, None)
3603            }
3604        } else {
3605            let (sel, w) = Self::moe_route_cfg(
3606                e, &logits, t, n_expert, n_used, m.active_experts.as_deref(),
3607            )?;
3608            (sel, w, None)
3609        };
3610        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
3611
3612        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
3613        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
3614        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
3615        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
3616        Self::trace_moe_input(e, il, t, n_embd, z)?;
3617
3618        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
3619        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
3620        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
3621        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
3622        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
3623        // wait for each pending block, so later copies can overlap the earlier expert kernels while
3624        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
3625        // T=1; batched forwards can have token-local consumers still in flight between selections.
3626        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
3627        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
3628        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
3629        let worker_disk_prefetch =
3630            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
3631        let promote_worker_h2d =
3632            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
3633        if promote_worker_h2d {
3634            let mut selected_blocks = Vec::with_capacity(n_used * 3);
3635            for &ex in sel_all.iter().take(n_used) {
3636                let ex = ex as u16;
3637                selected_blocks.extend([
3638                    BlockId::new(il, PROJ_GATE, ex),
3639                    BlockId::new(il, PROJ_UP, ex),
3640                    BlockId::new(il, PROJ_DOWN, ex),
3641                ]);
3642            }
3643            for &ex in sel_all.iter().take(n_used) {
3644                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
3645            }
3646            e.with_moe_cache(max_block, |cache, eng| {
3647                cache.promote_worker_reads_at_safe_boundary(
3648                    &selected_blocks,
3649                    &selected_blocks,
3650                    eng,
3651                )?;
3652                Ok(())
3653            })?;
3654        }
3655
3656        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
3657        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
3658        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
3659            let mut cnt = vec![0u32; n_expert];
3660            for &s in sel_all.iter() { cnt[s as usize] += 1; }
3661            let total = sel_all.len() as f64;
3662            let mut h = 0.0f64;
3663            let mut active = 0usize;
3664            for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
3665            let maxc = cnt.iter().copied().max().unwrap_or(0);
3666            println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
3667                     il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
3668        }
3669
3670        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
3671        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
3672        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
3673        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
3674        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
3675        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
3676        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
3677        // zeroed-then-accumulated exactly as before (fallback).
3678        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
3679        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
3680        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
3681        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
3682        let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled()
3683            && !cfg.swiglu_clamped_at(il as u32);
3684        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
3685        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
3686        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
3687        // archs the slabs were uploaded but never read, and every expert went through the
3688        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
3689        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
3690        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
3691        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
3692        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
3693        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
3694        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
3695        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
3696        // strictly worse than staging); under PP-2 without the prime walker this admits
3697        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
3698        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
3699        let slab_local = m.dev_exps.as_ref()
3700            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
3701        let slab_bases = slab_local.map(|d| {
3702            use cudarc::driver::DevicePtr;
3703            let s = e.stream();
3704            let (pg, _g0) = d.gate.device_ptr(&s);
3705            let (pu, _g1) = d.up.device_ptr(&s);
3706            let (pd, _g2) = d.down.device_ptr(&s);
3707            (pg as u64, pu as u64, pd as u64)
3708        });
3709        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
3710        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
3711        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
3712        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
3713        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
3714        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
3715        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
3716        // all-resident tokens, staged loop for misses), which is a dispatch-class
3717        // comparison, not a provenance one.
3718        let slab_fused_may_fire = slab_bases.is_some() && n_used <= 8 && gdec_enabled()
3719            && !cfg.swiglu_clamped_at(il as u32) && cfg.m3.is_none()
3720            && no_exp_macros && moe_q8;
3721        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
3722        // uninit; a token that falls through to any accumulating loop zeroes its own row.
3723        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
3724            e.uninit(t * n_embd)?
3725        } else {
3726            e.zeros(t * n_embd)?
3727        };
3728        // The router readback above already established a host boundary. Copy each small-t hidden
3729        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
3730        let cpu_input = if cpu_hybrid {
3731            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
3732        } else {
3733            None
3734        };
3735
3736        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
3737        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
3738        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
3739        // measured ~123 memsets/token of the decode wall).
3740        let g_len = m.gate_exps.max_expert_bytes();  // 860160 for the uniform 35B gate
3741        let u_len = m.up_exps.max_expert_bytes();    // 860160 for the uniform 35B up
3742        let d_len = m.down_exps.max_expert_bytes();  // 1114112 for the uniform 35B down
3743        let mut scratch_g: Option<CudaSlice<u8>> = None;
3744        let mut scratch_u: Option<CudaSlice<u8>> = None;
3745        let mut scratch_d: Option<CudaSlice<u8>> = None;
3746        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
3747        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
3748
3749        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
3750        // the copy stream before launching the current expert's compute. Pending slots stay invisible
3751        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
3752        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
3753        let page_window = moe_page_prefetch_window();
3754
3755        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
3756        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
3757        for tok in 0..t {
3758            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
3759            let w = &w_all[tok * n_used..(tok + 1) * n_used];
3760            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);  // CudaView<f32>
3761            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
3762
3763            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
3764            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
3765            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
3766            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
3767            // memcpy, zero admission, so no slot can move under the collected pointers) — any
3768            // miss falls through to the sequential loop below, which admits as before. In steady
3769            // state on a fully-resident rig every token-layer takes the grouped path.
3770            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
3771            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
3772            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
3773            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
3774            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
3775            // per-expert macro-scales the fused kernels don't fold — those fall through too.
3776            let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
3777                && m.down_exps.macros.is_none();
3778            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
3779            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
3780            // with pointers computed from the resident slab base + ex*stride instead of
3781            // collected SLRU slot addresses. No cache lock, no residency predicate — the
3782            // slab holds every expert by construction, so this arm never falls through
3783            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
3784            // staging both die). Bit-identity class: pointer provenance only, the same
3785            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
3786            // slab exists it is strictly better (no lock, no miss).
3787            if slab_fused_may_fire {
3788                let (pg, pu, pd) = slab_bases.unwrap();
3789                let mut gp = [0u64; 8];
3790                let mut up = [0u64; 8];
3791                let mut dp = [0u64; 8];
3792                for (j, &ex) in sel.iter().enumerate() {
3793                    let ex = ex as usize;
3794                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
3795                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
3796                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
3797                }
3798                let mut wv = [0f32; 8];
3799                wv[..n_used].copy_from_slice(w);
3800                if tok_q8.is_none() {
3801                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3802                }
3803                let (zq, zd) = tok_q8.as_ref().unwrap();
3804                let act = e.moe_gate_up_silu8_q8(crate::WPtr8(gp), crate::WPtr8(up), zq, zd,
3805                                                 n_embd, n_ff_exp, n_used,
3806                                                 m.gate_exps.qtype, m.up_exps.qtype,
3807                                                 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
3808                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
3809                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3810                e.moe_down8_fma_q8(crate::WPtr8(dp), crate::F32x8(wv), &aq2, &ad2, &mut dst,
3811                                   n_ff_exp, n_embd, n_used,
3812                                   m.down_exps.qtype, m.down_exps.row_bytes)?;
3813                continue;
3814            }
3815            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
3816                if tok_q8.is_none() {
3817                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3818                }
3819                let (zq, zd) = tok_q8.as_ref().unwrap();
3820                if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
3821                                           &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3822                    continue;
3823                }
3824            } else if gdec_may_fire && cfg.m3.is_none() && no_macros
3825                && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
3826                                        &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
3827                continue;
3828            }
3829
3830            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
3831            // slab pair could fire. This token fell through to a sequential axpy loop, which
3832            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
3833            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
3834            // has no fallible predicate), included for the allocation invariant's symmetry.
3835            if gdec_may_fire || slab_fused_may_fire {
3836                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3837                e.memset_zeros_view(&mut row)?;
3838            }
3839
3840            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
3841            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
3842            // stall this path exists to remove, while mixing projections would require another
3843            // activation round-trip. Weight addresses remain valid until this worker is joined at
3844            // the bottom of the token scope.
3845            let mut cpu_mask = vec![false; sel.len()];
3846            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
3847                let gpu_resident = if use_cache {
3848                    e.with_moe_cache(max_block, |cache, _| {
3849                        Ok(sel
3850                            .iter()
3851                            .map(|&expert| {
3852                                let expert = expert as u16;
3853                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
3854                                    .into_iter()
3855                                    .filter(|&projection| {
3856                                        cache
3857                                            .resident(BlockId::new(il, projection, expert))
3858                                            .is_some()
3859                                    })
3860                                    .count()
3861                            })
3862                            .collect::<Vec<_>>())
3863                    })?
3864                } else {
3865                    vec![0; sel.len()]
3866                };
3867                let mut cpu_selected = Vec::new();
3868                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
3869                    if gpu_resident[index] != 3 {
3870                        cpu_mask[index] = true;
3871                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
3872                        let expert = expert as usize;
3873                        cpu_selected.push((expert, route_weight));
3874                    }
3875                }
3876                if crate::cpu_experts::predictor_enabled() {
3877                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
3878                    // from this layer's MoE input and prefetches predicted-and-missing
3879                    // experts into the companion RAM cache. Never blocks this thread.
3880                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3881                    crate::cpu_experts::predictor_submit(il, row);
3882                }
3883                if cpu_selected.is_empty() {
3884                    None
3885                } else {
3886                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
3887                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
3888                        .map_err(std::io::Error::other)?;
3889                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
3890                }
3891            } else {
3892                None
3893            };
3894
3895            let worker_window = worker_disk_prefetch
3896                .then(worker_prefetch_window)
3897                .unwrap_or(0);
3898            for (j, &ex) in sel.iter().enumerate() {
3899                if cpu_mask[j] {
3900                    continue;
3901                }
3902                let ex = ex as usize;
3903                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
3904                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
3905                // fused form) and macro-carrying artifacts — still have their bytes in the
3906                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
3907                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
3908                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
3909                if let Some(d) = slab_local {
3910                    let gl = m.gate_exps.expert_layout(ex);
3911                    let ul = m.up_exps.expert_layout(ex);
3912                    let dl = m.down_exps.expert_layout(ex);
3913                    let (g0, u0, d0) = (ex * m.gate_exps.expert_stride,
3914                                        ex * m.up_exps.expert_stride,
3915                                        ex * m.down_exps.expert_stride);
3916                    let (gate, up) = if moe_q8 {
3917                        if tok_q8.is_none() {
3918                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3919                        }
3920                        let (zq, zd) = tok_q8.as_ref().unwrap();
3921                        (e.qmatvec_expert_q8(&d.gate, g0..g0 + gl.len, zq, zd, 1,
3922                                             m.gate_exps.in_f, m.gate_exps.out_f,
3923                                             gl.qtype, gl.row_bytes)?,
3924                         e.qmatvec_expert_q8(&d.up, u0..u0 + ul.len, zq, zd, 1,
3925                                             m.up_exps.in_f, m.up_exps.out_f,
3926                                             ul.qtype, ul.row_bytes)?)
3927                    } else {
3928                        (e.qmatvec_view(&d.gate, g0..g0 + gl.len, &zt, 1,
3929                                        m.gate_exps.in_f, m.gate_exps.out_f,
3930                                        gl.qtype, gl.row_bytes)?,
3931                         e.qmatvec_view(&d.up, u0..u0 + ul.len, &zt, 1,
3932                                        m.up_exps.in_f, m.up_exps.out_f,
3933                                        ul.qtype, ul.row_bytes)?)
3934                    };
3935                    let mut act = e.uninit(n_ff_exp)?;
3936                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3937                                      m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3938                    let y = if moe_q8 {
3939                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3940                        e.qmatvec_expert_q8(&d.down, d0..d0 + dl.len, &aq2, &ad2, 1,
3941                                            m.down_exps.in_f, m.down_exps.out_f,
3942                                            dl.qtype, dl.row_bytes)?
3943                    } else {
3944                        let actv = act.slice(0..n_ff_exp);
3945                        e.qmatvec_view(&d.down, d0..d0 + dl.len, &actv, 1,
3946                                       m.down_exps.in_f, m.down_exps.out_f,
3947                                       dl.qtype, dl.row_bytes)?
3948                    };
3949                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3950                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3951                    continue;
3952                }
3953                for next in page_prefetch_positions(j, sel.len(), page_window) {
3954                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
3955                }
3956                let keep = [
3957                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
3958                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
3959                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
3960                ];
3961                if worker_disk_prefetch && worker_window > 0 {
3962                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
3963                        Self::moe_prefetch_disk_expert(
3964                            e,
3965                            il,
3966                            sel[next] as usize,
3967                            m,
3968                            max_block,
3969                            &keep,
3970                        )?;
3971                    }
3972                } else if cache_dispatch
3973                    && !cpu_hybrid
3974                    && moe_prefetch_enabled()
3975                    && j + 1 < sel.len()
3976                {
3977                    let next = sel[j + 1] as usize;
3978                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
3979                }
3980                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
3981                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
3982                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
3983                    // layouts stay on the metadata-aware f32 path.
3984                    if (gate_q8 || up_q8) && tok_q8.is_none() {
3985                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
3986                    }
3987                    let gate = if gate_q8 {
3988                        let (zq, zd) = tok_q8.as_ref().unwrap();
3989                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
3990                    } else {
3991                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
3992                    };
3993                    let up = if up_q8 {
3994                        let (zq, zd) = tok_q8.as_ref().unwrap();
3995                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
3996                    } else {
3997                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
3998                    };
3999                    let mut act = e.uninit(n_ff_exp)?;
4000                    Self::ffn_act_lim(
4001                        e,
4002                        cfg,
4003                        &gate,
4004                        &up,
4005                        m.gate_exps.macro_scale(ex),
4006                        m.up_exps.macro_scale(ex),
4007                        lim_exp,
4008                        &mut act,
4009                        n_ff_exp,
4010                    )?;
4011                    let y = if down_q8 {
4012                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
4013                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
4014                    } else {
4015                        let actv = act.slice(0..n_ff_exp);
4016                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
4017                    };
4018                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4019                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
4020                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4021                } else if cache_dispatch {
4022                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
4023                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
4024                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
4025                    // only difference between HIT and MISS is whether the memcpy_htod ran.
4026                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
4027                    let up   = Self::moe_cached_gemm(e, il, PROJ_UP,   ex, m, max_block, &zt)?;
4028                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
4029                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4030                        m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
4031                    let actv = act.slice(0..n_ff_exp);
4032                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
4033                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4034                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
4035                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4036                } else if cache_frozen {
4037                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
4038                    // first prime. Reuse every fixed resident projection directly and stage only a
4039                    // true miss through the ordinary scratch slot. This preserves the established
4040                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
4041                    let gate = Self::moe_frozen_gemm(
4042                        e,
4043                        il,
4044                        PROJ_GATE,
4045                        ex,
4046                        m,
4047                        max_block,
4048                        &zt,
4049                        &mut scratch_g,
4050                        g_len,
4051                    )?;
4052                    let up = Self::moe_frozen_gemm(
4053                        e,
4054                        il,
4055                        PROJ_UP,
4056                        ex,
4057                        m,
4058                        max_block,
4059                        &zt,
4060                        &mut scratch_u,
4061                        u_len,
4062                    )?;
4063                    let mut act = e.uninit(n_ff_exp)?;
4064                    Self::ffn_act_lim(
4065                        e,
4066                        cfg,
4067                        &gate,
4068                        &up,
4069                        m.gate_exps.macro_scale(ex),
4070                        m.up_exps.macro_scale(ex),
4071                        lim_exp,
4072                        &mut act,
4073                        n_ff_exp,
4074                    )?;
4075                    let actv = act.slice(0..n_ff_exp);
4076                    let y = Self::moe_frozen_gemm(
4077                        e,
4078                        il,
4079                        PROJ_DOWN,
4080                        ex,
4081                        m,
4082                        max_block,
4083                        &actv,
4084                        &mut scratch_d,
4085                        d_len,
4086                    )?;
4087                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4088                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4089                } else {
4090                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
4091                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
4092                    // fully overwrites the byte range the GEMM reads).
4093                    if scratch_g.is_none() {
4094                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
4095                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
4096                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
4097                    }
4098                    let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
4099                                        scratch_d.as_mut().unwrap());
4100                    let gl = m.gate_exps.expert_layout(ex);
4101                    let ul = m.up_exps.expert_layout(ex);
4102                    let dl = m.down_exps.expert_layout(ex);
4103                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4104                    let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
4105                        m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
4106
4107                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4108                    let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
4109                        m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
4110
4111                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
4112                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4113                        m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
4114
4115                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4116                    let actv = act.slice(0..n_ff_exp);
4117                    let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
4118                        m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
4119
4120                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4121                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
4122                }
4123            }
4124            if let Some(worker) = cpu_worker {
4125                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
4126                let cpu_output = e.htod(&cpu_output)?;
4127                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4128                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4129            }
4130            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
4131                for (j, &ex) in sel.iter().enumerate() {
4132                    if cpu_mask[j] {
4133                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
4134                    }
4135                }
4136            }
4137        }
4138
4139        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
4140        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
4141        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
4142        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
4143        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4144            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4145        {
4146            let n_ff_sh = gate_shexp.out_features();  // 512
4147            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
4148            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
4149            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
4150            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
4151            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
4152            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
4153            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
4154            let verify_t = t > 1 && t < PRIME_MIN_T;
4155            let (sg_gate, sg_up) = if t == 1 {
4156                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
4157                    Some(pair) => pair,
4158                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
4159                }
4160            } else if verify_t {
4161                (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
4162            } else {
4163                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)   // [T, 512] each
4164            };
4165            let mut sa = e.uninit(t * n_ff_sh)?;  // activation fully overwrites
4166            Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp, &mut sa, t * n_ff_sh)?;
4167            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
4168                     else { e.matmul(down_shexp, &sa, t)? };     // [T, n_embd]
4169
4170            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
4171            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
4172            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
4173            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
4174            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
4175            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
4176            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
4177            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
4178            // expert's contribution into every token's residual, so under cross-request
4179            // concat prefill a session's hidden state depended on its co-arrivals' token
4180            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
4181            let g = match &m.gate_inp_shexp {
4182                Some(gate_inp_shexp) => {
4183                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4184                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4185                    } else {
4186                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4187                        let mut g = e.uninit(t)?;  // sigmoid fully overwrites
4188                        e.sigmoid(&gs, &mut g, t)?;
4189                        g
4190                    }
4191                }
4192                None => e.htod(&vec![1.0f32; t])?,
4193            };
4194            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
4195            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4196        }
4197
4198        Ok(moe_out)
4199    }
4200
4201    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
4202    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
4203    pub fn stage1_h2d_per_token(&self) -> u64 {
4204        use crate::hybrid::Ffn;
4205        let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
4206        let mut bytes = 0u64;
4207        for l in self.layers.iter() {
4208            if let Ffn::Moe(m) = &l.ffn {
4209                bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
4210                                   + m.down_exps.max_expert_bytes()) as u64;
4211            }
4212        }
4213        bytes
4214    }
4215
4216    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
4217    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
4218    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
4219    pub(crate) fn max_moe_block(&self) -> usize {
4220        use crate::hybrid::Ffn;
4221        let mut mx = 0usize;
4222        let mut scan = |ffn: &Ffn| {
4223            if let Ffn::Moe(m) = ffn {
4224                mx = mx.max(m.gate_exps.max_expert_bytes())
4225                       .max(m.up_exps.max_expert_bytes())
4226                       .max(m.down_exps.max_expert_bytes());
4227            }
4228        };
4229        for l in self.layers.iter() { scan(&l.ffn); }
4230        if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
4231        mx
4232    }
4233
4234    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
4235    /// but have no bytes and therefore consume no residency slot.
4236    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
4237        use crate::hybrid::Ffn;
4238        let mut sizes = Vec::new();
4239        let mut scan = |ffn: &Ffn| {
4240            let Ffn::Moe(m) = ffn else { return };
4241            for ex in 0..m.gate_exps.n_expert {
4242                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
4243                    continue;
4244                }
4245                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
4246                    let len = exps.expert_layout(ex).len;
4247                    if len > 0 {
4248                        sizes.push(len);
4249                    }
4250                }
4251            }
4252        };
4253        for layer in &self.layers {
4254            scan(&layer.ffn);
4255        }
4256        if let Some(mtp) = &self.mtp {
4257            scan(&mtp.ffn);
4258        }
4259        sizes
4260    }
4261
4262    /// Persist the frozen residency set so a later process can restage it directly and skip
4263    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
4264    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
4265    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
4266    /// post-freeze argmax gate still validates the serving assignment.
4267    pub fn save_cpu_expert_residency_profile(
4268        &self,
4269        e: &Engine,
4270        path: &std::path::Path,
4271    ) -> Result<(), Box<dyn std::error::Error>> {
4272        let Some(ids) = e.export_moe_residency() else {
4273            return Err("no MoE residency cache to persist".into());
4274        };
4275        let mut body = format!(
4276            "memra-freeze-profile v1 max_block={} blocks={}\n",
4277            self.max_moe_block(),
4278            ids.len()
4279        );
4280        for (layer, proj, ex) in &ids {
4281            body.push_str(&format!("{layer} {proj} {ex}\n"));
4282        }
4283        let tmp = path.with_extension("tmp");
4284        std::fs::write(&tmp, body)?;
4285        std::fs::rename(&tmp, path)?;
4286        println!(
4287            "[moe-cache] freeze profile saved: {} blocks -> {}",
4288            ids.len(),
4289            path.display()
4290        );
4291        Ok(())
4292    }
4293
4294    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
4295    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
4296    /// missing or its header does not match this model's slot geometry.
4297    pub fn restore_cpu_expert_residency_profile(
4298        &self,
4299        e: &Engine,
4300        path: &std::path::Path,
4301    ) -> Result<bool, Box<dyn std::error::Error>> {
4302        use crate::hybrid::Ffn;
4303        use crate::moe_cache::BlockId;
4304        let Ok(content) = std::fs::read_to_string(path) else {
4305            return Ok(false);
4306        };
4307        let mut lines = content.lines();
4308        let Some(header) = lines.next() else { return Ok(false) };
4309        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
4310        if !header.starts_with(&expected) {
4311            println!(
4312                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
4313                path.display()
4314            );
4315            return Ok(false);
4316        }
4317        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
4318            std::collections::HashMap::new();
4319        for line in lines {
4320            let mut fields = line.split_whitespace();
4321            let (Some(layer), Some(proj), Some(ex)) =
4322                (fields.next(), fields.next(), fields.next())
4323            else {
4324                continue;
4325            };
4326            let (Ok(layer), Ok(proj), Ok(ex)) =
4327                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
4328            else {
4329                continue;
4330            };
4331            by_layer
4332                .entry(layer)
4333                .or_default()
4334                .push(BlockId::new(layer, proj, ex));
4335        }
4336        let requested: usize = by_layer.values().map(Vec::len).sum();
4337        if requested == 0 {
4338            return Ok(false);
4339        }
4340        let max_block = self.max_moe_block();
4341        let mut restaged = 0usize;
4342        let mut stage_layer = |layer_index: u16,
4343                               ffn: &Ffn|
4344         -> Result<(), Box<dyn std::error::Error>> {
4345            let Ffn::Moe(m) = ffn else { return Ok(()) };
4346            let Some(ids) = by_layer.get(&layer_index) else {
4347                return Ok(());
4348            };
4349            e.with_moe_cache(max_block, |cache, eng| {
4350                for id in ids {
4351                    if cache.restage_block(*id, m, eng)? {
4352                        restaged += 1;
4353                    }
4354                }
4355                Ok(())
4356            })
4357        };
4358        for (index, layer) in self.layers.iter().enumerate() {
4359            stage_layer(index as u16, &layer.ffn)?;
4360        }
4361        if let Some(mtp) = self.mtp.as_ref() {
4362            stage_layer(u16::MAX, &mtp.ffn)?;
4363        }
4364        e.freeze_moe_cache();
4365        println!(
4366            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
4367            path.display()
4368        );
4369        Ok(true)
4370    }
4371
4372    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
4373    pub fn freeze_cpu_expert_residency(
4374        &self,
4375        e: &Engine,
4376    ) -> Result<(), Box<dyn std::error::Error>> {
4377        e.freeze_moe_cache();
4378        Ok(())
4379    }
4380
4381    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
4382    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
4383    /// the model's activation exactly.
4384    ///
4385    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
4386    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
4387    /// form for anything that can land on a clamped layer.
4388    pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4389               act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
4390        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
4391    }
4392
4393    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
4394    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
4395    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
4396    #[allow(clippy::too_many_arguments)]
4397    pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4398               gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
4399               -> Result<(), Box<dyn std::error::Error>> {
4400        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
4401    }
4402
4403    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
4404    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
4405    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
4406    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
4407    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
4408    ///                 arrays are SEPARATE and a layer can have one without the other.
4409    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
4410    /// already known live.
4411    #[allow(clippy::too_many_arguments)]
4412    pub(crate) fn ffn_act_lim(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
4413               gs: f32, us: f32, limit: Option<f32>, act: &mut CudaSlice<f32>, n: usize)
4414               -> Result<(), Box<dyn std::error::Error>> {
4415        if let Some(m3) = cfg.m3.as_ref() {
4416            debug_assert!(limit.is_none(), "m3 swigluoai and step35 clamp are different archs");
4417            return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
4418        }
4419        if let Some(l) = limit {
4420            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
4421        }
4422        if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
4423        e.silu_mul_scaled(gate, up, gs, us, act, n)
4424    }
4425
4426    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
4427    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
4428    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
4429    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
4430    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
4431    fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
4432                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4433        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None)
4434    }
4435
4436    /// DeepSeek-V3-class sigmoid routing (Step-3.7, MiniMax-M3, Hy3, GLM-DSA). Reference:
4437    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
4438    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
4439    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
4440    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
4441    /// Default uses one device kernel plus the same pinned sel/w readback contract as the fused
4442    /// softmax router. `MEMRA_SIG_ROUTER=0` restores the full-logit DtoH host oracle.
4443    #[allow(clippy::too_many_arguments)]
4444    fn moe_route_sigmoid_cfg(
4445        e: &Engine,
4446        logits: &CudaSlice<f32>,
4447        t: usize,
4448        n_expert: usize,
4449        n_used: usize,
4450        m: &MoeWeights,
4451        (sf, route_norm): (f32, bool),
4452    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4453        if sigmoid_router_enabled() {
4454            return e.moe_router_sigmoid_topk_host(
4455                logits,
4456                t,
4457                n_expert,
4458                n_used,
4459                m.active_count(),
4460                &m.exp_probs_b_dev,
4461                &m.active_experts_dev,
4462                sf,
4463                route_norm,
4464            );
4465        }
4466        let lg = e.dtoh(logits)?;
4467        Self::moe_route_sigmoid_host(
4468            &lg,
4469            t,
4470            n_expert,
4471            n_used,
4472            m.exp_probs_b.as_deref(),
4473            sf,
4474            route_norm,
4475            m.active_experts.as_deref(),
4476        )
4477    }
4478
4479    /// Softmax router for qwen35moe/OLMoE. An active-mask overlay retains the host oracle because
4480    /// the existing softmax device kernel has no mask input.
4481    fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
4482                     active: Option<&[bool]>)
4483                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4484        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
4485        // rollback) via the single-sync pinned readback — softmax arch only.
4486        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
4487            return e.moe_router_topk_host(logits, t, n_expert, n_used);
4488        }
4489        // Host oracle (the §D bit-identity reference).
4490        let lg = e.dtoh(logits)?;   // [T*n_expert] host
4491        let mut sel = vec![0u32; t * n_used];
4492        let mut w_out = vec![0f32; t * n_used];
4493        for tok in 0..t {
4494            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4495            // softmax over ALL n_expert (stable: subtract max)
4496            let maxl = row.iter().enumerate()
4497                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
4498                .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
4499            let mut probs = vec![0f32; n_expert];
4500            let mut den = 0f32;
4501            for i in 0..n_expert {
4502                if active.is_some_and(|mask| !mask[i]) { continue; }
4503                let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
4504            }
4505            for p in probs.iter_mut() { *p /= den; }
4506            // stable DESC sort: prob DESC, ascending-index tiebreak.
4507            let mut idx: Vec<usize> = (0..n_expert)
4508                .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
4509            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
4510            let sl = &idx[..n_used];
4511            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
4512            let mut ws: f32 = wv.iter().sum();
4513            ws = ws.max(6.103515625e-5_f32);  // F16 smallest normal, clamp BEFORE divide
4514            for x in wv.iter_mut() { *x /= ws; }
4515            for j in 0..n_used {
4516                sel[tok * n_used + j] = sl[j] as u32;
4517                w_out[tok * n_used + j] = wv[j];
4518            }
4519        }
4520        Ok((sel, w_out))
4521    }
4522
4523    #[allow(clippy::too_many_arguments)]
4524    fn moe_route_sigmoid_with_input(
4525        e: &Engine,
4526        logits: &CudaSlice<f32>,
4527        input: &CudaSlice<f32>,
4528        t: usize,
4529        n_expert: usize,
4530        n_used: usize,
4531        bias: Option<&[f32]>,
4532        (sf, route_norm): (f32, bool),
4533        active: Option<&[bool]>,
4534    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
4535        let (lg, input) = e.dtoh_pair(logits, input)?;
4536        let (sel, w) =
4537            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
4538        Ok((sel, w, input))
4539    }
4540
4541    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
4542    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
4543    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
4544    /// active mask, prebuilt projection descriptors) so no model reference escapes.
4545    pub fn start_moe_prefetch_predictor(
4546        &self,
4547        e: &Engine,
4548        cfg: &ModelConfig,
4549    ) -> Result<(), Box<dyn std::error::Error>> {
4550        use crate::hybrid::Ffn;
4551        let Some(sig) = cfg.sigmoid_router() else {
4552            return Err("prefetch predictor requires a sigmoid-router arch".into());
4553        };
4554        let resident: std::collections::HashSet<(u16, u8, u16)> = e
4555            .export_moe_residency()
4556            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
4557            .into_iter()
4558            .collect();
4559        let mut layers = Vec::new();
4560        for (index, layer) in self.layers.iter().enumerate() {
4561            let Ffn::Moe(m) = &layer.ffn else { continue };
4562            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
4563            let router = e.dtoh(data)?;
4564            let n_expert = m.gate_exps.n_expert;
4565            let n_embd = m.gate_exps.in_f;
4566            if router.len() != n_embd * n_expert {
4567                continue;
4568            }
4569            let build = |exps: &crate::model::HostExps| {
4570                (0..n_expert)
4571                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
4572                    .collect::<Vec<_>>()
4573            };
4574            layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
4575                router,
4576                bias: m.exp_probs_b.clone(),
4577                active: m.active_experts.clone(),
4578                n_embd,
4579                n_used: cfg
4580                    .moe
4581                    .as_ref()
4582                    .map(|moe| moe.expert_used_count as usize)
4583                    .ok_or("prefetch predictor requires MoE config")?,
4584                sig,
4585                weights_n_expert: n_expert,
4586                gate: build(&m.gate_exps),
4587                up: build(&m.up_exps),
4588                down: build(&m.down_exps),
4589            }));
4590        }
4591        crate::cpu_experts::start_prefetch_predictor(layers, resident)
4592            .map_err(|error| error.into())
4593    }
4594
4595    /// Sigmoid-routing oracle shared by the prefetch predictor and `kernel-check`: identical
4596    /// selection math to the rollback runtime, applied to host-computed logits.
4597    #[allow(clippy::too_many_arguments)]
4598    pub fn moe_route_sigmoid_host_public(
4599        logits: &[f32],
4600        t: usize,
4601        n_expert: usize,
4602        n_used: usize,
4603        bias: Option<&[f32]>,
4604        sf: f32,
4605        route_norm: bool,
4606        active: Option<&[bool]>,
4607    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4608        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
4609    }
4610
4611    #[allow(clippy::too_many_arguments)]
4612    fn moe_route_sigmoid_host(
4613        lg: &[f32],
4614        t: usize,
4615        n_expert: usize,
4616        n_used: usize,
4617        bias: Option<&[f32]>,
4618        sf: f32,
4619        route_norm: bool,
4620        active: Option<&[bool]>,
4621    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4622        let active_count = active
4623            .map(|mask| mask.iter().filter(|&&enabled| enabled).count())
4624            .unwrap_or(n_expert);
4625        crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4626        if lg.len() != t * n_expert {
4627            return Err(format!(
4628                "sigmoid router logits length mismatch: got {}, expected {}",
4629                lg.len(),
4630                t * n_expert,
4631            )
4632            .into());
4633        }
4634        let mut sel = vec![0u32; t * n_used];
4635        let mut w_out = vec![0f32; t * n_used];
4636        for tok in 0..t {
4637            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
4638            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
4639            // selection score = sigmoid + bias; weight = plain sigmoid.
4640            let selsc: Vec<f32> = match bias {
4641                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
4642                None => scores.clone(),
4643            };
4644            let mut idx: Vec<usize> = (0..n_expert)
4645                .filter(|&i| active.is_none_or(|mask| mask[i]))
4646                .collect();
4647            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
4648            let sl = &idx[..n_used];
4649            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
4650            if route_norm {
4651                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
4652                for x in wv.iter_mut() {
4653                    *x = *x / ws * sf;
4654                }
4655            } else {
4656                for x in wv.iter_mut() {
4657                    *x *= sf;
4658                }
4659            }
4660            for j in 0..n_used {
4661                sel[tok * n_used + j] = sl[j] as u32;
4662                w_out[tok * n_used + j] = wv[j];
4663            }
4664        }
4665        Ok((sel, w_out))
4666    }
4667
4668    /// Step-3.7 resident-slab dispatch: sigmoid top-k outputs remain on device and feed the
4669    /// existing slot-ordered q8 expert kernels directly. Mixed layouts, spill, remote slabs,
4670    /// macro-scaled experts, and observation modes are denied by the caller.
4671    #[allow(clippy::too_many_arguments)]
4672    fn moe_ffn_sigmoid_dev(
4673        e: &Engine,
4674        m: &MoeWeights,
4675        z: &CudaSlice<f32>,
4676        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
4677        logits: &CudaSlice<f32>,
4678        t: usize,
4679        cfg: &ModelConfig,
4680        il: u16,
4681        (scaling_factor, route_norm): (f32, bool),
4682    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4683        let moe = cfg.moe.as_ref().unwrap();
4684        let n_embd = cfg.n_embd as usize;
4685        let n_expert = moe.expert_count as usize;
4686        let n_used = moe.expert_used_count as usize;
4687        let n_ff_exp = moe.expert_ff_length as usize;
4688        let dev = m.dev_exps.as_ref().unwrap();
4689        debug_assert!(cfg.step35.is_some());
4690        debug_assert_eq!(dev.dev, e.ctx().ordinal());
4691        debug_assert!(m.has_uniform_expert_layout());
4692        debug_assert!(!m.has_macros);
4693
4694        let (sel_d, w_d) = e.moe_router_sigmoid_topk(
4695            logits,
4696            t,
4697            n_expert,
4698            n_used,
4699            m.active_count(),
4700            &m.exp_probs_b_dev,
4701            &m.active_experts_dev,
4702            scaling_factor,
4703            route_norm,
4704        )?;
4705        crate::moesd::record_device_routes(e, il, n_expert, n_used, &sel_d)?;
4706        let (gate_row_bytes, up_row_bytes) = if dev.gu_il {
4707            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
4708            (combined, combined)
4709        } else {
4710            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
4711        };
4712        let (zq, zd) = match (t, zq8) {
4713            (1, Some((q, d))) => (q.clone(), d.clone()),
4714            _ => e.quantize_q8_1(z, t, n_embd)?,
4715        };
4716        let n_pairs = t * n_used;
4717        let mut moe_out = if cfg.clamp_exp_at(il as u32).is_some() {
4718            // The final Step layers retain the established separate gate/up -> clamp -> down
4719            // arithmetic. Pair rows are derived from token position; selected expert ids and
4720            // routing weights remain the device router's buffers throughout.
4721            let pair_tok: Vec<i32> = (0..n_pairs)
4722                .map(|pair| (pair / n_used) as i32)
4723                .collect();
4724            let pair_tok_d = e.htod_i32(&pair_tok)?;
4725            let gate = e.moe_pairs_matvec_q8(
4726                &dev.ptr_row,
4727                0,
4728                &pair_tok_d,
4729                &sel_d,
4730                &zq,
4731                &zd,
4732                n_embd,
4733                n_ff_exp,
4734                n_expert,
4735                n_pairs,
4736                m.gate_exps.qtype,
4737                gate_row_bytes,
4738            )?;
4739            let up = e.moe_pairs_matvec_q8(
4740                &dev.ptr_row,
4741                1,
4742                &pair_tok_d,
4743                &sel_d,
4744                &zq,
4745                &zd,
4746                n_embd,
4747                n_ff_exp,
4748                n_expert,
4749                n_pairs,
4750                m.up_exps.qtype,
4751                up_row_bytes,
4752            )?;
4753            let mut act = e.uninit(n_pairs * n_ff_exp)?;
4754            Self::ffn_act_lim(
4755                e,
4756                cfg,
4757                &gate,
4758                &up,
4759                1.0,
4760                1.0,
4761                cfg.clamp_exp_at(il as u32),
4762                &mut act,
4763                n_pairs * n_ff_exp,
4764            )?;
4765            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4766            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
4767            let pair_self_d = e.htod_i32(&pair_self)?;
4768            let down = e.moe_pairs_matvec_q8(
4769                &dev.ptr_row,
4770                2,
4771                &pair_self_d,
4772                &sel_d,
4773                &aq2,
4774                &ad2,
4775                n_ff_exp,
4776                n_embd,
4777                n_expert,
4778                n_pairs,
4779                m.down_exps.qtype,
4780                m.down_exps.row_bytes,
4781            )?;
4782            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4783            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
4784            let tok_off_d = e.htod_i32(&tok_off)?;
4785            let tok_ids_d = e.htod_i32(&tok_ids)?;
4786            let mut output = e.uninit(t * n_embd)?;
4787            e.moe_pairs_scatter(
4788                &down,
4789                &w_d,
4790                &tok_off_d,
4791                &tok_ids_d,
4792                &mut output,
4793                t,
4794                n_embd,
4795            )?;
4796            output
4797        } else {
4798            let act = e.moe_gate_up_silu8_dev_q8_rows(
4799                &dev.ptr_row,
4800                &sel_d,
4801                &zq,
4802                &zd,
4803                t,
4804                n_embd,
4805                n_ff_exp,
4806                n_used,
4807                n_expert,
4808                m.gate_exps.qtype,
4809                m.up_exps.qtype,
4810                gate_row_bytes,
4811                up_row_bytes,
4812                &m.dev_macros,
4813            )?;
4814            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4815            let mut output = e.uninit(t * n_embd)?;
4816            e.moe_down8_fma_dev_q8_rows_g(
4817                &dev.ptr_row,
4818                &sel_d,
4819                &w_d,
4820                &aq2,
4821                &ad2,
4822                &mut output,
4823                t,
4824                n_ff_exp,
4825                n_embd,
4826                n_used,
4827                n_expert,
4828                m.down_exps.qtype,
4829                m.down_exps.row_bytes,
4830            )?;
4831            output
4832        };
4833
4834        if std::env::var("MEMRA_SIG_ROUTER_DISPATCH_TRACE").as_deref() == Ok("1") {
4835            eprintln!(
4836                "[sigrouter-dev] layer={il} tokens={t} experts={n_expert} used={n_used} clamp={} gu_il={}",
4837                cfg.clamp_exp_at(il as u32).is_some(),
4838                dev.gu_il,
4839            );
4840        }
4841        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
4842        Ok(moe_out)
4843    }
4844
4845    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
4846    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
4847    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
4848    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
4849    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
4850    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
4851    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
4852    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
4853    fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
4854                     t: usize, cfg: &ModelConfig)
4855                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4856        let moe = cfg.moe.as_ref().unwrap();
4857        let n_embd = cfg.n_embd as usize;
4858        let n_expert = moe.expert_count as usize;
4859        let n_used = moe.expert_used_count as usize;
4860        let n_ff_exp = moe.expert_ff_length as usize;
4861        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
4862        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
4863        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
4864        // that forgets the gate fails loudly in debug instead of returning wrong logits.
4865        debug_assert!(!cfg.swiglu_clamped_anywhere(),
4866                      "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU");
4867        let dev = m.dev_exps.as_ref().unwrap();
4868        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
4869        let (rbg_d, rbu_d) = if dev.gu_il {
4870            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
4871        } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
4872
4873        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
4874        let n_pairs = t * n_used;
4875        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
4876        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
4877        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
4878        let pair_ex:  Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
4879        let pair_w:   Vec<f32> = w_all.clone();
4880        let tok_off:  Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4881        let tok_ids:  Vec<i32> = (0..n_pairs as i32).collect();
4882        let pt = e.htod_i32(&pair_tok)?;
4883        let px = e.htod_i32(&pair_ex)?;
4884        let pw = e.htod(&pair_w)?;
4885        let toff = e.htod_i32(&tok_off)?;
4886        let tids = e.htod_i32(&tok_ids)?;
4887
4888        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
4889        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
4890        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
4891        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
4892        for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
4893        let mut ex_ids: Vec<i32> = Vec::new();
4894        let mut ex_off: Vec<i32> = vec![0];
4895        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
4896        for (ex, list) in by_ex.iter().enumerate() {
4897            if list.is_empty() { continue; }
4898            ex_ids.push(ex as i32);
4899            ex_pairs.extend_from_slice(list);
4900            ex_off.push(ex_pairs.len() as i32);
4901        }
4902        let n_active = ex_ids.len();
4903        let exi = e.htod_i32(&ex_ids)?;
4904        let exo = e.htod_i32(&ex_off)?;
4905        let exp_d = e.htod_i32(&ex_pairs)?;
4906        let _ = &px;   // pair-major twin keeps it; em path uses CSR
4907
4908        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
4909        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
4910        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
4911        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
4912        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
4913        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
4914        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
4915        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
4916        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
4917        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
4918        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
4919        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
4920        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
4921        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
4922        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
4923        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
4924        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
4925        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
4926        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4927        let mma_t = *MMA_T.get_or_init(|| {
4928            std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
4929        });
4930        let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
4931            && t >= mma_t
4932            && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
4933            && q8_expert_dec_supported(m.down_exps.qtype)
4934            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4935        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
4936        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
4937        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
4938        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
4939        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
4940        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
4941        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
4942        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
4943        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
4944        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
4945        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
4946        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
4947        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
4948        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
4949        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
4950        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
4951            && q8_expert_dec_supported(m.up_exps.qtype)
4952            && q8_expert_dec_supported(m.down_exps.qtype)
4953            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
4954        let f16g_mode = crate::moe_f16g_mode();
4955        let f16g = f16g_mode != 0 && t >= mma_t
4956            && (f16g_mode != 3 || !mma_capable)
4957            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
4958            && f16g_proj_ok(m.up_exps.qtype, n_embd)
4959            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
4960        if use_mma || f16g {
4961            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
4962            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
4963            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
4964            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
4965            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
4966            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
4967            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
4968            let y_down = if f16g {
4969                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
4970                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
4971                // permute at the very end back to pair-id order for the scatter.
4972                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
4973                let csr_tok_d = e.htod_i32(&csr_tok)?;
4974                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
4975                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
4976                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4977                                              m.gate_exps.qtype, rbg_d)?;
4978                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
4979                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4980                                              m.up_exps.qtype, rbu_d)?;
4981                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
4982                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
4983                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
4984                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
4985                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
4986                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
4987            } else {
4988            // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
4989            let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
4990            let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4991                                        n_embd, n_ff_exp, n_active, n_pairs, t,
4992                                        m.gate_exps.qtype, rbg_d)?;
4993            let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4994                                      n_embd, n_ff_exp, n_active, n_pairs, t,
4995                                      m.up_exps.qtype, rbu_d)?;
4996            // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
4997            // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
4998            // registers and writes ONLY the quantized scratch — the two-pass chain
4999            // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
5000            // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
5001            let a_scr = if crate::moe_fuse_actq_on() {
5002                e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
5003            } else {
5004                let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5005                e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
5006            };
5007            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5008            let pself = e.htod_i32(&pair_self)?;
5009            e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
5010                             n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
5011                             m.down_exps.qtype, m.down_exps.row_bytes)?
5012            };
5013            let mut moe_out = e.uninit(t * n_embd)?;
5014            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5015            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5016                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5017            {
5018                let n_ff_sh = gate_shexp.out_features();
5019                let sg_gate = e.matmul(gate_shexp, z, t)?;
5020                let sg_up = e.matmul(up_shexp, z, t)?;
5021                let mut sa = e.uninit(t * n_ff_sh)?;
5022                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5023                let sh = e.matmul(down_shexp, &sa, t)?;
5024                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
5025                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
5026                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
5027                // i.e. the one real prefill actually takes on a resident-expert MoE model,
5028                // so the concat-prime isolation fix has to land here as well.
5029                let g = match &m.gate_inp_shexp {
5030                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
5031                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5032                    }
5033                    Some(gate_inp_shexp) => {
5034                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5035                        let mut g = e.uninit(t)?;
5036                        e.sigmoid(&gs, &mut g, t)?;
5037                        g
5038                    }
5039                    None => e.htod(&vec![1.0f32; t])?,
5040                };
5041                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5042            }
5043            return Ok(moe_out);
5044        }
5045
5046        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
5047        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
5048        let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
5049        let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
5050                      inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5051            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
5052            let dec = dec && q8_expert_dec_supported(qtype);
5053            if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
5054                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
5055            else   { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
5056                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
5057        };
5058        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5059        let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
5060                          n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
5061        let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
5062                        n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
5063        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5064        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5065        // down consumes PAIR-major activation rows: pair_tok = identity.
5066        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5067        let pself = e.htod_i32(&pair_self)?;
5068        let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
5069                            n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
5070        let mut moe_out = e.uninit(t * n_embd)?;   // scatter fully overwrites per (token,col)
5071        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5072
5073        // SHARED EXPERT epilogue — same as the other paths.
5074        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5075        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5076        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5077            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5078        {
5079            let n_ff_sh = gate_shexp.out_features();
5080            // These decode-exact forms are required by the new Step resident arm. Keep the
5081            // established grouped shared-expert program for every other architecture: widening
5082            // this to Gemma changed its speculative acceptance despite green argmax gates.
5083            let step_exact = cfg.step35.is_some();
5084            let verify_t = step_exact && t > 1 && t < PRIME_MIN_T;
5085            let (sg_gate, sg_up) = if step_exact && t == 1 {
5086                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5087                    Some(pair) => pair,
5088                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5089                }
5090            } else if verify_t {
5091                let mut fused = None;
5092                if crate::spec::spec_fused_t() && (2..=4).contains(&t)
5093                    && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp)
5094                {
5095                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5096                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
5097                }
5098                match fused {
5099                    Some(pair) => pair,
5100                    None => (
5101                        e.matmul_decode_exact(gate_shexp, z, t)?,
5102                        e.matmul_decode_exact(up_shexp, z, t)?,
5103                    ),
5104                }
5105            } else {
5106                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
5107            };
5108            let mut sa = e.uninit(t * n_ff_sh)?;
5109            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5110            let sh = if verify_t {
5111                e.matmul_decode_exact(down_shexp, &sa, t)?
5112            } else {
5113                e.matmul(down_shexp, &sa, t)?
5114            };
5115            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
5116            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
5117            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
5118            // dispatch choice cannot change bits.
5119            let g = match &m.gate_inp_shexp {
5120                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
5121                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5122                }
5123                Some(gate_inp_shexp) => {
5124                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5125                    let mut g = e.uninit(t)?;
5126                    e.sigmoid(&gs, &mut g, t)?;
5127                    g
5128                }
5129                None => e.htod(&vec![1.0f32; t])?,
5130            };
5131            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5132        }
5133        Ok(moe_out)
5134    }
5135
5136    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
5137    #[allow(clippy::too_many_arguments)]
5138    #[allow(clippy::too_many_arguments)]
5139    fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
5140                   zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
5141                   t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
5142                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5143        let moe = cfg.moe.as_ref().unwrap();
5144        let n_embd = cfg.n_embd as usize;
5145        let n_expert = moe.expert_count as usize;
5146        let n_used = moe.expert_used_count as usize;
5147        let n_ff_exp = moe.expert_ff_length as usize;
5148        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
5149        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
5150        // clamped layers; assert both so a future caller that skips the gate fails loudly.
5151        debug_assert!(cfg.sigmoid_router().is_none(),
5152                      "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts");
5153        debug_assert!(!cfg.swiglu_clamped_at(il as u32),
5154                      "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form");
5155
5156        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
5157        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
5158        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
5159        // skipped entirely for macro-free experts (every k-quant GGUF).
5160        if m.has_macros {
5161            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
5162        }
5163
5164        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
5165        let mut moe_out = e.uninit(t * n_embd)?;
5166
5167        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
5168        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
5169        if let Some(dev) = m.dev_exps.as_ref() {
5170            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
5171            // the combined stride; up's base is offset in the ptr table. Down unchanged.
5172            let (rbg_d, rbu_d) = if dev.gu_il {
5173                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
5174            } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
5175            let q8 = moe_q8_enabled()
5176                && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
5177                && q8_expert_supported(m.down_exps.qtype);
5178            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
5179            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
5180            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
5181            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
5182            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
5183            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
5184            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
5185            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
5186            let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
5187                && n_ff_exp == 512 && n_used <= 8
5188                && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
5189                && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
5190            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
5191            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
5192            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
5193            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
5194            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
5195            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
5196            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
5197            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
5198            let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
5199                .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
5200            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
5201            let csr_arm = rows_arm && csr_mode > 0 && t <= 10
5202                && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
5203                && csr_qt(m.down_exps.qtype);
5204            if csr_arm {
5205                if csr_mode == 2 {
5206                    static ENGAGED: std::sync::Once = std::sync::Once::new();
5207                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
5208                }
5209                let n_pairs = t * n_used;
5210                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5211                let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
5212                                                         n_embd, n_ff_exp, n_used, n_expert,
5213                                                         m.gate_exps.qtype, m.up_exps.qtype,
5214                                                         rbg_d, rbu_d)?;
5215                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5216                // down stays on the _rows twin — BOTH CSR down variants measured negative
5217                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
5218                // 16-group rows have too little decode to amortize any dedup structure.
5219                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
5220                                            t, n_ff_exp, n_embd, n_used, n_expert,
5221                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
5222                if csr_mode == 2 {
5223                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
5224                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5225                                                                n_embd, n_ff_exp, n_used, n_expert,
5226                                                                m.gate_exps.qtype, m.up_exps.qtype,
5227                                                                rbg_d, rbu_d, &m.dev_macros)?;
5228                    let mut out_r = e.uninit(t * n_embd)?;
5229                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
5230                    e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
5231                                                t, n_ff_exp, n_embd, n_used, n_expert,
5232                                                m.down_exps.qtype, m.down_exps.row_bytes)?;
5233                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
5234                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
5235                    let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
5236                    let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
5237                    if ba + bo > 0 {
5238                        eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
5239                                  a1.len(), o1.len());
5240                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
5241                        let sel_h = e.dtoh_i32(&sel_d)?;
5242                        let mut shown = 0;
5243                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
5244                            if x.to_bits() != y.to_bits() && shown < 4 {
5245                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
5246                                let ex = sel_h[p];
5247                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
5248                                eprintln!("  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
5249                                shown += 1;
5250                            }
5251                        }
5252                        std::process::exit(3);
5253                    }
5254                }
5255            } else if rows_arm {
5256                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
5257                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
5258                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
5259                    use std::sync::atomic::{AtomicU64, Ordering};
5260                    static PAIRS: AtomicU64 = AtomicU64::new(0);
5261                    static UNIQ: AtomicU64 = AtomicU64::new(0);
5262                    static CALLS: AtomicU64 = AtomicU64::new(0);
5263                    let sel_h = e.dtoh_i32(&sel_d)?;
5264                    let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
5265                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
5266                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
5267                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
5268                    if c % 480 == 0 {
5269                        let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
5270                        eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
5271                                  q as f64 / p as f64);
5272                    }
5273                }
5274                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5275                let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5276                                                          n_embd, n_ff_exp, n_used, n_expert,
5277                                                          m.gate_exps.qtype, m.up_exps.qtype,
5278                                                          rbg_d, rbu_d, &m.dev_macros)?;
5279                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5280                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
5281                                            t, n_ff_exp, n_embd, n_used, n_expert,
5282                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
5283            } else {
5284            for tok in 0..t {
5285                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
5286                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
5287                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
5288                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5289                if q8 {
5290                    let (zq, zd) = match (t, zq8) {
5291                        (1, Some((q, d))) => (q.clone(), d.clone()),
5292                        _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
5293                    };
5294                    let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
5295                                                         n_embd, n_ff_exp, n_used, n_expert,
5296                                                         m.gate_exps.qtype, m.up_exps.qtype,
5297                                                         rbg_d, rbu_d, &m.dev_macros)?;
5298                    let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5299                    e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
5300                                           n_ff_exp, n_embd, n_used, n_expert,
5301                                           m.down_exps.qtype, m.down_exps.row_bytes)?;
5302                } else {
5303                    let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
5304                                                      n_used, n_expert,
5305                                                      m.gate_exps.qtype, m.up_exps.qtype,
5306                                                      rbg_d, rbu_d, &m.dev_macros)?;
5307                    e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
5308                                        n_ff_exp, n_embd, n_used, n_expert,
5309                                        m.down_exps.qtype, m.down_exps.row_bytes)?;
5310                }
5311            }
5312            }
5313        } else {
5314        // Launch under the cache lock: the row borrow lives as long as the closure, and the
5315        // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
5316        // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
5317        // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
5318        // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
5319        // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
5320        let q8 = moe_q8_enabled()
5321            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
5322            && q8_expert_supported(m.down_exps.qtype);
5323        e.with_moe_cache(max_block, |c, eng| {
5324            let row = c.layer_dev_row(il, n_expert, eng)?
5325                .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
5326            for tok in 0..t {
5327                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
5328                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
5329                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
5330                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5331                if q8 {
5332                    let (zq, zd) = match (t, zq8) {
5333                        (1, Some((q, d))) => (q.clone(), d.clone()),
5334                        _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
5335                    };
5336                    let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
5337                                                           n_embd, n_ff_exp, n_used, n_expert,
5338                                                           m.gate_exps.qtype, m.up_exps.qtype,
5339                                                           m.gate_exps.row_bytes, m.up_exps.row_bytes,
5340                                                           &m.dev_macros)?;
5341                    let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
5342                    eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
5343                                             n_ff_exp, n_embd, n_used, n_expert,
5344                                             m.down_exps.qtype, m.down_exps.row_bytes)?;
5345                } else {
5346                    let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
5347                                                        n_used, n_expert,
5348                                                        m.gate_exps.qtype, m.up_exps.qtype,
5349                                                        m.gate_exps.row_bytes, m.up_exps.row_bytes,
5350                                                        &m.dev_macros)?;
5351                    eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
5352                                          n_ff_exp, n_embd, n_used, n_expert,
5353                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
5354                }
5355            }
5356            // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
5357            c.hits += (t * 3 * n_used) as u64;
5358            Ok(())
5359        })?;
5360        }
5361
5362        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
5363        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
5364        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
5365        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
5366        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5367            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5368        {
5369            let n_ff_sh = gate_shexp.out_features();
5370            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
5371            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
5372            let verify_t = t > 1 && t < PRIME_MIN_T;
5373            let (sg_gate, sg_up) = if t == 1 {
5374                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
5375                    Some(pair) => pair,
5376                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
5377                }
5378            } else if verify_t {
5379                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
5380                // rides one shared quantize + one fused2 batched launch instead of two
5381                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
5382                let mut fused = None;
5383                if crate::spec::spec_fused_t() && (2..=4).contains(&t)
5384                    && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
5385                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5386                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
5387                }
5388                match fused {
5389                    Some(pair) => pair,
5390                    None => (e.matmul_decode_exact(gate_shexp, z, t)?,
5391                             e.matmul_decode_exact(up_shexp, z, t)?),
5392                }
5393            } else {
5394                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
5395            };
5396            let mut sa = e.uninit(t * n_ff_sh)?;  // silu_mul fully overwrites
5397            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
5398            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
5399                     else { e.matmul(down_shexp, &sa, t)? };
5400            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
5401            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
5402            // between the two arms; prefill keeps the batched cuBLASLt linear).
5403            let g = match &m.gate_inp_shexp {
5404                Some(gate_inp_shexp) => {
5405                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
5406                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
5407                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5408                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5409                    } else {
5410                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5411                        let mut g = e.uninit(t)?;
5412                        e.sigmoid(&gs, &mut g, t)?;
5413                        g
5414                    }
5415                }
5416                None => e.htod(&vec![1.0f32; t])?,
5417            };
5418            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
5419        }
5420
5421        Ok(moe_out)
5422    }
5423
5424    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
5425    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
5426    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
5427    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
5428    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
5429    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
5430    /// the collected raw pointers cannot move between collection and launch (single-threaded
5431    /// decode; the lock is held only for collection, launches are stream-ordered after any
5432    /// prior same-stream staging writes).
5433    #[allow(clippy::too_many_arguments)]
5434    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
5435    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
5436    #[allow(clippy::too_many_arguments)]
5437    fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5438                      zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
5439                      moe_out: &mut CudaSlice<f32>, tok: usize,
5440                      n_embd: usize, n_ff_exp: usize, n_used: usize)
5441                      -> Result<bool, Box<dyn std::error::Error>> {
5442        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5443        use cudarc::driver::DevicePtr;
5444        let ptrs = e.with_moe_cache(max_block, |c, eng| {
5445            let mut g = [0u64; 8];
5446            let mut u = [0u64; 8];
5447            let mut d = [0u64; 8];
5448            for (j, &ex) in sel.iter().enumerate() {
5449                let ex = ex as u16;
5450                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5451                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
5452                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5453                else { return Ok(None); };
5454                let __s = eng.stream();
5455                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5456                let (pu, _e1) = c.slot(su).device_ptr(&__s);
5457                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5458                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5459            }
5460            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5461                for &ex in sel {
5462                    let ex = ex as u16;
5463                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5464                        c.note_profile_hit(BlockId::new(il, proj, ex));
5465                    }
5466                }
5467            }
5468            c.hits += (3 * n_used) as u64;
5469            Ok(Some((g, u, d)))
5470        })?;
5471        let Some((g, u, d)) = ptrs else { return Ok(false) };
5472        let mut wv = [0f32; 8];
5473        wv[..n_used].copy_from_slice(w);
5474        let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
5475                                         n_embd, n_ff_exp, n_used,
5476                                         m.gate_exps.qtype, m.up_exps.qtype,
5477                                         m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5478        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
5479        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5480        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5481        e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
5482                           n_ff_exp, n_embd, n_used,
5483                           m.down_exps.qtype, m.down_exps.row_bytes)?;
5484        Ok(true)
5485    }
5486
5487    fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
5488                      zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
5489                      moe_out: &mut CudaSlice<f32>, tok: usize,
5490                      n_embd: usize, n_ff_exp: usize, n_used: usize)
5491                      -> Result<bool, Box<dyn std::error::Error>> {
5492        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
5493        use cudarc::driver::DevicePtr;
5494        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
5495        let ptrs = e.with_moe_cache(max_block, |c, eng| {
5496            let mut g = [0u64; 8];
5497            let mut u = [0u64; 8];
5498            let mut d = [0u64; 8];
5499            for (j, &ex) in sel.iter().enumerate() {
5500                let ex = ex as u16;
5501                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
5502                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
5503                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
5504                else { return Ok(None); };
5505                let __s = eng.stream();
5506                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
5507                let (pu, _e1) = c.slot(su).device_ptr(&__s);
5508                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
5509                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
5510            }
5511            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
5512                for &ex in sel {
5513                    let ex = ex as u16;
5514                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
5515                        c.note_profile_hit(BlockId::new(il, proj, ex));
5516                    }
5517                }
5518            }
5519            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
5520            Ok(Some((g, u, d)))
5521        })?;
5522        let Some((g, u, d)) = ptrs else { return Ok(false) };
5523        let mut wv = [0f32; 8];
5524        wv[..n_used].copy_from_slice(w);
5525        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
5526        let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
5527                                      n_embd, n_ff_exp, n_used,
5528                                      m.gate_exps.qtype, m.up_exps.qtype,
5529                                      m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5530        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5531        e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
5532                             n_ff_exp, n_embd, n_used,
5533                             m.down_exps.qtype, m.down_exps.row_bytes)?;
5534        Ok(true)
5535    }
5536
5537    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
5538    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
5539    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
5540    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
5541    fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5542                          max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
5543                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5544        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5545        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5546        let layout = exps.expert_layout(ex);
5547        let id = BlockId::new(il, proj, ex as u16);
5548        let source = exps.expert_source(ex);
5549        e.with_moe_cache(max_block, |c, eng| {
5550            let slot = c.dispatch_source(id, source, eng)?;
5551            let DispatchSlot::Resident(sl) = slot;
5552            let buf = c.slot(sl);
5553            eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
5554                                  layout.qtype, layout.row_bytes)
5555        })
5556    }
5557
5558    fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
5559                       max_block: usize, x: &cudarc::driver::CudaView<f32>)
5560                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5561        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
5562        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
5563        let layout = exps.expert_layout(ex);
5564        let id = BlockId::new(il, proj, ex as u16);
5565        let source = exps.expert_source(ex);
5566        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
5567        e.with_moe_cache(max_block, |c, eng| {
5568            let slot = c.dispatch_source(id, source, eng)?;
5569            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
5570            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
5571            let DispatchSlot::Resident(sl) = slot;
5572            let buf = c.slot(sl);
5573            eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
5574                             layout.qtype, layout.row_bytes)
5575        })
5576    }
5577
5578    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
5579    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
5580    /// so the current forward's backend assignment and output remain unchanged.
5581    fn moe_profile_admit_expert(
5582        e: &Engine,
5583        il: u16,
5584        ex: usize,
5585        m: &MoeWeights,
5586        max_block: usize,
5587    ) -> Result<(), Box<dyn std::error::Error>> {
5588        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5589        e.with_moe_cache(max_block, |cache, eng| {
5590            for (proj, exps) in [
5591                (PROJ_GATE, &m.gate_exps),
5592                (PROJ_UP, &m.up_exps),
5593                (PROJ_DOWN, &m.down_exps),
5594            ] {
5595                let id = BlockId::new(il, proj, ex as u16);
5596                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
5597            }
5598            Ok(())
5599        })
5600    }
5601
5602    /// Read a projection from the immutable residency set when present; otherwise use one
5603    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
5604    #[allow(clippy::too_many_arguments)]
5605    fn moe_frozen_gemm(
5606        e: &Engine,
5607        il: u16,
5608        proj: u8,
5609        ex: usize,
5610        m: &MoeWeights,
5611        max_block: usize,
5612        x: &cudarc::driver::CudaView<f32>,
5613        scratch: &mut Option<CudaSlice<u8>>,
5614        scratch_len: usize,
5615    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5616        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
5617        let exps = match proj {
5618            PROJ_GATE => &m.gate_exps,
5619            PROJ_UP => &m.up_exps,
5620            _ => &m.down_exps,
5621        };
5622        let layout = exps.expert_layout(ex);
5623        let id = BlockId::new(il, proj, ex as u16);
5624        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
5625            let Some(slot) = cache.resident(id) else {
5626                return Ok(None);
5627            };
5628            let buf = cache.slot(slot);
5629            Ok(Some(eng.qmatvec_view(
5630                buf,
5631                0..layout.len,
5632                x,
5633                1,
5634                exps.in_f,
5635                exps.out_f,
5636                layout.qtype,
5637                layout.row_bytes,
5638            )?))
5639        })? {
5640            return Ok(output);
5641        }
5642        if scratch.is_none() {
5643            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
5644        }
5645        let scratch = scratch.as_mut().unwrap();
5646        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
5647        e.qmatvec_view(
5648            scratch,
5649            0..layout.len,
5650            x,
5651            1,
5652            exps.in_f,
5653            exps.out_f,
5654            layout.qtype,
5655            layout.row_bytes,
5656        )
5657    }
5658
5659    fn moe_prefetch_expert(
5660        e: &Engine,
5661        il: u16,
5662        ex: usize,
5663        m: &MoeWeights,
5664        max_block: usize,
5665        keep: &[crate::moe_cache::BlockId],
5666    ) -> Result<(), Box<dyn std::error::Error>> {
5667        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5668        e.with_moe_cache(max_block, |c, eng| {
5669            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5670                                 (PROJ_DOWN, &m.down_exps)] {
5671                let id = BlockId::new(il, proj, ex as u16);
5672                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
5673            }
5674            Ok(())
5675        })
5676    }
5677
5678    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
5679    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
5680    fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
5681                                max_block: usize, keep: &[crate::moe_cache::BlockId])
5682                                -> Result<(), Box<dyn std::error::Error>> {
5683        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
5684        e.with_moe_cache(max_block, |c, eng| {
5685            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
5686                                 (PROJ_DOWN, &m.down_exps)] {
5687                let source = exps.expert_source(ex);
5688                if let crate::model::ExpertSource::Disk { .. } = &source {
5689                    let id = BlockId::new(il, proj, ex as u16);
5690                    let _ = c.prefetch_source(id, source, keep, eng)?;
5691                }
5692            }
5693            Ok(())
5694        })
5695    }
5696
5697    #[inline]
5698    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
5699        let _ = m.gate_exps.prefetch_expert_pages(ex);
5700        let _ = m.up_exps.prefetch_expert_pages(ex);
5701        let _ = m.down_exps.prefetch_expert_pages(ex);
5702    }
5703}
5704
5705// ================================================================================================
5706// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
5707//
5708// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
5709// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
5710// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
5711//
5712// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
5713// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
5714// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
5715// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
5716// identical to the per-token loop regardless of expert processing order.
5717//
5718// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
5719// ================================================================================================
5720
5721impl HybridModel {
5722    /// Resident-slab fast path for host-routed grouped prefill. Unclamped layers batch the
5723    /// sequential fused q8 program over the token axis; clamped layers use the separate
5724    /// expert-major q8 chain so `ffn_act_lim` remains authoritative.
5725    #[allow(clippy::too_many_arguments)]
5726    fn moe_ffn_grouped_resident_q8(
5727        e: &Engine,
5728        m: &MoeWeights,
5729        z: &CudaSlice<f32>,
5730        t: usize,
5731        cfg: &ModelConfig,
5732        il: u16,
5733        sel_all: &[u32],
5734        w_all: &[f32],
5735        table: &CudaSlice<u64>,
5736        gu_il: bool,
5737    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5738        let moe = cfg.moe.as_ref().unwrap();
5739        let n_embd = cfg.n_embd as usize;
5740        let n_expert = moe.expert_count as usize;
5741        let n_used = moe.expert_used_count as usize;
5742        let n_ff_exp = moe.expert_ff_length as usize;
5743        let n_pairs = t * n_used;
5744        debug_assert_eq!(sel_all.len(), n_pairs);
5745        debug_assert_eq!(w_all.len(), n_pairs);
5746        debug_assert!(
5747            m.gate_exps.macros.is_none()
5748                && m.up_exps.macros.is_none()
5749                && m.down_exps.macros.is_none(),
5750            "resident grouped q8 does not fold per-expert macro scales",
5751        );
5752
5753        // The rows twins run the resident sequential program verbatim on grid.z = token:
5754        // fused gate/up/SiLU per slot, batched activation quantization, then the original
5755        // slot-ordered down/FMA chain. Routing remains the sigmoid selector above (device by
5756        // default; MEMRA_SIG_ROUTER=0 is the host oracle); these kernels consume sel/w only and
5757        // never enter the softmax router.
5758        if !cfg.swiglu_clamped_at(il as u32) {
5759            let sel: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5760            let sel_d = e.htod_i32(&sel)?;
5761            let w_d = e.htod(w_all)?;
5762            let (gate_row_bytes, up_row_bytes) = if gu_il {
5763                let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5764                (combined, combined)
5765            } else {
5766                (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5767            };
5768            let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5769            let act = e.moe_gate_up_silu8_dev_q8_rows(
5770                table,
5771                &sel_d,
5772                &zq,
5773                &zd,
5774                t,
5775                n_embd,
5776                n_ff_exp,
5777                n_used,
5778                n_expert,
5779                m.gate_exps.qtype,
5780                m.up_exps.qtype,
5781                gate_row_bytes,
5782                up_row_bytes,
5783                &m.dev_macros,
5784            )?;
5785            let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5786            let mut moe_out = e.uninit(t * n_embd)?;
5787            e.moe_down8_fma_dev_q8_rows_g(
5788                table,
5789                &sel_d,
5790                &w_d,
5791                &aq2,
5792                &ad2,
5793                &mut moe_out,
5794                t,
5795                n_ff_exp,
5796                n_embd,
5797                n_used,
5798                n_expert,
5799                m.down_exps.qtype,
5800                m.down_exps.row_bytes,
5801            )?;
5802
5803            if std::env::var("MEMRA_MOE_STATS").is_ok() {
5804                let mut counts = vec![0usize; n_expert];
5805                for &expert in sel_all {
5806                    counts[expert as usize] += 1;
5807                }
5808                let mut sizes: Vec<usize> =
5809                    counts.into_iter().filter(|&count| count != 0).collect();
5810                sizes.sort_unstable();
5811                let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5812                println!(
5813                    "moe-grouped il={il} t={t} dispatch=resident-q8-rows active={}/{} \
5814                     m_e: min={} median={} mean={mean:.1} max={}",
5815                    sizes.len(),
5816                    n_expert,
5817                    sizes.first().copied().unwrap_or(0),
5818                    sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5819                    sizes.last().copied().unwrap_or(0),
5820                );
5821            }
5822            return Ok(moe_out);
5823        }
5824
5825        // Clamped layers keep gate/up, activation, and down as separate stages. The pair-major
5826        // matvec body is qmatvec_expert_q8 verbatim, while one launch covers every routed pair.
5827        // Pair ids stay in router slot order so scatter preserves the sequential FMA chain.
5828        let pair_tok: Vec<i32> = (0..n_pairs).map(|pair| (pair / n_used) as i32).collect();
5829        let pair_ex: Vec<i32> = sel_all.iter().map(|&expert| expert as i32).collect();
5830        let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5831        let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5832
5833        let mut by_expert: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
5834        for (pair, &expert) in pair_ex.iter().enumerate() {
5835            by_expert[expert as usize].push(pair as i32);
5836        }
5837
5838        let pair_tok_d = e.htod_i32(&pair_tok)?;
5839        let pair_ex_d = e.htod_i32(&pair_ex)?;
5840        let pair_w_d = e.htod(w_all)?;
5841        let tok_off_d = e.htod_i32(&tok_off)?;
5842        let tok_ids_d = e.htod_i32(&tok_ids)?;
5843
5844        let matvec = |
5845            proj: i32,
5846            pair_rows: &CudaSlice<i32>,
5847            aq: &CudaSlice<i8>,
5848            ad: &CudaSlice<f32>,
5849            in_f: usize,
5850            out_f: usize,
5851            qtype: i32,
5852            row_bytes: usize,
5853        | -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5854            e.moe_pairs_matvec_q8(
5855                table,
5856                proj,
5857                pair_rows,
5858                &pair_ex_d,
5859                aq,
5860                ad,
5861                in_f,
5862                out_f,
5863                n_expert,
5864                n_pairs,
5865                qtype,
5866                row_bytes,
5867            )
5868        };
5869
5870        let (gate_row_bytes, up_row_bytes) = if gu_il {
5871            let combined = m.gate_exps.row_bytes + m.up_exps.row_bytes;
5872            (combined, combined)
5873        } else {
5874            (m.gate_exps.row_bytes, m.up_exps.row_bytes)
5875        };
5876        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
5877        let gate = matvec(
5878            0,
5879            &pair_tok_d,
5880            &zq,
5881            &zd,
5882            n_embd,
5883            n_ff_exp,
5884            m.gate_exps.qtype,
5885            gate_row_bytes,
5886        )?;
5887        let up = matvec(
5888            1,
5889            &pair_tok_d,
5890            &zq,
5891            &zd,
5892            n_embd,
5893            n_ff_exp,
5894            m.up_exps.qtype,
5895            up_row_bytes,
5896        )?;
5897        let mut act = e.uninit(n_pairs * n_ff_exp)?;
5898        Self::ffn_act_lim(
5899            e,
5900            cfg,
5901            &gate,
5902            &up,
5903            1.0,
5904            1.0,
5905            cfg.clamp_exp_at(il as u32),
5906            &mut act,
5907            n_pairs * n_ff_exp,
5908        )?;
5909        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5910        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5911        let pair_self_d = e.htod_i32(&pair_self)?;
5912        let down = matvec(
5913            2,
5914            &pair_self_d,
5915            &aq2,
5916            &ad2,
5917            n_ff_exp,
5918            n_embd,
5919            m.down_exps.qtype,
5920            m.down_exps.row_bytes,
5921        )?;
5922        let mut moe_out = e.uninit(t * n_embd)?;
5923        e.moe_pairs_scatter(
5924            &down,
5925            &pair_w_d,
5926            &tok_off_d,
5927            &tok_ids_d,
5928            &mut moe_out,
5929            t,
5930            n_embd,
5931        )?;
5932
5933        if std::env::var("MEMRA_MOE_STATS").is_ok() {
5934            let mut sizes: Vec<usize> = by_expert
5935                .iter()
5936                .filter_map(|pairs| (!pairs.is_empty()).then_some(pairs.len()))
5937                .collect();
5938            sizes.sort_unstable();
5939            let mean = sizes.iter().sum::<usize>() as f64 / sizes.len().max(1) as f64;
5940            println!(
5941                "moe-grouped il={il} t={t} dispatch=resident-q8-clamped-pairs active={}/{} \
5942                 m_e: min={} median={} mean={mean:.1} max={}",
5943                sizes.len(),
5944                n_expert,
5945                sizes.first().copied().unwrap_or(0),
5946                sizes.get(sizes.len() / 2).copied().unwrap_or(0),
5947                sizes.last().copied().unwrap_or(0),
5948            );
5949        }
5950        Ok(moe_out)
5951    }
5952
5953    fn moe_ffn_grouped_add_shared(
5954        e: &Engine,
5955        m: &MoeWeights,
5956        z: &CudaSlice<f32>,
5957        t: usize,
5958        cfg: &ModelConfig,
5959        il: u16,
5960        moe_out: &mut CudaSlice<f32>,
5961    ) -> Result<(), Box<dyn std::error::Error>> {
5962        let n_embd = cfg.n_embd as usize;
5963        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
5964            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
5965        {
5966            let n_ff_sh = gate_shexp.out_features();
5967            let sg_gate = e.matmul(gate_shexp, z, t)?;
5968            let sg_up = e.matmul(up_shexp, z, t)?;
5969            let mut sa = e.uninit(t * n_ff_sh)?;
5970            Self::ffn_act_lim(
5971                e,
5972                cfg,
5973                &sg_gate,
5974                &sg_up,
5975                1.0,
5976                1.0,
5977                cfg.clamp_shexp_at(il as u32),
5978                &mut sa,
5979                t * n_ff_sh,
5980            )?;
5981            let sh = e.matmul(down_shexp, &sa, t)?;
5982            let gate = match &m.gate_inp_shexp {
5983                Some(gate_inp_shexp) => {
5984                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
5985                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
5986                    } else {
5987                        let raw = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
5988                        let mut gate = e.uninit(t)?;
5989                        e.sigmoid(&raw, &mut gate, t)?;
5990                        gate
5991                    }
5992                }
5993                None => e.htod(&vec![1.0f32; t])?,
5994            };
5995            e.add_scaled_rows(&sh, &gate, moe_out, n_embd, t)?;
5996        }
5997        Ok(())
5998    }
5999
6000    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
6001    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
6002    pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
6003                                  cfg: &ModelConfig, il: u16, max_block: usize)
6004                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6005        let moe = cfg.moe.as_ref().unwrap();
6006        let n_embd = cfg.n_embd as usize;
6007        let n_expert = moe.expert_count as usize;
6008        let n_used = moe.expert_used_count as usize;
6009        let n_ff_exp = moe.expert_ff_length as usize;
6010        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
6011        let lim_exp = cfg.clamp_exp_at(il as u32);
6012
6013        // 1. ROUTER: exactly the same m-invariant sigmoid selector as the sequential path
6014        // (device by default; MEMRA_SIG_ROUTER=0 is the host oracle). The grouped dispatch never
6015        // enters the softmax-only pairs/dev router.
6016        let logits = Self::moe_router_logits(e, m, z, t, cfg)?;
6017        if let Some(sig) = cfg.sigmoid_router() {
6018            Self::trace_sigmoid_router_logits(e, il, t, n_expert, n_used, &logits, m, sig)?;
6019        }
6020        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
6021            Self::moe_route_sigmoid_cfg(e, &logits, t, n_expert, n_used, m, sig)?
6022        } else {
6023            Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
6024                                m.active_experts.as_deref())?
6025        };
6026        crate::moesd::record_host_routes(il, n_expert, n_used, &sel_all)?;
6027        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
6028        Self::trace_moe_input(e, il, t, n_embd, z)?;
6029
6030        // Fits-VRAM Step35 under the PP stage split: use the expert-major q8 arithmetic directly
6031        // from the owning device's uniform resident slab. This is not `moe_ffn_pairs`: routing
6032        // already happened above and the activation is clamp-aware. Mixed layouts, remote slabs,
6033        // macro-scaled experts, q8-disabled configs, and spill fall through to metadata-aware A2.
6034        let no_exp_macros = m.gate_exps.macros.is_none()
6035            && m.up_exps.macros.is_none()
6036            && m.down_exps.macros.is_none();
6037        let resident_q8 = m.dev_exps.as_ref().filter(|dev| {
6038            m.has_uniform_expert_layout()
6039                && no_exp_macros
6040                && moe_q8_enabled()
6041                && q8_expert_supported(m.gate_exps.qtype)
6042                && q8_expert_supported(m.up_exps.qtype)
6043                && q8_expert_supported(m.down_exps.qtype)
6044                && moe_slab_enabled()
6045                && dev.dev == e.ctx().ordinal()
6046        });
6047        if let Some(dev) = resident_q8 {
6048            let mut moe_out = Self::moe_ffn_grouped_resident_q8(
6049                e,
6050                m,
6051                z,
6052                t,
6053                cfg,
6054                il,
6055                &sel_all,
6056                &w_all,
6057                &dev.ptr_row,
6058                dev.gu_il,
6059            )?;
6060            Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6061            return Ok(moe_out);
6062        }
6063
6064        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
6065        // For each expert e, we need: which tokens use it, their positions in z, their top-k
6066        // slot index (for bit-identical accumulation), and their weights.
6067        struct ExpertGroup {
6068            tok_indices: Vec<i32>,   // indices into z rows (0..T-1)
6069            slot_indices: Vec<i32>,  // top-k slot (0..n_used-1) for that token-expert pair
6070            weights: Vec<f32>,       // renormalized weight for that token-expert pair
6071        }
6072        let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
6073            tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
6074        }).collect();
6075
6076        for tok in 0..t {
6077            for j in 0..n_used {
6078                let ex = sel_all[tok * n_used + j] as usize;
6079                let w = w_all[tok * n_used + j];
6080                groups[ex].tok_indices.push(tok as i32);
6081                groups[ex].slot_indices.push(j as i32);
6082                groups[ex].weights.push(w);
6083            }
6084        }
6085
6086        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
6087        // Each token's 8 expert contributions land in their respective slots.
6088        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
6089        let mut wbuf = e.zeros(t * n_used)?;  // [T, n_used] weight buffer for FMA reduce
6090
6091        // Expert weight dimensions (used in both cache and staging paths).
6092        let g_len = m.gate_exps.max_expert_bytes();
6093        let u_len = m.up_exps.max_expert_bytes();
6094        let d_len = m.down_exps.max_expert_bytes();
6095        let moe_q8 = m.has_uniform_expert_layout()
6096            && moe_q8_enabled()
6097            && q8_expert_supported(m.gate_exps.qtype)
6098            && q8_expert_supported(m.up_exps.qtype)
6099            && q8_expert_supported(m.down_exps.qtype);
6100        // The per-expert slab view is legal only for the ordinary contiguous uniform layout.
6101        // Interleaved GU slabs require the pointer-table fast path above.
6102        let slab_local = m.dev_exps.as_ref().filter(|dev| {
6103            !dev.gu_il && moe_slab_enabled() && dev.dev == e.ctx().ordinal()
6104        });
6105        let use_cache =
6106            slab_local.is_none() && Engine::moe_cache_enabled() && !e.moe_cache_frozen();
6107        // The sequential no-cache/frozen staging oracle is f32. Use q8 only where sequential
6108        // also does: a local resident slab or a live SLRU dispatch.
6109        let grouped_q8 = moe_q8 && (slab_local.is_some() || use_cache);
6110
6111        // GPU scratch for staging (only allocated without a local slab or cache).
6112        let (mut scratch_g, mut scratch_u, mut scratch_d) = if slab_local.is_none() && !use_cache {
6113            (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
6114        } else {
6115            (None, None, None)
6116        };
6117
6118        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
6119        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
6120        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
6121        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
6122        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
6123        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
6124        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
6125        // at long prompts where every expert stages regardless. Order is FREE to change without
6126        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
6127        // regardless of expert processing order (the whole point of the slots).
6128        let mut order: Vec<usize> =
6129            (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
6130        order.sort_by(|&a, &b| groups[b].tok_indices.len()
6131            .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
6132        let mut m_dist: Vec<usize> = Vec::new();  // for stats
6133        let page_window = moe_page_prefetch_window();
6134        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
6135        if worker_disk_prefetch {
6136            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
6137                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
6138            }
6139        }
6140        for (order_pos, &ex) in order.iter().enumerate() {
6141            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
6142                Self::moe_prefetch_host_expert(order[next], m);
6143            }
6144            if worker_disk_prefetch {
6145                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
6146                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6147                    let keep = [
6148                        BlockId::new(il, PROJ_GATE, ex as u16),
6149                        BlockId::new(il, PROJ_UP, ex as u16),
6150                        BlockId::new(il, PROJ_DOWN, ex as u16),
6151                    ];
6152                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
6153                }
6154            }
6155            let grp = &groups[ex];
6156            let m_e = grp.tok_indices.len();
6157            m_dist.push(m_e);
6158            let gl = m.gate_exps.expert_layout(ex);
6159            let ul = m.up_exps.expert_layout(ex);
6160            let dl = m.down_exps.expert_layout(ex);
6161
6162            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
6163            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
6164            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
6165            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
6166            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
6167            let dmac = m.down_exps.macro_scale(ex);
6168            let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
6169                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
6170                e.htod(&scaled)?
6171            };
6172
6173            // GATHER: collect m_e activation rows from z into a contiguous buffer.
6174            let mut gathered = e.zeros(m_e * n_embd)?;
6175            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
6176            let gv = gathered.slice(0..m_e * n_embd);
6177
6178            // Compute gate/up/down from the local slab, metadata-aware cache, or staging. The q8
6179            // form keeps each gathered row in the sequential arithmetic class at `m=m_e`.
6180            let y = if let Some(dev) = slab_local {
6181                let gate_start = ex * m.gate_exps.expert_stride;
6182                let up_start = ex * m.up_exps.expert_stride;
6183                let down_start = ex * m.down_exps.expert_stride;
6184                if grouped_q8 {
6185                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6186                    let gate = e.qmatvec_expert_q8(
6187                        &dev.gate,
6188                        gate_start..gate_start + gl.len,
6189                        &zq,
6190                        &zd,
6191                        m_e,
6192                        m.gate_exps.in_f,
6193                        m.gate_exps.out_f,
6194                        gl.qtype,
6195                        gl.row_bytes,
6196                    )?;
6197                    let up = e.qmatvec_expert_q8(
6198                        &dev.up,
6199                        up_start..up_start + ul.len,
6200                        &zq,
6201                        &zd,
6202                        m_e,
6203                        m.up_exps.in_f,
6204                        m.up_exps.out_f,
6205                        ul.qtype,
6206                        ul.row_bytes,
6207                    )?;
6208                    let mut act = e.uninit(m_e * n_ff_exp)?;
6209                    Self::ffn_act_lim(
6210                        e,
6211                        cfg,
6212                        &gate,
6213                        &up,
6214                        m.gate_exps.macro_scale(ex),
6215                        m.up_exps.macro_scale(ex),
6216                        lim_exp,
6217                        &mut act,
6218                        m_e * n_ff_exp,
6219                    )?;
6220                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6221                    e.qmatvec_expert_q8(
6222                        &dev.down,
6223                        down_start..down_start + dl.len,
6224                        &aq2,
6225                        &ad2,
6226                        m_e,
6227                        m.down_exps.in_f,
6228                        m.down_exps.out_f,
6229                        dl.qtype,
6230                        dl.row_bytes,
6231                    )?
6232                } else {
6233                    let gate = e.qmatvec_view(
6234                        &dev.gate,
6235                        gate_start..gate_start + gl.len,
6236                        &gv,
6237                        m_e,
6238                        m.gate_exps.in_f,
6239                        m.gate_exps.out_f,
6240                        gl.qtype,
6241                        gl.row_bytes,
6242                    )?;
6243                    let up = e.qmatvec_view(
6244                        &dev.up,
6245                        up_start..up_start + ul.len,
6246                        &gv,
6247                        m_e,
6248                        m.up_exps.in_f,
6249                        m.up_exps.out_f,
6250                        ul.qtype,
6251                        ul.row_bytes,
6252                    )?;
6253                    let mut act = e.uninit(m_e * n_ff_exp)?;
6254                    Self::ffn_act_lim(
6255                        e,
6256                        cfg,
6257                        &gate,
6258                        &up,
6259                        m.gate_exps.macro_scale(ex),
6260                        m.up_exps.macro_scale(ex),
6261                        lim_exp,
6262                        &mut act,
6263                        m_e * n_ff_exp,
6264                    )?;
6265                    let actv = act.slice(0..m_e * n_ff_exp);
6266                    e.qmatvec_view(
6267                        &dev.down,
6268                        down_start..down_start + dl.len,
6269                        &actv,
6270                        m_e,
6271                        m.down_exps.in_f,
6272                        m.down_exps.out_f,
6273                        dl.qtype,
6274                        dl.row_bytes,
6275                    )?
6276                }
6277            } else if use_cache {
6278                use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6279                if grouped_q8 {
6280                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6281                    let gate = e.with_moe_cache(max_block, |cache, eng| {
6282                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
6283                        let slot =
6284                            cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
6285                        eng.qmatvec_expert_q8(
6286                            cache.buf(slot),
6287                            0..gl.len,
6288                            &zq,
6289                            &zd,
6290                            m_e,
6291                            m.gate_exps.in_f,
6292                            m.gate_exps.out_f,
6293                            gl.qtype,
6294                            gl.row_bytes,
6295                        )
6296                    })?;
6297                    let up = e.with_moe_cache(max_block, |cache, eng| {
6298                        let id = BlockId::new(il, PROJ_UP, ex as u16);
6299                        let slot =
6300                            cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
6301                        eng.qmatvec_expert_q8(
6302                            cache.buf(slot),
6303                            0..ul.len,
6304                            &zq,
6305                            &zd,
6306                            m_e,
6307                            m.up_exps.in_f,
6308                            m.up_exps.out_f,
6309                            ul.qtype,
6310                            ul.row_bytes,
6311                        )
6312                    })?;
6313                    let mut act = e.uninit(m_e * n_ff_exp)?;
6314                    Self::ffn_act_lim(
6315                        e,
6316                        cfg,
6317                        &gate,
6318                        &up,
6319                        m.gate_exps.macro_scale(ex),
6320                        m.up_exps.macro_scale(ex),
6321                        lim_exp,
6322                        &mut act,
6323                        m_e * n_ff_exp,
6324                    )?;
6325                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6326                    e.with_moe_cache(max_block, |cache, eng| {
6327                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
6328                        let slot =
6329                            cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
6330                        eng.qmatvec_expert_q8(
6331                            cache.buf(slot),
6332                            0..dl.len,
6333                            &aq2,
6334                            &ad2,
6335                            m_e,
6336                            m.down_exps.in_f,
6337                            m.down_exps.out_f,
6338                            dl.qtype,
6339                            dl.row_bytes,
6340                        )
6341                    })?
6342                } else {
6343                    let gate = e.with_moe_cache(max_block, |cache, eng| {
6344                        let id = BlockId::new(il, PROJ_GATE, ex as u16);
6345                        let slot =
6346                            cache.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
6347                        eng.qmatvec_view(
6348                            cache.buf(slot),
6349                            0..gl.len,
6350                            &gv,
6351                            m_e,
6352                            m.gate_exps.in_f,
6353                            m.gate_exps.out_f,
6354                            gl.qtype,
6355                            gl.row_bytes,
6356                        )
6357                    })?;
6358                    let up = e.with_moe_cache(max_block, |cache, eng| {
6359                        let id = BlockId::new(il, PROJ_UP, ex as u16);
6360                        let slot =
6361                            cache.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
6362                        eng.qmatvec_view(
6363                            cache.buf(slot),
6364                            0..ul.len,
6365                            &gv,
6366                            m_e,
6367                            m.up_exps.in_f,
6368                            m.up_exps.out_f,
6369                            ul.qtype,
6370                            ul.row_bytes,
6371                        )
6372                    })?;
6373                    let mut act = e.uninit(m_e * n_ff_exp)?;
6374                    Self::ffn_act_lim(
6375                        e,
6376                        cfg,
6377                        &gate,
6378                        &up,
6379                        m.gate_exps.macro_scale(ex),
6380                        m.up_exps.macro_scale(ex),
6381                        lim_exp,
6382                        &mut act,
6383                        m_e * n_ff_exp,
6384                    )?;
6385                    let actv = act.slice(0..m_e * n_ff_exp);
6386                    e.with_moe_cache(max_block, |cache, eng| {
6387                        let id = BlockId::new(il, PROJ_DOWN, ex as u16);
6388                        let slot =
6389                            cache.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
6390                        eng.qmatvec_view(
6391                            cache.buf(slot),
6392                            0..dl.len,
6393                            &actv,
6394                            m_e,
6395                            m.down_exps.in_f,
6396                            m.down_exps.out_f,
6397                            dl.qtype,
6398                            dl.row_bytes,
6399                        )
6400                    })?
6401                }
6402            } else {
6403                let sg = scratch_g.as_mut().unwrap();
6404                let su = scratch_u.as_mut().unwrap();
6405                let sd = scratch_d.as_mut().unwrap();
6406                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
6407                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
6408                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
6409                if grouped_q8 {
6410                    let (zq, zd) = e.quantize_q8_1(&gathered, m_e, n_embd)?;
6411                    let gate = e.qmatvec_expert_q8(
6412                        sg,
6413                        0..gl.len,
6414                        &zq,
6415                        &zd,
6416                        m_e,
6417                        m.gate_exps.in_f,
6418                        m.gate_exps.out_f,
6419                        gl.qtype,
6420                        gl.row_bytes,
6421                    )?;
6422                    let up = e.qmatvec_expert_q8(
6423                        su,
6424                        0..ul.len,
6425                        &zq,
6426                        &zd,
6427                        m_e,
6428                        m.up_exps.in_f,
6429                        m.up_exps.out_f,
6430                        ul.qtype,
6431                        ul.row_bytes,
6432                    )?;
6433                    let mut act = e.uninit(m_e * n_ff_exp)?;
6434                    Self::ffn_act_lim(
6435                        e,
6436                        cfg,
6437                        &gate,
6438                        &up,
6439                        m.gate_exps.macro_scale(ex),
6440                        m.up_exps.macro_scale(ex),
6441                        lim_exp,
6442                        &mut act,
6443                        m_e * n_ff_exp,
6444                    )?;
6445                    let (aq2, ad2) = e.quantize_q8_1(&act, m_e, n_ff_exp)?;
6446                    e.qmatvec_expert_q8(
6447                        sd,
6448                        0..dl.len,
6449                        &aq2,
6450                        &ad2,
6451                        m_e,
6452                        m.down_exps.in_f,
6453                        m.down_exps.out_f,
6454                        dl.qtype,
6455                        dl.row_bytes,
6456                    )?
6457                } else {
6458                    let gate = e.qmatvec_view(
6459                        sg,
6460                        0..gl.len,
6461                        &gv,
6462                        m_e,
6463                        m.gate_exps.in_f,
6464                        m.gate_exps.out_f,
6465                        gl.qtype,
6466                        gl.row_bytes,
6467                    )?;
6468                    let up = e.qmatvec_view(
6469                        su,
6470                        0..ul.len,
6471                        &gv,
6472                        m_e,
6473                        m.up_exps.in_f,
6474                        m.up_exps.out_f,
6475                        ul.qtype,
6476                        ul.row_bytes,
6477                    )?;
6478                    let mut act = e.uninit(m_e * n_ff_exp)?;
6479                    Self::ffn_act_lim(
6480                        e,
6481                        cfg,
6482                        &gate,
6483                        &up,
6484                        m.gate_exps.macro_scale(ex),
6485                        m.up_exps.macro_scale(ex),
6486                        lim_exp,
6487                        &mut act,
6488                        m_e * n_ff_exp,
6489                    )?;
6490                    let actv = act.slice(0..m_e * n_ff_exp);
6491                    e.qmatvec_view(
6492                        sd,
6493                        0..dl.len,
6494                        &actv,
6495                        m_e,
6496                        m.down_exps.in_f,
6497                        m.down_exps.out_f,
6498                        dl.qtype,
6499                        dl.row_bytes,
6500                    )?
6501                }
6502            };
6503
6504            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
6505            e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
6506                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6507        }
6508
6509        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
6510        let mut moe_out = e.zeros(t * n_embd)?;
6511        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
6512
6513        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
6514        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
6515            m_dist.sort_unstable();
6516            let active = m_dist.len();
6517            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
6518            let median = m_dist[active / 2];
6519            let max_m = *m_dist.last().unwrap();
6520            let min_m = m_dist[0];
6521            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
6522            println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
6523                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
6524                      above_gemm_threshold(>=16)={above16}/{active}");
6525        }
6526
6527        Self::moe_ffn_grouped_add_shared(e, m, z, t, cfg, il, &mut moe_out)?;
6528        Ok(moe_out)
6529    }
6530
6531    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
6532    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
6533    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
6534    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
6535    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
6536    /// expert-sum order identical to the sequential path.
6537    pub(crate) fn moe_ffn_lockstep(
6538        &self,
6539        e: &Engine,
6540        m: &MoeWeights,
6541        zbatch: &CudaSlice<f32>,
6542        mrows: usize,
6543        il: u16,
6544        max_block: usize,
6545    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6546        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
6547        let cfg = &self.cfg;
6548        let moe = cfg.moe.as_ref().unwrap();
6549        let n_embd = cfg.n_embd as usize;
6550        let n_expert = moe.expert_count as usize;
6551        let n_used = moe.expert_used_count as usize;
6552        let n_ff_exp = moe.expert_ff_length as usize;
6553        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
6554        let lim_exp = cfg.clamp_exp_at(il as u32);
6555        let lim_shexp = cfg.clamp_shexp_at(il as u32);
6556
6557        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
6558        if let Some(sig) = cfg.sigmoid_router() {
6559            Self::trace_sigmoid_router_logits(
6560                e, il, mrows, n_expert, n_used, &logits, m, sig,
6561            )?;
6562        }
6563        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
6564            Self::moe_route_sigmoid_cfg(e, &logits, mrows, n_expert, n_used, m, sig)?
6565        } else {
6566            Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
6567                                m.active_experts.as_deref())?
6568        };
6569        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
6570
6571        // Residency split at whole-expert granularity against the (frozen) cache.
6572        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
6573            Ok((0..n_expert)
6574                .map(|ex| {
6575                    [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
6576                        c.resident(BlockId::new(il, p, ex as u16)).is_some()
6577                    })
6578                })
6579                .collect())
6580        })?;
6581
6582        struct Group {
6583            rows: Vec<i32>,
6584            slots: Vec<i32>,
6585            weights: Vec<f32>,
6586        }
6587        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
6588        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
6589        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
6590            Default::default();
6591        for row in 0..mrows {
6592            for j in 0..n_used {
6593                let ex = sel_all[row * n_used + j] as usize;
6594                let w = w_all[row * n_used + j];
6595                if resident_expert[ex] {
6596                    let group = groups.entry(ex).or_insert_with(|| Group {
6597                        rows: Vec::new(),
6598                        slots: Vec::new(),
6599                        weights: Vec::new(),
6600                    });
6601                    group.rows.push(row as i32);
6602                    group.slots.push(j as i32);
6603                    group.weights.push(w);
6604                } else {
6605                    crate::cpu_experts::record_incomplete_gpu_residency(0);
6606                    cpu_rows[row].push((ex, w));
6607                    cpu_by_expert.entry(ex).or_default().push((row, w));
6608                }
6609            }
6610        }
6611
6612        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
6613        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
6614        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
6615        // order per row differs from the sequential single-call chunk — part of the
6616        // documented lockstep numeric class.
6617        let host_rows = e.dtoh(zbatch)?;
6618        let rows_ok = crate::cpu_experts::rows_supported();
6619        enum CpuPart {
6620            Single { row: usize },
6621            Rows { rows: Vec<usize> },
6622        }
6623        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
6624        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
6625        if rows_ok {
6626            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
6627                .into_iter()
6628                .filter(|(_, rows)| rows.len() >= 2)
6629                .collect();
6630            shared.sort_by_key(|(ex, _)| *ex);
6631            for (ex, mut row_weights) in shared {
6632                row_weights.sort_by_key(|(row, _)| *row);
6633                let inputs: Vec<(&[f32], f32)> = row_weights
6634                    .iter()
6635                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
6636                    .collect();
6637                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
6638                    .map_err(std::io::Error::other)?;
6639                for &(row, _) in &row_weights {
6640                    rows_served.insert((row, ex));
6641                }
6642                tickets.push((
6643                    CpuPart::Rows {
6644                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
6645                    },
6646                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
6647                ));
6648            }
6649        }
6650        for (row, selected) in cpu_rows.iter().enumerate() {
6651            let leftover: Vec<(usize, f32)> = selected
6652                .iter()
6653                .copied()
6654                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
6655                .collect();
6656            if leftover.is_empty() {
6657                continue;
6658            }
6659            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
6660            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
6661                .map_err(std::io::Error::other)?;
6662            tickets.push((
6663                CpuPart::Single { row },
6664                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
6665            ));
6666        }
6667
6668        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
6669        let mut wbuf = e.zeros(mrows * n_used)?;
6670        let mut order: Vec<usize> = groups.keys().copied().collect();
6671        order.sort_by(|&a, &b| {
6672            groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
6673        });
6674        for &ex in &order {
6675            let group = &groups[&ex];
6676            let m_e = group.rows.len();
6677            let gl = m.gate_exps.expert_layout(ex);
6678            let ul = m.up_exps.expert_layout(ex);
6679            let dl = m.down_exps.expert_layout(ex);
6680            let row_idx_d = e.htod_i32(&group.rows)?;
6681            let slot_idx_d = e.htod_i32(&group.slots)?;
6682            let dmac = m.down_exps.macro_scale(ex);
6683            let weight_d = if dmac == 1.0 {
6684                e.htod(&group.weights)?
6685            } else {
6686                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
6687                e.htod(&scaled)?
6688            };
6689            let mut gathered = e.zeros(m_e * n_embd)?;
6690            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
6691            let gv = gathered.slice(0..m_e * n_embd);
6692            let gate = e.with_moe_cache(max_block, |c, eng| {
6693                let slot = c
6694                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
6695                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6696                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
6697                    m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
6698            })?;
6699            let up = e.with_moe_cache(max_block, |c, eng| {
6700                let slot = c
6701                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
6702                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6703                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
6704                    m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
6705            })?;
6706            let mut act = e.zeros(m_e * n_ff_exp)?;
6707            Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
6708                m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
6709            let actv = act.slice(0..m_e * n_ff_exp);
6710            let y = e.with_moe_cache(max_block, |c, eng| {
6711                let slot = c
6712                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
6713                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
6714                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
6715                    m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
6716            })?;
6717            e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
6718                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
6719        }
6720        let mut moe_out = e.zeros(mrows * n_embd)?;
6721        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
6722
6723        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
6724        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
6725        for (part, ticket) in tickets {
6726            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
6727            let mut add_row = |row: usize, chunk: &[f32]| {
6728                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
6729                for (accumulator, value) in sum.iter_mut().zip(chunk) {
6730                    *accumulator += value;
6731                }
6732            };
6733            match part {
6734                CpuPart::Single { row } => add_row(row, &cpu_output),
6735                CpuPart::Rows { rows } => {
6736                    for (slot, row) in rows.into_iter().enumerate() {
6737                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
6738                    }
6739                }
6740            }
6741        }
6742        for (row, sum) in row_sums.into_iter().enumerate() {
6743            let Some(sum) = sum else { continue };
6744            let cpu_output = e.htod(&sum)?;
6745            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
6746            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
6747        }
6748
6749        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
6750            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
6751        {
6752            let n_ff_sh = gate_shexp.out_features();
6753            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
6754            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
6755            let mut sa = e.zeros(mrows * n_ff_sh)?;
6756            Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp,
6757                              &mut sa, mrows * n_ff_sh)?;
6758            let sh = e.matmul(down_shexp, &sa, mrows)?;
6759            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
6760            // decode matches the single-sequence decode chain bit-for-bit.
6761            let g = match &m.gate_inp_shexp {
6762                Some(gate_inp_shexp) => {
6763                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
6764                }
6765                None => e.htod(&vec![1.0f32; mrows])?,
6766            };
6767            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
6768        }
6769
6770        Ok(moe_out)
6771    }
6772}
6773
6774// ============================ gemma4 (R8 verified wiring) ==================================
6775// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
6776// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
6777// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
6778// gemma variants after the correctness gate).
6779impl HybridModel {
6780    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
6781    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
6782        let g = self.cfg.gemma4.as_ref().unwrap();
6783        let swa = g.swa_pattern[il];
6784        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
6785        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
6786        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
6787        // rows exact (softmax over one element) while every later position drifted).
6788        (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
6789         if swa { g.rope_base_swa } else { g.rope_base_global },
6790         1.0, swa)
6791    }
6792
6793    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
6794    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
6795    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
6796    fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
6797                       -> Result<(), Box<dyn std::error::Error>> {
6798        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
6799            // suppress_d is primary-owned by the audited invariant (serving picks the last/head
6800            // stage as primary, and this tail runs only after the last stage). The assert turns
6801            // that argued invariant into a checked one: any topology violating primary==head
6802            // trips here in debug instead of silently peer-reading a device-0 buffer.
6803            #[cfg(debug_assertions)]
6804            crate::debug_assert_tensor_stream_device(ids, &e.stream(),
6805                                                     "gemma4_suppress.suppress_d");
6806            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
6807        }
6808        Ok(())
6809    }
6810
6811    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
6812    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
6813    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
6814    /// only (v0): attends within `tokens` via the f32 sdpa.
6815    fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6816                         h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
6817                         cache: Option<&mut Cache>)
6818                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6819        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6820        let eps = self.cfg.rms_eps;
6821        let aux = self.gemma4_aux.as_ref().unwrap();
6822        let ones = aux.ones(e);
6823        #[cfg(debug_assertions)]
6824        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
6825                                                   "gemma4_attn_prime.ones");
6826
6827        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
6828        // (h stays borrowed across the triple, so the cache key can't go stale).
6829        e.mmq_act_begin();
6830        let q0 = e.matmul(&fa.wq, h, t)?;   // [t, nh*hd]
6831        let k0 = e.matmul(&fa.wk, h, t)?;   // [t, nkv*hd]
6832        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
6833        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
6834        let v0 = if swa { e.matmul(&fa.wv, h, t)? } else { e.clone_dtod(&k0)? };
6835
6836        let mut q = e.uninit(t * nh * hd)?;
6837        let mut k = e.uninit(t * nkv * hd)?;
6838        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
6839        let mut v = e.uninit(t * nkv * hd)?;
6840        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
6841        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
6842        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
6843        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6844        let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
6845            && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
6846        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
6847        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6848        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
6849        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
6850        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
6851        let v_f16 = emit && crate::fa_f16pv_on() && match hd {
6852            512 => true,
6853            256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
6854            _ => false,
6855        };
6856        if emit {
6857            e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6858                               ones, &mut q, &mut k, &mut v, &mut vb,
6859                               hd, nh * t, nkv * t, eps, v_f16)?;
6860        } else {
6861            e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6862                           ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
6863        }
6864
6865        let ff = if swa { None } else {
6866            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
6867        };
6868        #[cfg(debug_assertions)]
6869        if let Some(ff) = ff {
6870            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
6871                                                       "gemma4_attn_prime.rope_freqs");
6872        }
6873        if emit {
6874            e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
6875                               base, 1.0, ff)?;
6876        } else {
6877            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
6878        }
6879
6880        if let Some(cache) = cache {
6881            let kvl = cache.kv[il].as_mut().unwrap();
6882            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
6883            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
6884                                       kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
6885            kvl.len += t;
6886        }
6887        let mut attn = e.zeros(t * nh * hd)?;
6888        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
6889        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
6890        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
6891        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6892        if swa && t > win {
6893            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6894                if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6895                                             scale, true, win, v_f16)?; }
6896                else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
6897                                      win)?; }
6898            } else {
6899                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
6900            }
6901        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
6902            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6903        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
6904            if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
6905                                             scale, true, v_f16)?; }
6906            else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
6907        } else {
6908            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6909        }
6910        Ok(e.matmul(&fa.wo, &attn, t)?)
6911    }
6912
6913    /// Back-compat wrapper (pure prefill, no cache).
6914    fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6915                   h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
6916                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6917        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
6918    }
6919
6920    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
6921    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
6922    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
6923    /// the q8z epilogue is quantize_q8_1 verbatim).
6924    fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6925                     bits: &crate::hybrid::Gemma4MoeBits,
6926                     mq: &(CudaSlice<i8>, CudaSlice<f32>),
6927                     router_in: &CudaSlice<f32>, t: usize)
6928                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6929        let cfg = &self.cfg;
6930        let moe = cfg.moe.as_ref().unwrap();
6931        let n_embd = cfg.n_embd as usize;
6932        let n_expert = moe.expert_count as usize;
6933        let n_used = moe.expert_used_count as usize;
6934        let n_ff_exp = moe.expert_ff_length as usize;
6935        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
6936        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
6937        // the pair's 12us is kernel time, not launch gaps.
6938        let logits = if crate::router_kernel_on() {
6939            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
6940        } else {
6941            e.matmul(&m.gate_inp, router_in, t)?
6942        };
6943        let dev = m.dev_exps.as_ref().unwrap();
6944        let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
6945                                                    &bits.per_expert_scale_d)?;
6946        let (zq, zd) = mq;
6947        if t == 1 {
6948            let selv = sel_d.slice(0..n_used);
6949            let wv = w_d.slice(0..n_used);
6950            let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
6951                                                 n_embd, n_ff_exp, n_used, n_expert,
6952                                                 m.gate_exps.qtype, m.up_exps.qtype,
6953                                                 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
6954            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
6955            let mut moe_out = e.uninit(n_embd)?;
6956            e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
6957                                   &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
6958                                   n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
6959            return Ok(moe_out);
6960        }
6961        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
6962        let act = if csr {
6963            e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
6964                                           n_embd, n_ff_exp, n_used, n_expert,
6965                                           m.gate_exps.qtype, m.up_exps.qtype,
6966                                           m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6967        } else {
6968            e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
6969                                            n_embd, n_ff_exp, n_used, n_expert,
6970                                            m.gate_exps.qtype, m.up_exps.qtype,
6971                                            m.gate_exps.row_bytes, m.up_exps.row_bytes)?
6972        };
6973        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
6974        let mut moe_out = e.uninit(t * n_embd)?;
6975        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
6976        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
6977        e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
6978                                      n_ff_exp, n_embd, n_used, n_expert,
6979                                      m.down_exps.qtype, m.down_exps.row_bytes)?;
6980        Ok(moe_out)
6981    }
6982
6983    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
6984    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
6985    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
6986    fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
6987                  bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
6988                  router_in: &CudaSlice<f32>, t: usize)
6989                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6990        let cfg = &self.cfg;
6991        let moe = cfg.moe.as_ref().unwrap();
6992        let n_embd = cfg.n_embd as usize;
6993        let n_expert = moe.expert_count as usize;
6994        let n_used = moe.expert_used_count as usize;
6995        let n_ff_exp = moe.expert_ff_length as usize;
6996
6997        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
6998        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
6999        // batched matmul only at real prefill.
7000        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
7001            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
7002        } else {
7003            e.matmul(&m.gate_inp, router_in, t)?
7004        };
7005
7006        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
7007        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
7008        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
7009        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
7010        if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7011            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
7012            && expert_dp4a_supported(m.down_exps.qtype)
7013            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
7014            let dev = m.dev_exps.as_ref().unwrap();
7015            let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
7016                                                        &bits.per_expert_scale_d)?;
7017            if t == 1 {
7018                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
7019                let selv = sel_d.slice(0..n_used);
7020                let wv = w_d.slice(0..n_used);
7021                let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
7022                                                     n_embd, n_ff_exp, n_used, n_expert,
7023                                                     m.gate_exps.qtype, m.up_exps.qtype,
7024                                                     m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
7025                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
7026                let mut moe_out = e.uninit(n_embd)?;
7027                e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
7028                                       &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
7029                                       n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
7030                return Ok(moe_out);
7031            }
7032            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
7033            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
7034            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
7035            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
7036            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
7037            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
7038            let act = if csr {
7039                e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
7040                                               n_embd, n_ff_exp, n_used, n_expert,
7041                                               m.gate_exps.qtype, m.up_exps.qtype,
7042                                               m.gate_exps.row_bytes, m.up_exps.row_bytes)?
7043            } else {
7044                e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
7045                                                n_embd, n_ff_exp, n_used, n_expert,
7046                                                m.gate_exps.qtype, m.up_exps.qtype,
7047                                                m.gate_exps.row_bytes, m.up_exps.row_bytes)?
7048            };
7049            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
7050            let mut moe_out = e.uninit(t * n_embd)?;
7051            e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
7052                                          n_ff_exp, n_embd, n_used, n_expert,
7053                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
7054            return Ok(moe_out);
7055        }
7056
7057        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
7058        for (i, &sx) in sel_all.iter().enumerate() {
7059            w_all[i] *= bits.per_expert_scale[sx as usize];
7060        }
7061
7062        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
7063        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
7064        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
7065        if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7066            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
7067            && expert_dp4a_supported(m.down_exps.qtype)
7068            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
7069            let dev = m.dev_exps.as_ref().unwrap();
7070            let n_pairs = t * n_used;
7071            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
7072            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
7073            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
7074            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
7075            let pt = e.htod_i32(&pair_tok)?;
7076            let pw = e.htod(&w_all)?;
7077            let toff = e.htod_i32(&tok_off)?;
7078            let tids = e.htod_i32(&tok_ids)?;
7079            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
7080            for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
7081            let mut ex_ids: Vec<i32> = Vec::new();
7082            let mut ex_off: Vec<i32> = vec![0];
7083            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
7084            for (ex, list) in by_ex.iter().enumerate() {
7085                if list.is_empty() { continue; }
7086                ex_ids.push(ex as i32);
7087                ex_pairs.extend_from_slice(list);
7088                ex_off.push(ex_pairs.len() as i32);
7089            }
7090            let n_active = ex_ids.len();
7091            let exi = e.htod_i32(&ex_ids)?;
7092            let exo = e.htod_i32(&ex_off)?;
7093            let exp_d = e.htod_i32(&ex_pairs)?;
7094            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
7095            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
7096            // end-to-end (gelu is elementwise), one row permute before the scatter. The
7097            // ragged down k (704) needs no padding here — cublas takes any k.
7098            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
7099            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
7100            // Hopper default — see moe_f16g_gemma_on.
7101            if crate::moe_f16g_gemma_on()
7102                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
7103                && f16g_proj_ok(m.up_exps.qtype, n_embd)
7104                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
7105                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
7106                let csr_tok_d = e.htod_i32(&csr_tok)?;
7107                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
7108                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
7109                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
7110                                              m.gate_exps.qtype, m.gate_exps.row_bytes)?;
7111                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
7112                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
7113                                              m.up_exps.qtype, m.up_exps.row_bytes)?;
7114                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
7115                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
7116                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
7117                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
7118                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
7119                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
7120                let mut moe_out = e.uninit(t * n_embd)?;
7121                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
7122                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
7123                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
7124                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
7125                    eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
7126                              scan(&yd), scan(&mo));
7127                }
7128                return Ok(moe_out);
7129            }
7130            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
7131            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
7132            let mma = n_embd % 256 == 0
7133                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
7134            let (gate, up) = if mma {
7135                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
7136                (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
7137                                  n_embd, n_ff_exp, n_active, n_pairs, t,
7138                                  m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7139                 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
7140                                  n_embd, n_ff_exp, n_active, n_pairs, t,
7141                                  m.up_exps.qtype, m.up_exps.row_bytes)?)
7142            } else {
7143                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
7144                (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
7145                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
7146                                           m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7147                 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
7148                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
7149                                           m.up_exps.qtype, m.up_exps.row_bytes)?)
7150            };
7151            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
7152            let pself = e.htod_i32(&pair_self)?;
7153            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
7154            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
7155            // to the 256-val superblock (768) while the act quantizer's zero padding
7156            // makes every padded-k product exactly zero (weight overread bytes multiply
7157            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
7158            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
7159            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
7160            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
7161            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
7162            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
7163            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
7164            let y_down = if mma {
7165                let in_pad = n_ff_exp.div_ceil(256) * 256;
7166                let a_scr = if crate::moe_fuse_actq_on() {
7167                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
7168                } else {
7169                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
7170                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
7171                };
7172                e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
7173                                 in_pad, n_embd, n_active, n_pairs, n_pairs,
7174                                 m.down_exps.qtype, m.down_exps.row_bytes)?
7175            } else {
7176                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
7177                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
7178                e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
7179                                          n_ff_exp, n_embd, n_expert, n_active, n_pairs,
7180                                          m.down_exps.qtype, m.down_exps.row_bytes)?
7181            };
7182            let mut moe_out = e.uninit(t * n_embd)?;
7183            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
7184            return Ok(moe_out);
7185        }
7186
7187        let g_len = m.gate_exps.expert_stride;
7188        let u_len = m.up_exps.expert_stride;
7189        let d_len = m.down_exps.expert_stride;
7190        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
7191        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
7192        // the spill fallback.
7193        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
7194        let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
7195            (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
7196        };
7197        let mut moe_out = e.zeros(t * n_embd)?;
7198        for tok in 0..t {
7199            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
7200            let w = &w_all[tok * n_used..(tok + 1) * n_used];
7201            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
7202            for (j, &ex) in sel.iter().enumerate() {
7203                let ex = ex as usize;
7204                let gate = match dev {
7205                    Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
7206                        m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
7207                    None => {
7208                        let sg = sg.as_mut().unwrap();
7209                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
7210                        e.qmatvec_view(sg, 0..g_len, &zt, 1,
7211                            m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
7212                    }
7213                };
7214                let up = match dev {
7215                    Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
7216                        m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
7217                    None => {
7218                        let su = su.as_mut().unwrap();
7219                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
7220                        e.qmatvec_view(su, 0..u_len, &zt, 1,
7221                            m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
7222                    }
7223                };
7224                let mut act = e.uninit(n_ff_exp)?;
7225                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
7226                let actv = act.slice(0..n_ff_exp);
7227                let y = match dev {
7228                    Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
7229                        m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
7230                    None => {
7231                        let sd = sd.as_mut().unwrap();
7232                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
7233                        e.qmatvec_view(sd, 0..d_len, &actv, 1,
7234                            m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
7235                    }
7236                };
7237                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
7238                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
7239            }
7240        }
7241        Ok(moe_out)
7242    }
7243
7244    /// One gemma4 trunk layer (R8): x -> x_next.
7245    fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
7246                    x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
7247                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7248        let n_embd = self.cfg.n_embd as usize;
7249        let eps = self.cfg.rms_eps;
7250
7251        let mut h = e.zeros(t * n_embd)?;
7252        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7253        let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7254        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
7255        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
7256        let mut cur = e.zeros(t * n_embd)?;
7257        e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
7258        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
7259    }
7260
7261    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
7262    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
7263    /// layer scale — shared verbatim by the prefill, decode and verify paths.
7264    fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7265                             cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
7266                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7267        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
7268    }
7269
7270    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
7271    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
7272    fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7273                               cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7274                               next_norm: Option<&CudaSlice<f32>>)
7275                               -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
7276        let n_embd = self.cfg.n_embd as usize;
7277        let bits = layer.gemma4.as_ref().unwrap();
7278        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
7279        let mut xn = e.uninit(t * n_embd)?;
7280        match next_norm {
7281            Some(w) => {
7282                let mut hn = e.uninit(t * n_embd)?;
7283                e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
7284                                     n_embd, t, self.cfg.rms_eps)?;
7285                Ok((xn, Some(hn)))
7286            }
7287            None => {
7288                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
7289                Ok((xn, None))
7290            }
7291        }
7292    }
7293
7294    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
7295    /// norm — returns (sn, attn_out) for the closing add+scale variants.
7296    fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7297                              cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
7298                              -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7299        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
7300    }
7301
7302    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
7303    /// means `cur` is the RAW attention output and the dense entry runs
7304    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
7305    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
7306    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
7307    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
7308    fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7309                                 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7310                                 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
7311                                 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7312        let n_embd = self.cfg.n_embd as usize;
7313        let eps = self.cfg.rms_eps;
7314        let bits = layer.gemma4.as_ref().unwrap();
7315
7316        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
7317        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
7318        let Some(mbits) = bits.moe_bits.as_ref() else {
7319            let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
7320            else { panic!("gemma4 dense layer without Dense ffn") };
7321            let mut attn_out = e.uninit(t * n_embd)?;
7322            let mut zsh = e.uninit(t * n_embd)?;
7323            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
7324            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
7325            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7326            match pre_norm {
7327                Some(wa) if t == 1 => {
7328                    zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
7329                                                            bits.ffn_norm.float_data(),
7330                                                            &mut attn_out, &mut zsh,
7331                                                            n_embd, t, eps)?);
7332                }
7333                Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
7334                                                   &mut attn_out, &mut zsh, n_embd, t, eps)?,
7335                None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
7336                                       &mut zsh, n_embd, t, eps)?,
7337            }
7338            let n_ff = ffn_gate.out_features();
7339            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
7340            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
7341            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
7342            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
7343            // rescue segment C — the megakernel front is closed for the dense tail.
7344            let (gate, up) = if t == 1 {
7345                let (zq, zd) = match zpair {
7346                    Some(p) => p,
7347                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
7348                };
7349                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
7350                    Some(p) => p,
7351                    None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
7352                             e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
7353                }
7354            } else {
7355                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
7356                // launch for the verify's gate+up — the up segment's blocks fill SMs as
7357                // the gate segment drains (the launch-tail mechanism behind the b-tier
7358                // plateau; first positive after six falsified in-kernel variants).
7359                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7360                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
7361                let fused = if f2b {
7362                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
7363                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
7364                } else { None };
7365                match fused {
7366                    Some(p) => p,
7367                    None => {
7368                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
7369                        e.mmq_act_begin();
7370                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
7371                    }
7372                }
7373            };
7374            let mut act = e.uninit(t * n_ff)?;
7375            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
7376            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
7377            let f0 = if e.uses_q8_1_fast(ffn_down) {
7378                let upv = e.view(&up, t * n_ff);
7379                let up_all = upv.slice(0..t * n_ff);
7380                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
7381                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
7382            } else {
7383                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
7384                e.matmul(ffn_down, &act, t)?
7385            };
7386            if defer_post_norm { return Ok((f0, attn_out)); }
7387            let mut sn = e.uninit(t * n_embd)?;
7388            e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
7389            return Ok((sn, attn_out));
7390        };
7391
7392        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
7393        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
7394        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
7395        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
7396        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
7397        let mut attn_out = e.uninit(t * n_embd)?;
7398        let mut router_in = e.uninit(t * n_embd)?;
7399        let fast_moe = match &layer.ffn {
7400            crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
7401                && expert_dp4a_supported(m.gate_exps.qtype)
7402                && expert_dp4a_supported(m.up_exps.qtype)
7403                && expert_dp4a_supported(m.down_exps.qtype)
7404                && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
7405            _ => false,
7406        };
7407        let q8z = t < PRIME_MIN_T && fast_moe;
7408        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
7409            let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
7410                                               &mbits.router_scale_pre,
7411                                               mbits.pre_ffw_norm_2.float_data(),
7412                                               &mut attn_out, &mut router_in, n_embd, t, eps)?;
7413            (None, Some(z0), Some(m2))
7414        } else {
7415            let mut zsh = e.uninit(t * n_embd)?;
7416            let mut moe_in = e.uninit(t * n_embd)?;
7417            e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
7418                            mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
7419                            &mut router_in, &mut moe_in, n_embd, t, eps)?;
7420            (Some((zsh, moe_in)), None, None)
7421        };
7422        let attn_out2 = attn_out;
7423        #[allow(unused_variables)]
7424        let attn_out = &attn_out2;
7425        let n_ff = mbits.shared_gate.out_features();
7426        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
7427            if t == 1 {
7428                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
7429                    Some(p) => p,
7430                    None => {
7431                        let h0 = e.zeros(0)?;
7432                        (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
7433                         e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
7434                    }
7435                }
7436            } else {
7437                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
7438                let h0 = e.zeros(0)?;
7439                (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
7440                 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
7441            }
7442        } else {
7443            let (zsh, _) = zsh_f32.as_ref().unwrap();
7444            (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
7445        };
7446        let mut act = e.uninit(t * n_ff)?;
7447        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
7448        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
7449        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
7450        let moe0 = match (&moe_q8, &zsh_f32) {
7451            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
7452            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
7453            _ => unreachable!(),
7454        };
7455        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
7456        let mut mlp = e.uninit(t * n_embd)?;
7457        let mut moe = e.uninit(t * n_embd)?;
7458        e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
7459                     mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
7460
7461        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
7462        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
7463        let mut sum = e.uninit(t * n_embd)?;
7464        let mut sn = e.uninit(t * n_embd)?;
7465        e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
7466                       n_embd, t, eps)?;
7467        Ok((sn, attn_out2))
7468    }
7469
7470    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
7471    fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7472                                cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
7473                                next_norm: Option<&CudaSlice<f32>>)
7474                                -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
7475        let n_embd = self.cfg.n_embd as usize;
7476        let bits = layer.gemma4.as_ref().unwrap();
7477        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
7478        let mut xn = e.uninit(t * n_embd)?;
7479        match next_norm {
7480            Some(w) => {
7481                let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
7482                                                     n_embd, t, self.cfg.rms_eps)?;
7483                Ok((xn, Some(pair)))
7484            }
7485            None => {
7486                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
7487                Ok((xn, None))
7488            }
7489        }
7490    }
7491
7492    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
7493    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
7494    fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7495                      -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7496        // E4B routes to its own forward regardless of the caller's entry point (forward /
7497        // forward_last / prime paths all funnel here for gemma4).
7498        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
7499        let n_embd = self.cfg.n_embd as usize;
7500        let t = tokens.len();
7501        let pos: Vec<i32> = (0..t as i32).collect();
7502        let pos_d = e.htod_i32(&pos)?;
7503
7504        let mut x = self.embed(e, tokens)?;
7505        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7506        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
7507        // the bring-up bisect vs llama-eval-callback node stats.
7508        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
7509        let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
7510            let h = e.dtoh(x)?;
7511            let bad = h.iter().filter(|v| !v.is_finite()).count();
7512            let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
7513            eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
7514            Ok(())
7515        };
7516        if probe { stat(e, &x, "embed")?; }
7517        for (il, layer) in self.layers.iter().enumerate() {
7518            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
7519            if probe { stat(e, &x, &format!("L{il}"))?; }
7520        }
7521        let mut hn = e.zeros(t * n_embd)?;
7522        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
7523        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7524        let n_vocab = self.output.out_features();
7525        let logits = if last_only {
7526            let hv = e.view(&hn, t * n_embd);
7527            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
7528            let mut hlast = e.zeros(n_embd)?;
7529            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
7530            let mut ld = e.matmul(&self.output, &hlast, 1)?;
7531            e.softcap(&mut ld, cap, n_vocab)?;
7532            self.gemma4_suppress(e, &mut ld, 1)?;
7533            e.dtoh(&ld)?
7534        } else {
7535            let mut ld = e.matmul(&self.output, &hn, t)?;
7536            e.softcap(&mut ld, cap, t * n_vocab)?;
7537            self.gemma4_suppress(e, &mut ld, t)?;
7538            e.dtoh(&ld)?
7539        };
7540        Ok(logits)
7541    }
7542
7543    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
7544    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
7545    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
7546    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
7547    pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7548                               -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7549        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
7550        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
7551        // whole worker process on this line. The worker now primes gemma4 monolithically and
7552        // routes continuation suffixes tokenwise; this is the per-request backstop.
7553        if cache.pos != 0 {
7554            return Err("gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
7555                        — prime the full prompt in one call or decode tokenwise".into());
7556        }
7557        let n_embd = self.cfg.n_embd as usize;
7558        let eps = self.cfg.rms_eps;
7559        let t = tokens.len();
7560        let pos: Vec<i32> = (0..t as i32).collect();
7561        let pos_d = e.htod_i32(&pos)?;
7562        let mut x = self.embed(e, tokens)?;
7563        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7564        for (il, layer) in self.layers.iter().enumerate() {
7565            let mut h = e.zeros(t * n_embd)?;
7566            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7567            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
7568            let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
7569            let mut cur = e.zeros(t * n_embd)?;
7570            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
7571            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
7572            self.dflash_tap(e, cache, il, &x, t)?;
7573        }
7574        cache.pos += t;
7575        let hiddens = e.clone_dtod(&x)?;
7576        let xv = e.view(&x, t * n_embd);
7577        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
7578        let mut h_seed = e.zeros(n_embd)?;
7579        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
7580        let mut hn = e.uninit(n_embd)?;
7581        e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7582        let mut ld = e.matmul(&self.output, &hn, 1)?;
7583        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7584        e.softcap(&mut ld, cap, self.output.out_features())?;
7585        self.gemma4_suppress(e, &mut ld, 1)?;
7586        let logits = e.dtoh(&ld)?;
7587        Ok((logits, h_seed, hiddens))
7588    }
7589
7590    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
7591    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
7592    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
7593    /// fused norm emits q8 directly — the f32 h never materializes).
7594    fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
7595                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7596                          pos_d: &CudaSlice<i32>, cache: &mut Cache)
7597                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7598        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7599        let eps = self.cfg.rms_eps;
7600        let aux = self.gemma4_aux.as_ref().unwrap();
7601        let ones = aux.ones(e);
7602        #[cfg(debug_assertions)]
7603        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
7604                                                   "gemma4_decode_attn.ones");
7605        let (hq, hdq) = (hq, hdq);
7606        let h0 = e.zeros(0)?;
7607        let h = &h0;
7608        let (q0, k0, v0) = if swa {
7609            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
7610                Some(t3) => t3,
7611                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7612                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
7613                         e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
7614            }
7615        } else {
7616            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
7617                Some(p) => p,
7618                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7619                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
7620            };
7621            let v0 = e.clone_dtod(&k0)?;
7622            (q0, k0, v0)
7623        };
7624        let mut q = e.uninit(nh * hd)?;
7625        let mut k = e.uninit(nkv * hd)?;
7626        let mut v = e.uninit(nkv * hd)?;
7627        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
7628        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
7629        let ff = if swa { None } else {
7630            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
7631        };
7632        #[cfg(debug_assertions)]
7633        if let Some(ff) = ff {
7634            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
7635                                                       "gemma4_decode_attn.rope_freqs");
7636        }
7637        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7638                            ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7639                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
7640        let kvl = cache.kv[il].as_mut().unwrap();
7641        e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
7642                              kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
7643        kvl.len += 1;
7644        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
7645        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
7646        // positional). Globals attend the full history.
7647        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7648        let mut attn = e.uninit(nh * hd)?;
7649        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
7650        if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
7651            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7652            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7653            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7654            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
7655            let base = kvl.len as i32;
7656            e.i32_set_k(&mut kvl.len_d, base)?;
7657            e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
7658                             kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
7659                             false, None)?;
7660            return Ok(e.matmul(&fa.wo, &attn, 1)?);
7661        }
7662        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
7663        if swa && kvl.len > win && hd == 256
7664            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
7665            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
7666            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
7667            let base = kvl.len as i32;
7668            e.i32_set_k(&mut kvl.len_d, base)?;
7669            e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
7670                               win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
7671            return Ok(e.matmul(&fa.wo, &attn, 1)?);
7672        }
7673        let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
7674        let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7675                                     (off_tok + t_kv) * kvl.k_tok_bytes);
7676        let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7677                                     (off_tok + t_kv) * kvl.v_tok_bytes);
7678        e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7679                    kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
7680        Ok(e.matmul(&fa.wo, &attn, 1)?)
7681    }
7682
7683    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
7684    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
7685    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
7686    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
7687    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
7688    /// in-graph; the driver gates).
7689    #[allow(clippy::too_many_arguments)]
7690    pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7691                                 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7692                                 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7693                                 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
7694                                 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7695        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7696        self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
7697                                        n_vocab, cap_bucket_max, &mut tok_out)?;
7698        Ok(tok_out)
7699    }
7700
7701    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
7702    /// every replay; pass `token_d` itself for the self-feeding graph loop).
7703    #[allow(clippy::too_many_arguments)]
7704    pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
7705                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7706                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7707                                      n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7708                                      tok_out: &mut CudaSlice<u32>)
7709                                      -> Result<(), Box<dyn std::error::Error>> {
7710        let n_embd = self.cfg.n_embd as usize;
7711        let eps = self.cfg.rms_eps;
7712        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7713        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7714        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7715        let n_layers = self.layers.len();
7716        for (il, layer) in self.layers.iter().enumerate() {
7717            let (hq, hdq) = match h_carry.take() {
7718                Some(p) => p,
7719                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
7720            };
7721            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7722            let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
7723            let mut cur = e.uninit(n_embd)?;
7724            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
7725            let next_norm = if il + 1 < n_layers {
7726                Some(self.layers[il + 1].attn_norm.float_data())
7727            } else { None };
7728            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
7729            x = xn;
7730            h_carry = hn;
7731        }
7732        let mut hn = e.uninit(n_embd)?;
7733        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7734        let mut logits = e.matmul(&self.output, &hn, 1)?;
7735        self.gemma4_suppress(e, &mut logits, 1)?;   // cap skipped (monotonic); the mask is not
7736        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
7737        e.inc_seqlen(pos_d)?;
7738        if cap_bucket_max.is_none() { cache.pos += 1; }
7739        Ok(())
7740    }
7741
7742    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
7743    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
7744    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
7745    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
7746
7747    /// Build the slot set (call OUTSIDE any capture).
7748    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
7749        let n_embd = self.cfg.n_embd as usize;
7750        let n_vocab = self.output.out_features();
7751        let n_layers = self.layers.len();
7752        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
7753        for il in 0..n_layers {
7754            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
7755            qmax = qmax.max(nh * hd);
7756            kvmax = kvmax.max(nkv * hd);
7757            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
7758                ffmax = ffmax.max(ffn_gate.out_features());
7759            }
7760        }
7761        Ok(G4DcSlots {
7762            x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
7763            hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
7764            q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
7765            q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
7766            attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
7767            attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
7768            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
7769            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
7770            zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
7771            gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
7772            act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
7773            f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
7774            hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
7775        })
7776    }
7777
7778    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
7779    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
7780    fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
7781                         aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
7782                         -> Result<(), Box<dyn std::error::Error>> {
7783        use crate::model::GpuTensor;
7784        let (bytes, qtype, row_bytes, scale, rp) = match w {
7785            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
7786                (bytes, *qtype, *row_bytes, *scale, *rp),
7787            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
7788        };
7789        let (mbytes, mrp) = match w {
7790            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
7791            _ => (bytes, rp),
7792        };
7793        e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
7794                            qtype, row_bytes, scale, mrp, y)
7795    }
7796
7797    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
7798    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
7799    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
7800    #[allow(clippy::too_many_arguments)]
7801    pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
7802                                         pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7803                                         embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7804                                         n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
7805                                         sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
7806                                         ring: Option<(&mut CudaSlice<u32>, usize)>)
7807                                         -> Result<(), Box<dyn std::error::Error>> {
7808        let n_embd = self.cfg.n_embd as usize;
7809        let eps = self.cfg.rms_eps;
7810        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
7811        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
7812        let n_layers = self.layers.len();
7813        let mut has_carry = false;
7814        for il in 0..n_layers {
7815            if !has_carry {
7816                e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
7817                                     eps, &mut sl.hq, &mut sl.hd_)?;
7818            }
7819            has_carry = true;
7820            let layer = &self.layers[il];
7821            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
7822            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
7823            e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
7824            let next_norm = if il + 1 < n_layers {
7825                Some(self.layers[il + 1].attn_norm.float_data())
7826            } else { None };
7827            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
7828            std::mem::swap(&mut sl.x, &mut sl.xn);
7829        }
7830        e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
7831        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7832        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
7833        {
7834            let (zq, zd) = (&sl.zq, &sl.zd);
7835            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
7836            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
7837            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
7838        }
7839        self.gemma4_suppress(e, &mut sl.logits, 1)?;
7840        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
7841        if let Some((ring, base)) = ring {
7842            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
7843            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
7844            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
7845            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
7846        }
7847        e.inc_seqlen(pos_d)?;
7848        if cap_bucket_max.is_none() { cache.pos += 1; }
7849        Ok(())
7850    }
7851
7852    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
7853    #[allow(clippy::too_many_arguments)]
7854    fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
7855                                     il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
7856                                     cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
7857                                     -> Result<(), Box<dyn std::error::Error>> {
7858        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
7859        let eps = self.cfg.rms_eps;
7860        let aux = self.gemma4_aux.as_ref().unwrap();
7861        let ones = aux.ones(e);
7862        #[cfg(debug_assertions)]
7863        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
7864                                                   "gemma4_decode_attn_dc_slotted.ones");
7865        {
7866            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
7867            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
7868            if swa {
7869                if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
7870                                            &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
7871                    return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
7872                }
7873            } else {
7874                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
7875                    return Err("slotted step: fused2 unavailable".into());
7876                }
7877                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
7878                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
7879            }
7880        }
7881        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
7882        // kernel-for-kernel (graph stream-identity gate).
7883        let ff = if swa { None } else {
7884            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
7885        };
7886        #[cfg(debug_assertions)]
7887        if let Some(ff) = ff {
7888            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
7889                                                       "gemma4_decode_attn_dc_slotted.rope_freqs");
7890        }
7891        let kvl = cache.kv[il].as_mut().unwrap();
7892        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7893        if crate::Engine::qkv_append_on() {
7894            // append fold (2026-07-23): mirrors dc_into.
7895            e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
7896                fa.k_norm.float_data(), ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7897                pos_d, nh, nkv, base, 1.0, ff, eps,
7898                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
7899        } else {
7900            e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7901                                ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
7902                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
7903            e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
7904                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
7905                                     kv_fp8)?;
7906        }
7907        e.inc_seqlen(&mut kvl.len_d)?;
7908        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
7909        let k_view = e.view_u8(&kvl.k, kvl.k.len());
7910        let v_view = e.view_u8(&kvl.v, kvl.v.len());
7911        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
7912        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7913        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
7914        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
7915        // the dc_into arm branch-for-branch (stream gate).
7916        let mut fa_q8 = false;
7917        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
7918            e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
7919                             1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7920                             Some((&kvl.len_d, -1)), false, false,
7921                             Some((&mut sl.zq, &mut sl.zd)))?;
7922            fa_q8 = true;
7923        } else if swa && b_swa > win && hd == 256 && rows_on {
7924            e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
7925                               &kvl.len_d, -1, 1, scale, win,
7926                               kvl.k_tok_bytes, kvl.v_tok_bytes,
7927                               Some((&mut sl.zq, &mut sl.zd)))?;
7928            fa_q8 = true;
7929        } else {
7930            let b = if swa { b_swa } else { b_glob };
7931            e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
7932                           scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
7933                           swa && crate::Engine::wkv_on())?;
7934        }
7935        if !fa_q8 {
7936            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
7937            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
7938        }
7939        {
7940            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7941            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7942            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
7943        }
7944        Ok(())
7945    }
7946
7947    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
7948    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
7949    fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
7950                                 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
7951                                 -> Result<(), Box<dyn std::error::Error>> {
7952        let n_embd = self.cfg.n_embd as usize;
7953        let eps = self.cfg.rms_eps;
7954        let bits = layer.gemma4.as_ref().unwrap();
7955        let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
7956        else { return Err("slotted tail: dense ffn only".into()) };
7957        e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
7958                       &mut sl.zsh, n_embd, 1, eps)?;
7959        let n_ff = ffn_gate.out_features();
7960        {
7961            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
7962            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
7963        }
7964        {
7965            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
7966            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
7967            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
7968                return Err("slotted tail: ffn fused2 unavailable".into());
7969            }
7970        }
7971        debug_assert!(e.uses_q8_1_fast(ffn_down));
7972        {
7973            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
7974            let upv = e.view(upr, n_ff);
7975            let up_all = upv.slice(0..n_ff);
7976            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
7977            e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
7978                                      &mut sl.actq, &mut sl.actd)?;
7979        }
7980        {
7981            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
7982            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
7983            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
7984        }
7985        e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
7986        match next_norm {
7987            Some(w) => {
7988                e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
7989                                               &mut sl.xn, n_embd, 1, eps,
7990                                               &mut sl.hq, &mut sl.hd_)?;
7991            }
7992            None => {
7993                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
7994            }
7995        }
7996        Ok(())
7997    }
7998
7999    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
8000    #[allow(clippy::too_many_arguments)]
8001    fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8002                             hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8003                             pos_d: &CudaSlice<i32>, cache: &mut Cache,
8004                             cap_bucket_max: Option<(usize, usize)>)
8005                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8006        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8007        let eps = self.cfg.rms_eps;
8008        let aux = self.gemma4_aux.as_ref().unwrap();
8009        let ones = aux.ones(e);
8010        #[cfg(debug_assertions)]
8011        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8012                                                   "gemma4_decode_attn_dc.ones");
8013        let (q0, k0, v0) = if swa {
8014            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
8015                Some(t3) => t3,
8016                None => {
8017                    let h0 = e.zeros(0)?;
8018                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
8019                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
8020                     e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
8021                }
8022            }
8023        } else {
8024            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
8025                Some(p) => p,
8026                None => {
8027                    let h0 = e.zeros(0)?;
8028                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
8029                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
8030                }
8031            };
8032            let v0 = e.clone_dtod(&k0)?;
8033            (q0, k0, v0)
8034        };
8035        let mut q = e.uninit(nh * hd)?;
8036        let mut k = e.uninit(nkv * hd)?;
8037        let mut v = e.uninit(nkv * hd)?;
8038        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
8039        let ff = if swa { None } else {
8040            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8041        };
8042        #[cfg(debug_assertions)]
8043        if let Some(ff) = ff {
8044            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8045                                                       "gemma4_decode_attn_dc.rope_freqs");
8046        }
8047        let kvl = cache.kv[il].as_mut().unwrap();
8048        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
8049        if crate::Engine::qkv_append_on() {
8050            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
8051            e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
8052                fa.k_norm.float_data(), ones, &mut q, &mut k, &mut v, hd, nh, nkv,
8053                pos_d, nh, nkv, base, 1.0, ff, eps,
8054                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
8055        } else {
8056            e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8057                                ones, &mut q, &mut k, &mut v, hd, nh, nkv,
8058                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
8059            e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
8060                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
8061        }
8062        e.inc_seqlen(&mut kvl.len_d)?;
8063        let mut attn = e.uninit(nh * hd)?;
8064        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
8065        // rides g4_matvec_m1_into instead of matmul's internal quantize.
8066        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8067        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
8068        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
8069        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
8070        // (gemma4_e4b_attn, +0.65% valid window).
8071        match cap_bucket_max {
8072            None => {
8073                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
8074                // decode (SWA layers attend the last `sliding_window` keys); the device
8075                // counters carry only the append slot + the graph seam.
8076                kvl.len += 1;
8077                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8078                if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
8079                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8080                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
8081                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
8082                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
8083                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
8084                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8085                    e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
8086                                     scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8087                                     Some((&kvl.len_d, -1)), false, false,
8088                                     Some((&mut aq8, &mut ad8)))?;
8089                    fa_q8 = Some((aq8, ad8));
8090                } else if swa && kvl.len > win && hd == 256
8091                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8092                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
8093                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
8094                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
8095                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8096                    e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
8097                                       1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
8098                                       Some((&mut aq8, &mut ad8)))?;
8099                    fa_q8 = Some((aq8, ad8));
8100                } else {
8101                    let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
8102                                          else { (0, kvl.len) };
8103                    let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
8104                                                 (off_tok + t_kv) * kvl.k_tok_bytes);
8105                    let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
8106                                                 (off_tok + t_kv) * kvl.v_tok_bytes);
8107                    e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
8108                                kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
8109                }
8110            }
8111            Some((b_swa, b_glob)) => {
8112                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
8113                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
8114                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
8115                // the RUNG max for the rows family (kernels derive per-replay splits from
8116                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
8117                let k_view = e.view_u8(&kvl.k, kvl.k.len());
8118                let v_view = e.view_u8(&kvl.v, kvl.v.len());
8119                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
8120                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8121                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
8122                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8123                    e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
8124                                     1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8125                                     Some((&kvl.len_d, -1)), false, false,
8126                                     Some((&mut aq8, &mut ad8)))?;
8127                    fa_q8 = Some((aq8, ad8));
8128                } else if swa && b_swa > win && hd == 256 && rows_on {
8129                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
8130                    e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8131                                       &kvl.len_d, -1, 1, scale, win,
8132                                       kvl.k_tok_bytes, kvl.v_tok_bytes,
8133                                       Some((&mut aq8, &mut ad8)))?;
8134                    fa_q8 = Some((aq8, ad8));
8135                } else {
8136                    let b = if swa { b_swa } else { b_glob };
8137                    e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
8138                                   scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8139                                   swa && crate::Engine::wkv_on())?;
8140                }
8141            }
8142        }
8143        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
8144        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
8145        if let Some((aq8, ad8)) = fa_q8 {
8146            let mut y = e.uninit(fa.wo.out_features())?;
8147            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
8148            return Ok(y);
8149        }
8150        Ok(e.matmul(&fa.wo, &attn, 1)?)
8151    }
8152
8153    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
8154    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
8155    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
8156    /// views in-graph); caller gates and falls back to the dc-eager loop.
8157    pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
8158                                 cache: &mut Cache, max_new: usize, eos: &[u32],
8159                                 mut on_token: impl FnMut(u32) -> bool)
8160                                 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
8161        if self.is_gemma4_e4b() {
8162            return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
8163        }
8164        use crate::decode::StopReason;
8165        let n_vocab = self.output.out_features();
8166        let n_embd = self.cfg.n_embd as usize;
8167        let embd_gpu = self.embd_gpu.get_or_init(|| {
8168            e.upload_u8(&self.embd.raw).expect("embed table upload")
8169        });
8170        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8171        for kvl in cache.kv.iter_mut().flatten() {
8172            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8173        }
8174        let mut token_d = e.stream().clone_htod(&[first_token])?;
8175        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
8176        let g4 = self.cfg.gemma4.as_ref().unwrap();
8177        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
8178        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
8179        let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
8180            .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
8181        let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
8182            .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
8183        let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
8184                                                  (cudarc::driver::CudaGraph,
8185                                                   Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
8186        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
8187        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
8188        let mut slots = self.g4_dc_slots(e)?;
8189        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
8190        // baked at the door entry (the modulo keeps every capture valid indefinitely).
8191        const RING: usize = 64;
8192        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
8193        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
8194        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
8195        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
8196        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
8197        const DRAIN: usize = 1;
8198        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
8199        let ring_base = prompt_pos;
8200        let mut out = Vec::with_capacity(max_new);
8201        let mut reason = StopReason::MaxNew;
8202        let mut next = first_token;
8203        let mut captures = 0usize;
8204        for _ in 0..max_new {
8205            out.push(next);
8206            if eos.contains(&next) { reason = StopReason::Eos; break; }
8207            if !on_token(next) { reason = StopReason::Callback; break; }
8208            let t_kv = cache.pos + 1;
8209            // Bucket key per ARM (graph arc step 3):
8210            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
8211            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
8212            //    the component collapses to a single marker).
8213            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
8214            //    at/above it — the kernel derives splits from len_d per replay, so buckets
8215            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
8216            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8217            let f512 = crate::fa512_min_tkv();
8218            let key_s = if t_kv > win { (true, usize::MAX) }
8219                        else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
8220            let (key_g, rung_end) = if t_kv >= f512 {
8221                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
8222                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
8223                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
8224                ((true, end), end)
8225            } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
8226            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
8227            if !graphs.contains_key(&key) {
8228                let bucket_max = (t_kv, rung_end);
8229                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
8230                let snap = cache.snapshot(e)?;
8231                let pos_save = e.dtoh_i32_one(&pos_d)?;
8232                let len_save: Vec<Option<i32>> = cache.kv.iter()
8233                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
8234                let tok_save = e.dtoh_u32_one(&token_d)?;
8235                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
8236                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
8237                // regression class, and this door's measured -8.8%. The keeper pins warmup
8238                // transients so the captured graph holds kernel nodes only.
8239                let graph = {
8240                    let tok_ref = &mut token_d;
8241                    let pos_ref = &mut pos_d;
8242                    let cache_ref = &mut *cache;
8243                    let slots_ref = &mut slots;
8244                    let ring_ref = &mut ring;
8245                    e.capture_graph_retained_flags(
8246                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
8247                        |e| {
8248                        // self-feeding: the argmax writes token_d itself.
8249                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
8250                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
8251                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
8252                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
8253                                                           cache_ref, n_vocab, Some(bucket_max),
8254                                                           sl, tok_ref, Some((rg, ring_base)))
8255                    })?
8256                };
8257                cache.rollback(e, &snap, 0)?;
8258                e.set_i32_one(&mut pos_d, pos_save)?;
8259                for (il, ls) in len_save.iter().enumerate() {
8260                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
8261                        e.set_i32_one(&mut kvl.len_d, *v)?;
8262                    }
8263                }
8264                e.set_u32_one(&mut token_d, tok_save)?;
8265                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
8266                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
8267                        eprintln!("[graph-census] {c:?}");
8268                    }
8269                }
8270                graphs.insert(key, graph);
8271                captures += 1;
8272            }
8273            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
8274            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
8275            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
8276            // the budget; capture warmups already emitted their tokens through the ring.
8277            let mut chunk = 1usize;
8278            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
8279                .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
8280            while chunk < drain_cap && out.len() + chunk < max_new {
8281                let t_next = cache.pos + 1 + chunk;
8282                let key_s2 = if t_next > win { (true, usize::MAX) }
8283                             else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
8284                let key_g2 = if t_next >= f512 {
8285                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
8286                } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
8287                if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
8288                chunk += 1;
8289            }
8290            let g = &graphs.get(&key).unwrap().0;
8291            for _ in 0..chunk { g.launch()?; }
8292            e.stream().synchronize()?;
8293            let ringh = e.dtoh_u32(&ring)?;
8294            for j in 0..chunk {
8295                let pos_j = cache.pos + j;
8296                let tok_j = ringh[(pos_j - ring_base) % RING];
8297                cache.pos += 0; // advanced below in one shot
8298                if j + 1 == chunk { next = tok_j; }
8299                else {
8300                    out.push(tok_j);
8301                    if eos.contains(&tok_j) || !on_token(tok_j) {
8302                        reason = if eos.contains(&tok_j) { StopReason::Eos }
8303                                 else { StopReason::Callback };
8304                        // roll device/host state back to the stop point.
8305                        let keep = cache.pos + j + 1;
8306                        e.set_i32_one(&mut pos_d, keep as i32)?;
8307                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
8308                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
8309                            kvl.len = keep;
8310                        }
8311                        cache.pos = keep;
8312                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
8313                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
8314                        }
8315                        return Ok((out, reason));
8316                    }
8317                }
8318            }
8319            cache.pos += chunk;
8320            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
8321        }
8322        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
8323            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
8324        }
8325        Ok((out, reason))
8326    }
8327
8328    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
8329    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
8330    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
8331    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
8332    /// logits (host) + advances cache.pos by t.
8333    pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
8334                                       cache: &mut Cache)
8335                                       -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8336        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
8337    }
8338
8339    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
8340    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
8341    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
8342    pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
8343                                          cache: &mut Cache)
8344                                          -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8345        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
8346        let t = tokens.len();
8347        let n_vocab = self.output.out_features();
8348        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
8349        for i in 0..t {
8350            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
8351        }
8352        Ok((e.dtoh_u32(&toks)?, hn))
8353    }
8354
8355    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
8356    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
8357    pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
8358                                              pos0: usize, cache: &mut Cache)
8359                                              -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8360        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
8361        let n_vocab = self.output.out_features();
8362        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
8363        for i in 0..t {
8364            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
8365        }
8366        Ok((vam, hn))
8367    }
8368
8369    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
8370    /// llama's h_nextn convention).
8371    pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
8372                                         cache: &mut Cache)
8373                                         -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8374        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
8375        let t = tokens.len();
8376        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8377        e.softcap(&mut ld, cap, t * self.output.out_features())?;
8378        Ok((e.dtoh(&ld)?, hn))
8379    }
8380
8381    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
8382    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
8383    pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
8384                                        -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
8385        Ok(VerifyStreamScratch {
8386            pos_d: e.htod_i32(&vec![0i32; cap])?,
8387            row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
8388        })
8389    }
8390
8391    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
8392    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
8393    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
8394    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
8395    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
8396    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
8397    /// sync, exactly the turnaround the burst exists to remove.
8398    pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
8399                                            ctr: &CudaSlice<i32>, hint: usize,
8400                                            cache: &mut Cache,
8401                                            scr: &mut VerifyStreamScratch)
8402                                            -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8403        let n_embd = self.cfg.n_embd as usize;
8404        let eps = self.cfg.rms_eps;
8405        assert!(t <= scr.row_ctrs.len() && t <= 64);
8406        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
8407        for i in 0..t {
8408            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
8409        }
8410        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
8411        let embd_gpu = self.embd_gpu.get_or_init(|| {
8412            e.upload_u8(&self.embd.raw).expect("embed table upload")
8413        });
8414        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8415        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
8416        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
8417        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8418        let n_layers = self.layers.len();
8419        for (il, layer) in self.layers.iter().enumerate() {
8420            let (hq, hdq) = match h_carry.take() {
8421                Some(p) => p,
8422                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8423            };
8424            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8425            let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
8426                                                    hint, row_ctrs)?;
8427            let mut cur = e.uninit(t * n_embd)?;
8428            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8429            let next_norm = if il + 1 < n_layers {
8430                Some(self.layers[il + 1].attn_norm.float_data())
8431            } else { None };
8432            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8433            x = xn;
8434            h_carry = hn;
8435            self.dflash_tap(e, cache, il, &x, t)?;
8436        }
8437        let mut hn = e.uninit(t * n_embd)?;
8438        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8439        let ld = e.matmul(&self.output, &hn, t)?;
8440        let n_vocab = self.output.out_features();
8441        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
8442        for i in 0..t {
8443            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
8444        }
8445        Ok((vam, hn))
8446    }
8447
8448    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
8449    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
8450    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
8451    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
8452    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
8453    /// kernel later if it shows in the profile).
8454    fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
8455                  -> Result<(), Box<dyn std::error::Error>> {
8456        let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
8457        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
8458        let h = taps.hidden;
8459        let n_taps = taps.layer_ids.len();
8460        debug_assert_eq!(taps.t, t);
8461        let xv = e.view(x, t * h);
8462        for r in 0..t {
8463            let row = xv.slice(r * h..(r + 1) * h);
8464            e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
8465        }
8466        Ok(())
8467    }
8468
8469    fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
8470                           tok_dev: Option<&CudaSlice<u32>>)
8471                           -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8472        let n_embd = self.cfg.n_embd as usize;
8473        let eps = self.cfg.rms_eps;
8474        let t = tokens.len();
8475        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8476        let pos_d = e.htod_i32(&pos)?;
8477        let mut x = match tok_dev {
8478            Some(td) => {
8479                let embd_gpu = self.embd_gpu.get_or_init(|| {
8480                    e.upload_u8(&self.embd.raw).expect("embed table upload")
8481                });
8482                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
8483                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
8484            }
8485            None => e.htod(&self.embd.gather(n_embd, tokens))?,
8486        };
8487        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
8488        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8489        let n_layers = self.layers.len();
8490        for (il, layer) in self.layers.iter().enumerate() {
8491            let (hq, hdq) = match h_carry.take() {
8492                Some(p) => p,
8493                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
8494            };
8495            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8496            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
8497            let mut cur = e.uninit(t * n_embd)?;
8498            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
8499            let next_norm = if il + 1 < n_layers {
8500                Some(self.layers[il + 1].attn_norm.float_data())
8501            } else { None };
8502            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
8503            x = xn;
8504            h_carry = hn;
8505            self.dflash_tap(e, cache, il, &x, t)?;
8506        }
8507        let mut hn = e.uninit(t * n_embd)?;
8508        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8509        let mut ld = e.matmul(&self.output, &hn, t)?;
8510        self.gemma4_suppress(e, &mut ld, t)?;   // before the per-row argmax consumers
8511        cache.pos += t;
8512        Ok((ld, hn))
8513    }
8514
8515    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
8516    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
8517    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
8518    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
8519    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
8520    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
8521    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
8522    #[allow(clippy::too_many_arguments)]
8523    fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8524                                 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8525                                 pos_d: &CudaSlice<i32>, t: usize,
8526                                 cache: &mut Cache, hint: usize,
8527                                 row_ctrs: &[CudaSlice<i32>])
8528                                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8529        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8530        let eps = self.cfg.rms_eps;
8531        let aux = self.gemma4_aux.as_ref().unwrap();
8532        let ones = aux.ones(e);
8533        #[cfg(debug_assertions)]
8534        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8535                                                   "gemma4_verify_attn_stream.ones");
8536        let h0 = e.zeros(0)?;
8537        let h = &h0;
8538        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
8539        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
8540        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8541        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8542        let fused_qkv = if f2b {
8543            if swa {
8544                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8545                    .map(|(a, b, c)| (a, b, Some(c)))
8546            } else {
8547                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8548                    .map(|(a, b)| (a, b, None))
8549            }
8550        } else { None };
8551        let (q0, k0, v0) = match fused_qkv {
8552            Some((a, b, cv)) => {
8553                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8554                (a, b, v)
8555            }
8556            None => {
8557                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8558                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8559                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8560                         else { e.clone_dtod(&k0)? };
8561                (q0, k0, v0)
8562            }
8563        };
8564        let mut q = e.uninit(t * nh * hd)?;
8565        let mut k = e.uninit(t * nkv * hd)?;
8566        let mut v = e.uninit(t * nkv * hd)?;
8567        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
8568        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
8569        let ff = if swa { None } else {
8570            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8571        };
8572        #[cfg(debug_assertions)]
8573        if let Some(ff) = ff {
8574            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8575                                                       "gemma4_verify_attn_stream.rope_freqs");
8576        }
8577        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8578                            ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8579                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
8580        let kvl = cache.kv[il].as_mut().unwrap();
8581        // append at the DEVICE slot; the counter advances by t on-device.
8582        e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
8583                                      kvl.kv_dim_k, kvl.kv_dim_v,
8584                                      kvl.k_tok_bytes, kvl.v_tok_bytes,
8585                                      (!swa && crate::Engine::gkv_on())
8586                                          || (swa && crate::Engine::wkv_on()))?;
8587        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
8588        // the sole len writer after this round's attention (base stays = old len, plus = 0).
8589        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8590        let mut attn = e.uninit(t * nh * hd)?;
8591        let k_view = e.view_u8(&kvl.k, kvl.k.len());
8592        let v_view = e.view_u8(&kvl.v, kvl.v.len());
8593        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
8594        // and a stable window regime — the same rung/regime keys as the draft graph).
8595        if swa && hint + 1 >= win {
8596            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
8597            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
8598            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8599                               &kvl.len_d, 0, t, scale, win,
8600                               kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8601        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
8602            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
8603            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
8604            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
8605            // Burst entry gates the horizon onto one side of the crossover, so hint decides
8606            // for every row.
8607            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
8608            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
8609            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
8610            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
8611            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
8612            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
8613            // any bucket >= the live length is exact.
8614            let bucket = (hint + t + 2).next_power_of_two()
8615                .min(crate::fa512_min_tkv().saturating_sub(1));
8616            let qv = e.view(&q, t * nh * hd);
8617            for i in 0..t {
8618                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
8619                let mut q_one = e.uninit(nh * hd)?;
8620                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8621                let mut a_one = e.uninit(nh * hd)?;
8622                e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
8623                               &row_ctrs[i], bucket, scale,
8624                               kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
8625                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8626            }
8627        } else if hd == 512 {
8628            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
8629            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
8630            e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
8631                             kvl.k_tok_bytes, kvl.v_tok_bytes,
8632                             Some((&kvl.len_d, 0)), false, false, None)?;
8633        } else {
8634            // hd256 under-window: v4 device-len rows twin.
8635            e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8636                                &kvl.len_d, hint + t, t, scale,
8637                                kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8638                                swa && crate::Engine::wkv_on())?;
8639        }
8640        Ok(e.matmul(&fa.wo, &attn, t)?)
8641    }
8642
8643    fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
8644                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
8645                          pos_d: &CudaSlice<i32>, t: usize,
8646                          cache: &mut Cache)
8647                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8648        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
8649        let eps = self.cfg.rms_eps;
8650        let aux = self.gemma4_aux.as_ref().unwrap();
8651        let ones = aux.ones(e);
8652        #[cfg(debug_assertions)]
8653        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
8654                                                   "gemma4_verify_attn.ones");
8655        let n_embd = self.cfg.n_embd as usize;
8656        let _ = n_embd;
8657
8658        let h0 = e.zeros(0)?;
8659        let h = &h0;
8660        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
8661        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
8662        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8663        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
8664        let fused_qkv = if f2b {
8665            if swa {
8666                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
8667                    .map(|(a, b, c)| (a, b, Some(c)))
8668            } else {
8669                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
8670                    .map(|(a, b)| (a, b, None))
8671            }
8672        } else { None };
8673        let (q0, k0, v0) = match fused_qkv {
8674            Some((a, b, cv)) => {
8675                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
8676                (a, b, v)
8677            }
8678            None => {
8679                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
8680                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
8681                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
8682                         else { e.clone_dtod(&k0)? };
8683                (q0, k0, v0)
8684            }
8685        };
8686        let mut q = e.uninit(t * nh * hd)?;
8687        let mut k = e.uninit(t * nkv * hd)?;
8688        let mut v = e.uninit(t * nkv * hd)?;
8689        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
8690        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
8691        let ff = if swa { None } else {
8692            Some(aux.rope_freqs(e).expect("gemma4 global rope needs rope_freqs.weight"))
8693        };
8694        #[cfg(debug_assertions)]
8695        if let Some(ff) = ff {
8696            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
8697                                                       "gemma4_verify_attn.rope_freqs");
8698        }
8699        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
8700                            ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
8701                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
8702        let kvl = cache.kv[il].as_mut().unwrap();
8703        let base_len = kvl.len;
8704        e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
8705                                   kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()))?;
8706        kvl.len += t;
8707        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
8708        let mut attn = e.uninit(t * nh * hd)?;
8709        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
8710        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
8711        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
8712            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
8713            // decode rides the SAME symbol at t=1 (parity law).
8714            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
8715        if rows_ok && (!swa || base_len + t <= win) {
8716            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8717            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8718            if hd == 512 {
8719                // device-len twin: sync the counter to the verify base (async arg-store).
8720                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8721                e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
8722                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8723                                 Some((&kvl.len_d, 0)), false,
8724                                 swa && crate::Engine::wkv_on(), None)?;
8725            } else {
8726                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
8727                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
8728                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
8729                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8730                e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
8731                                    &kvl.len_d, base_len + t, t, scale,
8732                                    kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
8733                                    swa && crate::Engine::wkv_on())?;
8734            }
8735            return Ok(e.matmul(&fa.wo, &attn, t)?);
8736        }
8737        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
8738        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
8739        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
8740        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
8741        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
8742        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
8743        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
8744        if hd == 256 && swa && base_len + 1 >= win
8745            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8746            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
8747            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
8748            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
8749            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
8750                               t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8751            return Ok(e.matmul(&fa.wo, &attn, t)?);
8752        }
8753        for i in 0..t {
8754            let avail = base_len + i + 1;
8755            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
8756            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
8757                                         (off_tok + t_kv) * kvl.k_tok_bytes);
8758            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
8759                                         (off_tok + t_kv) * kvl.v_tok_bytes);
8760            let qi = e.view(&q, t * nh * hd);
8761            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
8762            let mut q_one = e.uninit(nh * hd)?;
8763            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
8764            let mut a_one = e.uninit(nh * hd)?;
8765            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
8766            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
8767            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
8768            if swa && avail > win && hd == 256
8769                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8770                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8771                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8772                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8773                e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
8774                                   1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
8775            } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
8776                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
8777                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
8778                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
8779                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
8780                e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
8781                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
8782                                 Some((&kvl.len_d, 0)), false, false, None)?;
8783            } else {
8784                e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
8785                            kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
8786            }
8787            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
8788        }
8789        Ok(e.matmul(&fa.wo, &attn, t)?)
8790    }
8791
8792    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
8793    /// h_seed = pre-output_norm hidden). Advances cache.pos.
8794    pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
8795                                       -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8796        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
8797        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
8798        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
8799        // unsplit rather than guessing a fence.
8800        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
8801            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
8802        }
8803        if crate::pp::pp_cuts(self.layers.len()).is_some() {
8804            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
8805        }
8806        let n_embd = self.cfg.n_embd as usize;
8807        let eps = self.cfg.rms_eps;
8808        let pos_d = e.htod_i32(&[cache.pos as i32])?;
8809        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8810        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8811        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
8812        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
8813        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8814        let n_layers = self.layers.len();
8815        for (il, layer) in self.layers.iter().enumerate() {
8816            let (hq, hdq) = match h_carry.take() {
8817                Some(p) => p,
8818                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
8819            };
8820            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8821            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
8822            let mut cur = e.uninit(n_embd)?;
8823            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8824            let next_norm = if il + 1 < n_layers {
8825                Some(self.layers[il + 1].attn_norm.float_data())
8826            } else { None };
8827            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8828            x = xn;
8829            h_carry = hn;
8830        }
8831        let mut hn = e.uninit(n_embd)?;
8832        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8833        let h_seed = e.clone_dtod(&x)?;
8834        let mut ld = e.matmul(&self.output, &hn, 1)?;
8835        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8836        e.softcap(&mut ld, cap, self.output.out_features())?;   // R4 on device (262k host tanh ~ms/step)
8837        self.gemma4_suppress(e, &mut ld, 1)?;
8838        let logits = e.dtoh(&ld)?;
8839        cache.pos += 1;
8840        Ok((logits, h_seed))
8841    }
8842
8843    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
8844    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
8845    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
8846    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
8847    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
8848    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
8849    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
8850    fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
8851                            pos_d: &CudaSlice<i32>, cache: &mut Cache)
8852                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8853        let n_embd = self.cfg.n_embd as usize;
8854        let eps = self.cfg.rms_eps;
8855        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
8856        for il in lo..hi {
8857            let layer = &self.layers[il];
8858            let (hq, hdq) = match h_carry.take() {
8859                Some(p) => p,
8860                // range head: il == lo — norm against THIS layer's attn_norm.
8861                None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
8862            };
8863            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
8864            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
8865            let mut cur = e.uninit(n_embd)?;
8866            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
8867            let next_norm = if il + 1 < hi {
8868                Some(self.layers[il + 1].attn_norm.float_data())
8869            } else { None };
8870            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
8871            x = xn;
8872            h_carry = hn;
8873        }
8874        Ok(x)
8875    }
8876
8877    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
8878    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
8879    /// boundary handoff — same choreography as the generic arm (decode.rs), same
8880    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
8881    /// stage 1 = layers [split, n) + output_norm + softcapped head.
8882    /// Each stage uploads its own copy of the step's position scalar on its own stream.
8883    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
8884    fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
8885                                -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8886        if crate::pp::pp2_streams_off() {
8887            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
8888        }
8889        let rt = crate::pp::Pp2Rt::get(e)?;
8890        let e0 = rt.engine(0, e);
8891        let e1 = rt.engine(1, e);
8892        let n_embd = self.cfg.n_embd as usize;
8893        let eps = self.cfg.rms_eps;
8894        let pos = cache.pos as i32;
8895
8896        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
8897        let slot = {
8898            let _st0 = rt.enter(0);
8899            let pos_d = e0.htod_i32(&[pos])?;
8900            #[cfg(debug_assertions)]
8901            crate::debug_assert_tensor_stream_device(&pos_d, &e0.stream(),
8902                                                       "gemma4_decode_step_h_pp2.stage0.pos_d");
8903            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
8904            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8905            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
8906            rt.tx(0, &x, n_embd)?
8907        };
8908
8909        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
8910        let _st1 = rt.enter(1);
8911        let pos_d = e1.htod_i32(&[pos])?;
8912        #[cfg(debug_assertions)]
8913        crate::debug_assert_tensor_stream_device(&pos_d, &e1.stream(),
8914                                                   "gemma4_decode_step_h_pp2.stage1.pos_d");
8915        let x = rt.rx(0, slot, n_embd)?;
8916        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
8917
8918        let mut hn = e1.uninit(n_embd)?;
8919        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8920        let h_seed = e1.clone_dtod(&x)?;
8921        let mut ld = e1.matmul(&self.output, &hn, 1)?;
8922        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8923        e1.softcap(&mut ld, cap, self.output.out_features())?;
8924        self.gemma4_suppress(e1, &mut ld, 1)?;
8925        let logits = e1.dtoh(&ld)?;
8926        cache.pos += 1;
8927        Ok((logits, h_seed))
8928    }
8929
8930    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
8931    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
8932    fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
8933                                           split: usize)
8934                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8935        let n_embd = self.cfg.n_embd as usize;
8936        let eps = self.cfg.rms_eps;
8937        let pos_d = e.htod_i32(&[cache.pos as i32])?;
8938
8939        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
8940        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
8941        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
8942        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
8943
8944        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
8945        let boundary_tx = e.clone_dtod(&x)?;
8946        let boundary_rx = e.clone_dtod(&boundary_tx)?;
8947
8948        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
8949        let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
8950
8951        let mut hn = e.uninit(n_embd)?;
8952        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
8953        let h_seed = e.clone_dtod(&x)?;
8954        let mut ld = e.matmul(&self.output, &hn, 1)?;
8955        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
8956        e.softcap(&mut ld, cap, self.output.out_features())?;
8957        self.gemma4_suppress(e, &mut ld, 1)?;
8958        let logits = e.dtoh(&ld)?;
8959        cache.pos += 1;
8960        Ok((logits, h_seed))
8961    }
8962}
8963
8964// ============================ step35 (Step-3.7-Flash) ==================================
8965// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
8966// FAMILY and not a few branches inside the generic `full_attn*` chain:
8967//
8968//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
8969//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
8970//      shapes and the FA head counts would be wrong on 33 of 45 layers.
8971//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
8972//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
8973//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
8974//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
8975//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
8976//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
8977//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
8978//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
8979//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
8980//
8981// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
8982impl HybridModel {
8983    /// Step35's declarative geometry row. Missing rows are artifact/config errors; never
8984    /// synthesize a drafter or trunk layer from a neighboring class.
8985    pub(crate) fn step35_geom(&self, il: usize) -> memra_gguf::config::LayerGeometry {
8986        let geometry = self.cfg.layer_geometry(il as u32)
8987            .unwrap_or_else(|| panic!("step35 layer {il} has no geometry-table row"));
8988        debug_assert_eq!(
8989            geometry.attention_gate,
8990            memra_gguf::config::AttentionGateKind::SeparateHead
8991        );
8992        geometry
8993    }
8994
8995    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
8996    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
8997    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
8998    ///
8999    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input
9000    /// on single-sequence paths. `gt_pre` supplies the already-projected per-head gate for a
9001    /// cross-request batch; exactly one of `hg` or `gt_pre` must be present.
9002    /// `cache`:
9003    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
9004    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
9005    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
9006    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
9007    ///     contract, lane/chunkinv-flip).
9008    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
9009    ///     q/k/v, no cache side effect.
9010    ///
9011    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
9012    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
9013    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
9014    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
9015    /// still contains must be masked per query. memra's window convention
9016    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
9017    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
9018    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
9019    ///
9020    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
9021    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
9022    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
9023    ///
9024    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
9025    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
9026    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
9027    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
9028    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
9029    /// hidden rows, and the generated text — a function of the chunk size:
9030    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
9031    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
9032    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
9033    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
9034    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
9035    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
9036    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
9037    ///   one-token change in a documented machine-config knob changed the answer.
9038    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
9039    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
9040    /// the same rows moves the logits by ~1.8.
9041    ///
9042    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
9043    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
9044    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
9045    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
9046    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
9047    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
9048    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
9049    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
9050    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
9051    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
9052    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
9053    /// those with t_kv <= win = 512.
9054    #[allow(clippy::too_many_arguments)]
9055    fn step35_attn_pre_wo(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
9056                          hg: Option<&CudaSlice<f32>>, gt_pre: Option<&CudaSlice<f32>>,
9057                          pos_d: &CudaSlice<i32>, t: usize,
9058                          cache: Option<&mut Cache>, il: usize, seq_end: usize)
9059                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9060        let geometry = self.step35_geom(il);
9061        let hd = geometry.head_dim_k as usize;
9062        let nkv = geometry.n_head_kv as usize;
9063        let nh = geometry.n_head as usize;
9064        let rbase = geometry.rope_base;
9065        let scale = geometry.attention_scale();
9066        let swa = geometry.window.is_some();
9067        let eps = self.cfg.rms_eps;
9068        let win = geometry.window.unwrap_or(0) as usize;
9069        let n_rot = geometry.n_rot as usize;
9070
9071        let v = g3.pop().unwrap();
9072        let k0 = g3.pop().unwrap();
9073        let q0 = g3.pop().unwrap();
9074
9075        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
9076        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
9077        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
9078        let mut q = e.uninit(t * nh * hd)?;
9079        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
9080        let mut k = e.uninit(t * nkv * hd)?;
9081        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
9082        let ff = if geometry.rope_factors {
9083            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
9084        } else {
9085            None
9086        };
9087        #[cfg(debug_assertions)]
9088        if let Some(ff) = ff {
9089            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9090                                                       "step35_attn_pre_wo.rope_freqs");
9091        }
9092        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
9093
9094        let mut attn = e.uninit(t * nh * hd)?;
9095        match cache {
9096            Some(cache) => {
9097                let base_len = cache.kv[il].as_ref().unwrap().len;
9098                // Read per layer call, never in a measured default.
9099                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
9100                let legacy_calllocal =
9101                    std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
9102                let off = if swa {
9103                    let raw = base_len.saturating_sub(win - 1);
9104                    if legacy_tkv || legacy_calllocal { raw } else { raw & !31usize }
9105                } else {
9106                    0
9107                };
9108                {
9109                    let kvl = cache.kv[il].as_mut().unwrap();
9110                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
9111                    let write_row = e.prepare_kv_append(kvl, off, t)?;
9112                    e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, write_row, t,
9113                                               kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
9114                                               kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
9115                    kvl.len += t;
9116                    let new_len = kvl.len as i32;
9117                    e.set_i32_one(&mut kvl.len_d, new_len)?;
9118                }
9119                let kvl = cache.kv[il].as_ref().unwrap();
9120                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
9121                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
9122                // unaligned view offset here. Both halves are load-bearing for the canaries:
9123                // on the FA default the predicate arms agree bitwise wherever they can differ
9124                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
9125                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
9126                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
9127                // on the current FA path: its tile grid starts at the chunk/call boundary.
9128                // SWA: trim the view to the oldest key any query in this chunk can reach —
9129                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
9130                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
9131                // kernel's online-softmax recurrence groups keys into BK tiles relative to
9132                // the VIEW START — so an unaligned off regroups the same absolute keys into
9133                // different tiles at different chunk sizes = different (m,l) rounding =
9134                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
9135                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
9136                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
9137                // size; the <=31 extra leading keys are older than EVERY query's window
9138                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
9139                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
9140                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
9141                // the floor arm's bits do not move either (gated: G2f, battery 2).
9142                let t_kv = base_len + t - off;
9143                let physical = kvl.physical_rows(off, off + t_kv)?;
9144                let k_view = e.view_u8_range(&kvl.k, physical.start * kvl.k_tok_bytes,
9145                                             physical.end * kvl.k_tok_bytes);
9146                let v_view = e.view_u8_range(&kvl.v, physical.start * kvl.v_tok_bytes,
9147                                             physical.end * kvl.v_tok_bytes);
9148                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
9149                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
9150                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
9151                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
9152                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
9153                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
9154                // construction, so the invariance assertion MUST break under it (the seam whose
9155                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
9156                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
9157                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
9158                // cached (probes flip it in-process). Never on in a measured default run.
9159                let swa_naive = if legacy_tkv { t_kv > win } else { seq_end > win };
9160                if swa && swa_naive {
9161                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
9162                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
9163                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
9164                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
9165                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
9166                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
9167                    // identically to the unwindowed one modulo the mask, which is the point.
9168                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
9169                    // selected on `seq_end` like every arm here, so the class is uniform for
9170                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
9171                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
9172                    // the f32 floor (the previous numeric config, kept as the A/B seam).
9173                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
9174                        e.sdpa_naive_w_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh,
9175                                                      nkv, t, t_kv, scale, true, win,
9176                                                      kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9177                    } else {
9178                        e.fa_prefill_view_ws_w_hd128(&q, &k_view, &v_view, &mut attn, hd, nh,
9179                                                     nkv, t, t_kv, scale, true, win,
9180                                                     kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9181                    }
9182                } else if std::env::var("MEMRA_NOFA").is_ok() {
9183                    e.sdpa_naive_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9184                                                t, t_kv, scale, true,
9185                                                kvl.k_tok_bytes, kvl.v_tok_bytes)?;
9186                } else {
9187                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
9188                    // reach past the window, so the window mask is a no-op under causal and every
9189                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
9190                    // request either way, which is what makes the chunk size arithmetic-free.
9191                    e.fa_prefill_view_ws(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9192                                         t, t_kv, scale, true,
9193                                         kvl.k_tok_bytes, kvl.v_tok_bytes,
9194                                         crate::Engine::kv_fp8_on())?;
9195                }
9196            }
9197            None => {
9198                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
9199                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
9200                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
9201                // seq_end here too or it re-opens the same door.
9202                debug_assert_eq!(seq_end, t, "step35 cacheless prefill is monolithic (seq_end == t)");
9203                if swa && seq_end > win {
9204                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
9205                } else if std::env::var("MEMRA_NOFA").is_ok() {
9206                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9207                } else {
9208                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9209                }
9210            }
9211        }
9212
9213        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
9214        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
9215        let gw = fa.attn_gate.as_ref()
9216            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
9217        let gt_owned = if gt_pre.is_none() {
9218            Some(e.matmul(
9219                gw,
9220                hg.ok_or("step35 attention needs hg when gt_pre is absent")?,
9221                t,
9222            )?)
9223        } else {
9224            None
9225        };
9226        let gt = gt_pre.or(gt_owned.as_ref()).unwrap();
9227        let mut ag = e.uninit(t * nh * hd)?;
9228        e.attn_head_gate(&attn, gt, &mut ag, None, hd, nh, t)?;
9229        Ok(ag)
9230    }
9231
9232    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
9233    /// `forward_last`, t2probe). Post-`wo`.
9234    pub(crate) fn step35_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
9235                              pos_d: &CudaSlice<i32>, t: usize, il: usize)
9236                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9237        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
9238        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
9239        let ag = self.step35_attn_pre_wo(e, fa, g3, Some(h), None, pos_d, t, None, il, t)?;
9240        Ok(e.matmul(&fa.wo, &ag, t)?)
9241    }
9242
9243    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
9244    /// resident quantized cache, attend through the cache view). Post-`wo`.
9245    ///
9246    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
9247    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
9248    /// own extent.
9249    #[allow(clippy::too_many_arguments)]
9250    pub(crate) fn step35_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
9251                                    hx: Option<&CudaSlice<u8>>, pos_d: &CudaSlice<i32>, t: usize,
9252                                    cache: &mut Cache, il: usize, seq_end: usize)
9253                                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9254        let g3 = match hx {
9255            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
9256            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
9257        };
9258        let ag = self.step35_attn_pre_wo(
9259            e,
9260            fa,
9261            g3,
9262            Some(h),
9263            None,
9264            pos_d,
9265            t,
9266            Some(cache),
9267            il,
9268            seq_end,
9269        )?;
9270        Ok(e.matmul(&fa.wo, &ag, t)?)
9271    }
9272
9273    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
9274    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
9275    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
9276    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
9277    /// requiring `attn_gate`).
9278    ///
9279    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
9280    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
9281    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
9282    #[allow(clippy::too_many_arguments)]
9283    pub(crate) fn step35_decode_attn(&self, e: &Engine, fa: &FullAttnLayer, il: usize,
9284                          h: &CudaSlice<f32>,
9285                          pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9286                          pos_d: &CudaSlice<i32>, cache: &mut Cache)
9287                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9288        let geometry = self.step35_geom(il);
9289        let hd = geometry.head_dim_k as usize;
9290        let nkv = geometry.n_head_kv as usize;
9291        let nh = geometry.n_head as usize;
9292        let rbase = geometry.rope_base;
9293        let scale = geometry.attention_scale();
9294        let swa = geometry.window.is_some();
9295        let eps = self.cfg.rms_eps;
9296        let win = geometry.window.unwrap_or(0) as usize;
9297        let n_rot = geometry.n_rot as usize;
9298        let n_embd = self.cfg.n_embd as usize;
9299        let gw = fa.attn_gate.as_ref()
9300            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
9301
9302        let (q0, k0, v0, gt) = match pre_q {
9303            Some((hq, hdq)) => {
9304                debug_assert!(e.uses_q8_1_fast(gw),
9305                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
9306                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast");
9307                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
9308                    Some(t3) => t3,
9309                    None => (e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
9310                             e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
9311                             e.matmul_pre(&fa.wv, hq, hdq, h, 1)?),
9312                };
9313                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
9314                (a, b, c, gt)
9315            }
9316            None => {
9317                if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
9318                    && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw) {
9319                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
9320                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
9321                        Some(t3) => t3,
9322                        None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
9323                                 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
9324                                 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
9325                    };
9326                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
9327                    (a, b, c, gt)
9328                } else {
9329                    (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
9330                     e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
9331                }
9332            }
9333        };
9334
9335        let mut q = e.uninit(nh * hd)?;
9336        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
9337        let mut k = e.uninit(nkv * hd)?;
9338        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
9339        let ff = if swa { None } else {
9340            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
9341        };
9342        #[cfg(debug_assertions)]
9343        if let Some(ff) = ff {
9344            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9345                                                       "step35_decode_attn.rope_freqs");
9346        }
9347        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
9348
9349        if std::env::var("MEMRA_NOFA").is_ok() {
9350            return Err("MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
9351                        cache; unset MEMRA_NOFA to use fa_decode".into());
9352        }
9353        let kvl = cache.kv[il].as_mut().unwrap();
9354        let next_len = kvl.len + 1;
9355        let (off, t_kv) = if swa && next_len > win { (next_len - win, win) } else { (0, next_len) };
9356        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
9357        e.append_kv_quantized(&k, &v0, &mut kvl.k, &mut kvl.v, write_row,
9358                              kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
9359                              crate::Engine::kv_fp8_on())?;
9360        kvl.len = next_len;
9361        let physical = kvl.physical_rows(off, off + t_kv)?;
9362        let k_view = e.view_u8_range(&kvl.k, physical.start * kvl.k_tok_bytes,
9363                                     physical.end * kvl.k_tok_bytes);
9364        let v_view = e.view_u8_range(&kvl.v, physical.start * kvl.v_tok_bytes,
9365                                     physical.end * kvl.v_tok_bytes);
9366        let mut attn = e.uninit(nh * hd)?;
9367        e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
9368                          kvl.k_tok_bytes, kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
9369
9370        let mut ag = e.uninit(nh * hd)?;
9371        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
9372        Ok(e.matmul(&fa.wo, &ag, 1)?)
9373    }
9374}
9375
9376// ===================================================================================== //
9377//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
9378//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
9379//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
9380//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
9381//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
9382//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
9383// ===================================================================================== //
9384impl HybridModel {
9385    pub fn is_gemma4_e4b(&self) -> bool {
9386        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
9387    }
9388
9389    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
9390    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
9391    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
9392    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
9393        let g = self.cfg.gemma4.as_ref().unwrap();
9394        let swa = g.swa_pattern[il];
9395        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
9396        let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
9397        let nh = fa.wq.out_features() / hd;
9398        let nkv = fa.wk.out_features() / hd;
9399        (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
9400    }
9401
9402    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
9403    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
9404        self.layers[il].gemma4.as_ref()
9405            .and_then(|b| b.e4b.as_ref())
9406            .and_then(|e4| e4.kv_share.map(|t| t as usize))
9407    }
9408
9409    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
9410    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
9411    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
9412    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
9413    fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
9414                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9415        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
9416        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
9417    }
9418
9419    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
9420    fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
9421                             x_scaled: &CudaSlice<f32>, t: usize)
9422                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9423        let aux = self.gemma4_aux.as_ref().unwrap();
9424        let m = aux.e4b.as_ref().unwrap();
9425        let n_embd = self.cfg.n_embd as usize;
9426        let n_layer = self.layers.len();
9427        let width = m.n_epl * n_layer;
9428        let tbl = m.tok_tbl_gpu.get_or_init(|| {
9429            e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
9430        });
9431        let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
9432                                             m.tok_embd_row_bytes)?;
9433        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
9434        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
9435        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
9436        let mut pn = e.uninit(t * width)?;
9437        e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
9438                   self.cfg.rms_eps)?;
9439        let mut out = e.uninit(t * width)?;
9440        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
9441        Ok(out)
9442    }
9443
9444    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
9445    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
9446    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
9447    /// already holds this forward's rows — the target runs earlier in the stack).
9448    #[allow(clippy::too_many_arguments)]
9449    fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
9450                       hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
9451                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
9452                       dc_bucket: Option<usize>)
9453                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9454        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
9455        let eps = self.cfg.rms_eps;
9456        let aux = self.gemma4_aux.as_ref().unwrap();
9457        let ones = aux.ones(e);
9458        #[cfg(debug_assertions)]
9459        crate::debug_assert_tensor_stream_device(ones, &e.stream(),
9460                                                   "gemma4_e4b_attn.ones");
9461        let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
9462        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
9463        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
9464        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
9465        let h0 = e.zeros(0)?;
9466        let h = &h0;
9467
9468        let ff = if swa { None } else {
9469            Some(aux.rope_freqs(e).expect("e4b global rope needs rope_freqs.weight"))
9470        };
9471        #[cfg(debug_assertions)]
9472        if let Some(ff) = ff {
9473            crate::debug_assert_tensor_stream_device(ff, &e.stream(),
9474                                                       "gemma4_e4b_attn.rope_freqs");
9475        }
9476        let share = self.gemma4_e4b_kv_target(il);
9477        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
9478        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
9479        let mut q;
9480        if let Some(_tgt) = share {
9481            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
9482            q = e.uninit(t * nh * hd)?;
9483            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
9484            // empty; q0 stands in for the unused k/v pointers).
9485            let mut kdummy = e.uninit(1)?;
9486            let mut vdummy = e.uninit(1)?;
9487            e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
9488                                fa.q_norm.float_data(), ones,
9489                                &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
9490                                pos_d, nh, 1, base, 1.0, ff, eps)?;
9491        } else {
9492            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
9493            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
9494            // q|k|v rows — the cat norm+rope twin consumes it directly.
9495            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
9496            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
9497            q = e.uninit(t * nh * hd)?;
9498            let mut k = e.uninit(t * nkv * hd)?;
9499            let mut v = e.uninit(t * nkv * hd)?;
9500            if t == 1 && cat.is_some() {
9501                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
9502                e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
9503                                        ones, &mut q, &mut k, &mut v, hd, nh, nkv,
9504                                        pos_d, nh, nkv, base, 1.0, ff, eps)?;
9505            } else {
9506                let (q0, k0, v0) = match if t == 1 {
9507                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
9508                } else {
9509                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
9510                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
9511                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9512                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
9513                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
9514                    } else { None }
9515                } {
9516                    Some(triple) => triple,
9517                    None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
9518                             e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
9519                             e.matmul_pre(&fa.wv, hq, hdq, h, t)?),   // E4B: real v (K != V)
9520                };
9521                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
9522                // the normed rows; V ones-rms, never roped).
9523                e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
9524                                    fa.k_norm.float_data(), ones, &mut q, &mut k, &mut v,
9525                                    hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
9526            }
9527            let kvl = cache.kv[il].as_mut().unwrap();
9528            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
9529            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
9530            // degenerate tok-0 stream, 2026-07-12).
9531            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9532            if dc_bucket.is_some() {
9533                // DC arm (graph serving): append at the len_d slot, advance the counter
9534                // in-stream — replay-correct, no host len in the launch args. Host mirrors
9535                // are NOT touched here (the replay loop owns them; a bump at capture-record
9536                // time would double-count the capture iteration).
9537                debug_assert!(t == 1);
9538                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
9539                e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
9540                                                 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
9541                                                 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
9542            } else {
9543                e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
9544                                           kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
9545                                           kvl.v_tok_bytes, cls)?;
9546                kvl.len += t;
9547            }
9548            kv_f32 = Some((k, v));
9549        }
9550        // attention: per-row causal fa over the (own or target) quantized cache. The cache
9551        // already contains this forward's rows in both arms; row i attends [.., base+i].
9552        let kvl_idx = share.unwrap_or(il);
9553        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
9554        let base_len = kvl.len - t;   // pre-append length (target appended this forward too)
9555        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
9556        let mut attn = e.uninit(t * nh * hd)?;
9557        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
9558        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
9559        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
9560        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
9561        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
9562        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
9563        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
9564        //     rows (the T=K verify kernel; the target appended this forward's rows already).
9565        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
9566        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
9567        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
9568        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
9569            if let Some((kf, vf)) = &kv_f32 {
9570                if hd == 256 && t <= win {
9571                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
9572                    return Ok(e.matmul(&fa.wo, &attn, t)?);
9573                }
9574                if hd == 256 && swa && t > win {
9575                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9576                                   win)?;
9577                    return Ok(e.matmul(&fa.wo, &attn, t)?);
9578                }
9579                if hd == 512 && !swa {
9580                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
9581                                       true)?;
9582                    return Ok(e.matmul(&fa.wo, &attn, t)?);
9583                }
9584            } else if share.is_some() {
9585                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9586                let k_view = e.view_u8(&kvl.k, kvl.k.len());
9587                let v_view = e.view_u8(&kvl.v, kvl.v.len());
9588                if hd == 256 && (!swa || t <= win) {
9589                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
9590                    e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
9591                                      scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9592                    return Ok(e.matmul(&fa.wo, &attn, t)?);
9593                }
9594                // remaining shared classes (swa above the window; hd512 globals): dequant
9595                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
9596                let kv_dim = nkv * hd;
9597                let mut kf = e.uninit(t * kv_dim)?;
9598                let mut vf = e.uninit(t * kv_dim)?;
9599                e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
9600                                         t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9601                if hd == 512 {
9602                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
9603                                       true)?;
9604                } else {
9605                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
9606                                   win)?;
9607                }
9608                return Ok(e.matmul(&fa.wo, &attn, t)?);
9609            }
9610        }
9611        if let Some(bucket) = dc_bucket {
9612            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
9613            // fa_decode_dc over the live counter. len_d already advanced past this token
9614            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
9615            // counter (advanced when the target ran earlier in the stack).
9616            assert!(t == 1);
9617            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
9618            // and under the window every live t_kv sits below it — cap the capture bucket
9619            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
9620            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
9621            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
9622            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
9623                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
9624            } else { bucket };
9625            let k_view = e.view_u8(&kvl.k, kvl.k.len());
9626            let v_view = e.view_u8(&kvl.v, kvl.v.len());
9627            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
9628            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
9629            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
9630            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
9631            // captured into the dc graph like any other launch. Extending the cascade to
9632            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
9633            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
9634            // MEMRA_WPF=0 rollback seam.
9635            if crate::Engine::wpf_level() >= 1 {
9636                e.prefetch_weight_l2(&fa.wo)?;
9637            }
9638            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
9639            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
9640            if e.uses_q8_1_fast(&fa.wo) {
9641                let mut oq = e.alloc_i8_uninit(nh * hd)?;
9642                let mut od = e.zeros(nh * hd / 32)?;
9643                e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9644                                  &kvl.len_d, bucket, scale,
9645                                  kvl.k_tok_bytes, kvl.v_tok_bytes, g,
9646                                  Some((&mut oq, &mut od)))?;
9647                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
9648            }
9649            e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
9650                           &kvl.len_d, bucket, scale,
9651                           kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
9652            return Ok(e.matmul(&fa.wo, &attn, t)?);
9653        }
9654        for i in 0..t {
9655            let avail = base_len + i + 1;
9656            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
9657            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
9658                                         (off_tok + t_kv) * kvl.k_tok_bytes);
9659            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
9660                                         (off_tok + t_kv) * kvl.v_tok_bytes);
9661            let qv = e.view(&q, t * nh * hd);
9662            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
9663            let mut q_one = e.uninit(nh * hd)?;
9664            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
9665            let mut a_one = e.uninit(nh * hd)?;
9666            // read class MUST match the append class (globals are e4m3 under gkv): the
9667            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
9668            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
9669            e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
9670                        kvl.k_tok_bytes, kvl.v_tok_bytes,
9671                        (!swa && crate::Engine::gkv_on())
9672                            || (swa && crate::Engine::wkv_on()))?;
9673            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
9674        }
9675        Ok(e.matmul(&fa.wo, &attn, t)?)
9676    }
9677
9678    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
9679    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
9680    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
9681    /// layer; does NOT advance cache.pos (caller owns pos).
9682    fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
9683                        head_last: bool)
9684                        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9685        let n_embd = self.cfg.n_embd as usize;
9686        let t = tokens.len();
9687        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9688        let pos_d = e.htod_i32(&pos)?;
9689        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
9690        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9691        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
9692        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
9693    }
9694
9695    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
9696    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
9697    /// eager chain by construction: SAME functions, not twins).
9698    fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
9699                             pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
9700                             dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
9701                             -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9702        let n_embd = self.cfg.n_embd as usize;
9703        let eps = self.cfg.rms_eps;
9704        let n_layer = self.layers.len();
9705        let mut x = x_in;
9706        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
9707        let n_epl = aux_e4b.n_epl;
9708
9709        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
9710        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
9711        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
9712        // head rides matmul_pre too. First layer's pair comes from a standalone fused
9713        // norm+quant.
9714        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
9715        for il in 0..n_layer {
9716            let layer = &self.layers[il];
9717            let (hq, hdq) = match h_carry.take() {
9718                Some(p) => p,
9719                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
9720            };
9721            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
9722            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
9723            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
9724            let bits = layer.gemma4.as_ref().unwrap();
9725            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
9726            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
9727            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
9728            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
9729            // the fused single-phase reduction is NOT FP-order-identical to the unfused
9730            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
9731            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
9732            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
9733            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
9734            // gate dropped, decode AND verify ride the same fused chain — parity by
9735            // construction, VERIFY-GATE 0.000e0.
9736            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
9737            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
9738                e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
9739            let mut resid = e.uninit(t * n_embd)?;
9740            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
9741            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
9742            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
9743            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
9744            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
9745            let g = if fuse_exit {
9746                // sn here = RAW f0 (post_ffw deferred).
9747                let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
9748                                                  &attn_out, &mut resid, n_embd, t,
9749                                                  self.cfg.rms_eps)?;
9750                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
9751            } else {
9752                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
9753                e.matmul(&e4b.inp_gate, &resid, t)?
9754            };
9755            let mut act = e.uninit(t * n_epl)?;
9756            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
9757                let ipv = e.view(&inp_pl, n_epl * n_layer);
9758                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
9759                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
9760                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
9761            } else {
9762                let mut inp_this = e.uninit(t * n_epl)?;
9763                e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
9764                                    il * n_epl)?;
9765                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
9766                e.matmul(&e4b.proj, &act, t)?
9767            };
9768            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
9769            // ONE launch (glue-fusion lane; last layer emits through output_norm).
9770            let next_norm = if il + 1 < n_layer {
9771                self.layers[il + 1].attn_norm.float_data()
9772            } else {
9773                self.output_norm.float_data()
9774            };
9775            let mut xn = e.uninit(t * n_embd)?;
9776            let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
9777                                                         &resid, bits.layer_scale, next_norm,
9778                                                         &mut xn, n_embd, t, eps)?;
9779            h_carry = Some(pair);
9780            x = xn;
9781        }
9782        // the head consumes the last layer's fused (output_norm) emit. head_last callers
9783        // (prime, last_only forward) need only the final row's logits — the all-T head is
9784        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
9785        let (oq, odq) = h_carry.take().unwrap();
9786        let h0 = e.zeros(0)?;
9787        let hm = if head_last { 1 } else { t };
9788        let (hq, hd) = if head_last && t > 1 {
9789            let mut q1 = e.uninit_i8(n_embd)?;
9790            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
9791            let nb = n_embd / 32;
9792            let mut d1 = e.uninit(nb)?;
9793            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
9794            (q1, d1)
9795        } else {
9796            (oq, odq)
9797        };
9798        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
9799        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
9800        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
9801        // Logit-returning callers (host logits / spec prime) keep the capped emit.
9802        if cap_logits {
9803            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
9804            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
9805        }
9806        self.gemma4_suppress(e, &mut ld, hm)?;  // mask both capped and argmax-only consumers
9807        Ok((ld, x))
9808    }
9809
9810    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
9811    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
9812    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
9813    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
9814    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
9815    /// covers exactly the layers that appended).
9816    pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
9817                                                  t: usize, pos0: usize, cache: &mut Cache)
9818                                                  -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9819        let n_embd = self.cfg.n_embd as usize;
9820        let eps = self.cfg.rms_eps;
9821        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
9822        let pos_d = e.htod_i32(&pos)?;
9823        let embd_gpu = self.embd_gpu.get_or_init(|| {
9824            e.upload_u8(&self.embd.raw).expect("embed table upload")
9825        });
9826        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
9827        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
9828        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
9829        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
9830        let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
9831                                                  false)?;
9832        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
9833        // emit is already capped, matching the eager chain bit-for-bit).
9834        let n_vocab = self.output.out_features();
9835        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
9836        for i in 0..t {
9837            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
9838        }
9839        let mut hn = e.uninit(t * n_embd)?;
9840        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9841        cache.pos += t;
9842        Ok((vam, hn))
9843    }
9844
9845    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
9846    /// prime path — mirror of `gemma4_decode_step_t_h`).
9847    pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
9848                                             cache: &mut Cache)
9849                                             -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9850        let n_embd = self.cfg.n_embd as usize;
9851        let eps = self.cfg.rms_eps;
9852        let t = tokens.len();
9853        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
9854        let mut hn = e.uninit(t * n_embd)?;
9855        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9856        cache.pos += t;
9857        Ok((e.dtoh(&ld)?, hn))
9858    }
9859
9860    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
9861    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
9862    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
9863    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
9864    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
9865    pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
9866                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9867                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9868                                      n_vocab: usize, bucket: usize)
9869                                      -> Result<(), Box<dyn std::error::Error>> {
9870        let n_embd = self.cfg.n_embd as usize;
9871        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9872        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9873        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9874        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
9875                                                  false, false)?;
9876        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
9877        e.inc_seqlen(pos_d)?;
9878        Ok(())
9879    }
9880
9881    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
9882    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
9883    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
9884    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
9885    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
9886    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
9887    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
9888    #[allow(clippy::too_many_arguments)]
9889    pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
9890                                     pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
9891                                     embd_qt: i32, embd_rb: usize, cache: &mut Cache,
9892                                     n_vocab: usize)
9893                                     -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9894        let n_embd = self.cfg.n_embd as usize;
9895        let eps = self.cfg.rms_eps;
9896        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
9897        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
9898        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
9899        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
9900                                                  false)?;
9901        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
9902        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
9903        e.inc_seqlen(pos_d)?;
9904        cache.pos += 1;
9905        let _ = eps;
9906        Ok(tok_out)
9907    }
9908
9909    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
9910    /// pre-output_norm hidden). Advances cache.pos.
9911    pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
9912                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9913        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
9914        let logits = e.dtoh(&ld)?;
9915        cache.pos += 1;
9916        Ok((logits, x))
9917    }
9918
9919    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
9920    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
9921    /// fast; the prefill fa arms come later.
9922    pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
9923                                   -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9924        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
9925        // process-kill as gemma4_prime — refuse per-request.
9926        if cache.pos != 0 {
9927            return Err("e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
9928                        call or decode tokenwise".into());
9929        }
9930        let n_embd = self.cfg.n_embd as usize;
9931        let t = tokens.len();
9932        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
9933        cache.pos += t;
9934        let last = e.dtoh(&ld)?;   // head_last: ld is already the final row only
9935        let xv = e.view(&x, t * n_embd);
9936        let row = xv.slice((t - 1) * n_embd..t * n_embd);
9937        let mut h_seed = e.uninit(n_embd)?;
9938        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
9939        Ok((last, h_seed, x))
9940    }
9941
9942    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
9943    pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
9944                                     -> Result<Vec<f32>, Box<dyn std::error::Error>> {
9945        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
9946        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
9947        Ok(e.dtoh(&ld)?)   // head_last already reduced to the final row when last_only
9948    }
9949}
9950
9951#[cfg(test)]
9952mod prime_chunk_schedule_tests {
9953    use super::{
9954        dynamic_prime_chunk_ranges, fixed_prime_chunk_ranges, fixed_prime_chunk_ranges_for_ring,
9955        PRIME_MIN_T,
9956        PRIME_PIPE_MIN_CHUNK,
9957    };
9958
9959    fn sizes(ranges: &[(usize, usize)]) -> Vec<usize> {
9960        ranges.iter().map(|(start, end)| end - start).collect()
9961    }
9962
9963    fn auto_chunk(t: usize) -> usize {
9964        t.div_ceil(8).max(PRIME_PIPE_MIN_CHUNK).min(4096)
9965    }
9966
9967    #[test]
9968    fn fixed_schedule_retains_measured_geometry() {
9969        assert_eq!(
9970            sizes(&fixed_prime_chunk_ranges(461, 128)),
9971            vec![128, 128, 128, 77]
9972        );
9973        assert_eq!(
9974            sizes(&fixed_prime_chunk_ranges(1833, 230)),
9975            vec![230, 230, 230, 230, 230, 230, 230, 223]
9976        );
9977        assert_eq!(
9978            sizes(&fixed_prime_chunk_ranges(4096, 512)),
9979            vec![512; 8]
9980        );
9981        let capped = sizes(&fixed_prime_chunk_ranges_for_ring(8200, 4096, true));
9982        assert_eq!(capped, vec![4096, 4088, 16]);
9983        assert!(capped.iter().all(|&rows| rows <= 4096));
9984        assert_eq!(
9985            sizes(&fixed_prime_chunk_ranges_for_ring(4100, 4096, false)),
9986            vec![4100],
9987            "flag-off schedule remains byte-for-byte the legacy monolithic tail",
9988        );
9989    }
9990
9991    #[test]
9992    fn dynamic_schedule_matches_registered_shapes() {
9993        let cases = [
9994            (461, vec![64, 141, 132, 124]),
9995            (1833, vec![115, 269, 260, 252, 244, 237, 231, 225]),
9996            (4096, vec![256, 602, 580, 563, 545, 531, 516, 503]),
9997        ];
9998        for (t, expected) in cases {
9999            let chunk = auto_chunk(t);
10000            let fixed = fixed_prime_chunk_ranges(t, chunk);
10001            assert_eq!(
10002                sizes(&dynamic_prime_chunk_ranges(t, chunk, &fixed)),
10003                expected
10004            );
10005        }
10006    }
10007
10008    #[test]
10009    fn dynamic_schedule_covers_exactly_and_shrinks_after_fill() {
10010        for t in 256..=8192 {
10011            let chunk = auto_chunk(t);
10012            let fixed = fixed_prime_chunk_ranges(t, chunk);
10013            let dynamic = dynamic_prime_chunk_ranges(t, chunk, &fixed);
10014            assert_eq!(dynamic.len(), fixed.len(), "T={t}");
10015            assert_eq!(dynamic.first().unwrap().0, 0, "T={t}");
10016            assert_eq!(dynamic.last().unwrap().1, t, "T={t}");
10017            for pair in dynamic.windows(2) {
10018                assert_eq!(pair[0].1, pair[1].0, "T={t}");
10019            }
10020            assert!(
10021                dynamic
10022                    .iter()
10023                    .all(|(start, end)| end - start >= PRIME_MIN_T),
10024                "T={t} sizes={:?}",
10025                sizes(&dynamic)
10026            );
10027            if dynamic.len() >= 3 {
10028                let chunk_sizes = sizes(&dynamic);
10029                assert!(
10030                    chunk_sizes[0] < chunk_sizes[1],
10031                    "T={t} sizes={chunk_sizes:?}"
10032                );
10033                assert!(
10034                    chunk_sizes[1..].windows(2).all(|pair| pair[0] >= pair[1]),
10035                    "T={t} sizes={chunk_sizes:?}"
10036                );
10037            }
10038        }
10039    }
10040}
10041
10042#[cfg(test)]
10043mod page_prefetch_tests {
10044    use super::{
10045        grouped_worker_prefetch_position, page_prefetch_positions,
10046        page_prefetch_window_from_values, worker_prefetch_positions,
10047    };
10048
10049    #[test]
10050    fn page_prefetch_window_keeps_existing_opt_in_default() {
10051        assert_eq!(page_prefetch_window_from_values(false, None), 0);
10052        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
10053        assert_eq!(page_prefetch_window_from_values(true, None), 1);
10054        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
10055        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
10056        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
10057    }
10058
10059    #[test]
10060    fn rolling_page_prefetch_advises_each_future_expert_once() {
10061        let advised: Vec<_> = (0..7)
10062            .flat_map(|position| page_prefetch_positions(position, 7, 3))
10063            .collect();
10064        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
10065
10066        let one_ahead: Vec<_> = (0..4)
10067            .flat_map(|position| page_prefetch_positions(position, 4, 1))
10068            .collect();
10069        assert_eq!(one_ahead, vec![1, 2, 3]);
10070        assert!(page_prefetch_positions(0, 4, 0).is_empty());
10071    }
10072
10073    #[test]
10074    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
10075        assert_eq!(grouped_worker_prefetch_position(0, None), None);
10076        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
10077            .chain((0..4).filter_map(|position| {
10078                grouped_worker_prefetch_position(4, Some(position))
10079            }))
10080            .collect();
10081        assert_eq!(positions, vec![0, 1, 2, 3]);
10082        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
10083    }
10084
10085    #[test]
10086    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
10087        let queued: Vec<_> = (0..8)
10088            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
10089            .collect();
10090        assert_eq!(queued, (0..8).collect::<Vec<_>>());
10091
10092        let one_at_a_time: Vec<_> = (0..4)
10093            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
10094            .collect();
10095        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
10096        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
10097    }
10098}
10099
10100pub struct G4DcSlots {
10101    x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
10102    hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
10103    q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
10104    q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
10105    attn: CudaSlice<f32>, o: CudaSlice<f32>,
10106    attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
10107    zq: CudaSlice<i8>, zd: CudaSlice<f32>,
10108    gate: CudaSlice<f32>, up: CudaSlice<f32>,
10109    act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
10110    f0: CudaSlice<f32>, sn: CudaSlice<f32>,
10111    hn: CudaSlice<f32>, logits: CudaSlice<f32>,
10112}