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
41/// task #18 (attn side): one sequence's pre-attention outputs (post-rope q/k, v, out-gate).
42pub(crate) struct AttnPre {
43    pub q: cudarc::driver::CudaSlice<f32>,
44    pub k: cudarc::driver::CudaSlice<f32>,
45    pub v: cudarc::driver::CudaSlice<f32>,
46    pub gate: Option<cudarc::driver::CudaSlice<f32>>,
47}
48
49/// task #18: one sequence's GDN prep outputs (the scan inputs).
50pub(crate) struct GdnPrep {
51    pub hk: usize,
52    pub q_l2: cudarc::driver::CudaSlice<f32>,
53    pub k_l2: cudarc::driver::CudaSlice<f32>,
54    pub v_g: cudarc::driver::CudaSlice<f32>,
55    pub beta: cudarc::driver::CudaSlice<f32>,
56    pub g_log: cudarc::driver::CudaSlice<f32>,
57    pub kb16: Option<cudarc::driver::CudaSlice<u8>>,
58    pub qb16: Option<cudarc::driver::CudaSlice<u8>>,
59}
60
61/// Device scratch for the burst verify stream (see `verify_stream_scratch`).
62pub(crate) struct VerifyStreamScratch {
63    pub pos_d: CudaSlice<i32>,
64    pub row_ctrs: Vec<CudaSlice<i32>>,
65}
66use crate::hybrid::{HybridModel, Mixer, FullAttnLayer, LinearAttnLayer, MoeWeights};
67
68struct MoeInputTraceWriter {
69    dir: std::path::PathBuf,
70    index: std::fs::File,
71    payloads: std::collections::HashMap<u16, (std::fs::File, u64)>,
72}
73
74static MOE_INPUT_TRACE_WRITER: std::sync::OnceLock<
75    std::sync::Mutex<Option<MoeInputTraceWriter>>,
76> = std::sync::OnceLock::new();
77
78/// STAGE-2 GROUPED DECODE gate (MEMRA_MOE_GDEC, default ON; `=0` restores the sequential
79/// per-expert launch chain). See `moe_gdec_token`.
80fn gdec_enabled() -> bool {
81    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
82    *E.get_or_init(|| std::env::var("MEMRA_MOE_GDEC").map(|v| v != "0").unwrap_or(true))
83}
84
85/// SLAB-LOCAL RESIDENT ARM gate (lane/pp-leverb 2026-08-08, MEMRA_MOE_SLAB, default ON;
86/// `=0` restores the SLRU dispatch even when resident slabs exist). Read PER CALL, never
87/// memoized — probes A/B the two provenances in one process (the MEMRA_PRIME_PP pattern).
88/// See `moe_ffn_sequential_zq8`'s slab_local arm: the sigmoid-router archs (step35/M3/Hy3)
89/// are denied every `dev_exps` consumer (pairs/dev route softmax), so before this arm the
90/// fits-VRAM resident slabs were UPLOADED for them but never READ — the SLRU kept staging
91/// the same bytes beside a dead copy (37 GB H2D per pp4096 prime on the Step SKU, anatomy
92/// receipt). The arm reads the SAME bytes through the SAME kernels; only the pointer
93/// PROVENANCE changes (slab base + ex*stride vs SLRU slot address) — the bit-identity class
94/// `moe_ffn_dev`'s resident arm already documents against its SLRU arm.
95fn moe_slab_enabled() -> bool {
96    std::env::var("MEMRA_MOE_SLAB").as_deref() != Ok("0")
97}
98
99/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
100/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
101fn moe_prefetch_enabled() -> bool {
102    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
103    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
104        || crate::spill_pread::worker_enabled())
105}
106
107/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
108/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
109/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
110/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
111fn moe_page_prefetch_window() -> usize {
112    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
113    *W.get_or_init(|| page_prefetch_window_from_values(
114        std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
115        std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
116    ))
117}
118
119fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
120    if !enabled {
121        return 0;
122    }
123    raw_window
124        .and_then(|value| value.parse().ok())
125        .unwrap_or(1)
126}
127
128/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
129/// full window; each later position adds one expert at the far edge. Thus widening the window does
130/// not repeatedly issue `MADV_WILLNEED` for the same range.
131fn page_prefetch_positions(
132    position: usize,
133    len: usize,
134    window: usize,
135) -> std::ops::Range<usize> {
136    if window == 0 || position >= len {
137        return len..len;
138    }
139    let (start, count) = if position == 0 {
140        (1, window)
141    } else {
142        (position.saturating_add(window), 1)
143    };
144    let start = start.min(len);
145    start..start.saturating_add(count).min(len)
146}
147
148/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
149/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
150fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
151    let position = current.map_or(0, |position| position.saturating_add(1));
152    (position < order_len).then_some(position)
153}
154
155/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
156/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
157/// window. Position zero primes the current expert too: its three independent reads can run in
158/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
159fn worker_prefetch_window() -> usize {
160    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
161    *WINDOW.get_or_init(|| {
162        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
163        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
164            .ok()
165            .and_then(|value| value.parse::<usize>().ok())
166            .unwrap_or(automatic.max(1))
167    })
168}
169
170/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
171/// this includes the current expert when the window is seeded so all three current projections
172/// enter the CPU pool together.
173fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
174    if window == 0 || position >= len {
175        return len..len;
176    }
177    let (start, count) = if position == 0 {
178        (0, window)
179    } else {
180        (position.saturating_add(window).saturating_sub(1), 1)
181    };
182    let start = start.min(len);
183    start..start.saturating_add(count).min(len)
184}
185
186/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
187/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
188/// expert weight pointers come from the per-layer device table. Requires the fused router (the
189/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
190fn moe_dev_enabled() -> bool {
191    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
192    *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
193        && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
194}
195
196/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
197/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
198/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
199/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
200fn moe_q8_enabled() -> bool {
201    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
202    *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
203}
204
205/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
206/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
207fn expert_dp4a_supported(qt: i32) -> bool {
208    qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
209        || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
210}
211
212fn q8_expert_supported(qt: i32) -> bool {
213    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
214    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
215    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
216    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
217    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
218    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
219    let kq = *KQ.get_or_init(|| {
220        std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
221    });
222    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
223    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
224    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
225    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
226    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
227    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
228    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
229    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
230        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
231}
232
233/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
234/// k-quant tensors must fall to the _em dot path instead.
235fn q8_expert_dec_supported(qt: i32) -> bool {
236    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
237}
238
239/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
240/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
241/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
242/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
243/// q35 layers, which is why that cell measured FLAT.
244fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
245    match qt {
246        crate::QT_Q4_0 => in_f % 32 == 0,
247        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
248        | crate::QT_Q6_K => in_f % 256 == 0,
249        _ => false,
250    }
251}
252
253/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
254/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
255fn moe_prewarm_enabled() -> bool {
256    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
257    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
258}
259
260/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
261/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
262/// can then vote for and exercise those experts on GPU before the cache is frozen.
263fn cpu_expert_profile_admit_enabled() -> bool {
264    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
265    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
266}
267
268/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
269/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
270/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
271pub const PRIME_MIN_T: usize = 16;
272
273impl HybridModel {
274    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
275    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
276    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
277    /// (it forces a dtoh + host hash per layer).
278    fn prime_trace_path() -> Option<&'static str> {
279        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
280        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
281            .as_deref()
282    }
283
284    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
285    pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
286        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
287        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
288        let cfg = &self.cfg;
289        let n_embd = cfg.n_embd as usize;
290        let t = tokens.len();
291        let eps = cfg.rms_eps;
292        let pos: Vec<i32> = (0..t as i32).collect();
293        let pos_d = e.htod_i32(&pos)?;
294
295        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
296
297        for (il, layer) in self.layers.iter().enumerate() {
298            // attn_norm
299            let mut h = e.uninit(t * n_embd)?;
300            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
301
302            let mixed = match &layer.mixer {
303                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
304                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
305                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
306            };
307
308            // residual 1
309            let mut x1 = e.uninit(t * n_embd)?;
310            e.add(&x, &mixed, &mut x1, t * n_embd)?;
311
312            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
313            let mut z = e.uninit(t * n_embd)?;
314            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
315            let ffn_out = match &layer.ffn {
316                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
317                    let n_ff = ffn_gate.out_features();
318                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
319                    let up = g2.pop().unwrap();
320                    let gate = g2.pop().unwrap();
321                    let mut act = e.uninit(t * n_ff)?;
322                    // A DENSE FFN reads the SHEXP clamp array: upstream's one `build_ffn` serves
323                    // both the dense MLP and the shared expert, and its limit is
324                    // `swiglu_clamp_shexp[il]` (llama-graph.cpp:1751). step35's leading dense
325                    // blocks 0-2 therefore key off clamp_shexp, not clamp_exp.
326                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
327                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
328                    e.matmul(ffn_down, &act, t)?
329                }
330                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
331            };
332            let mut x2 = e.uninit(t * n_embd)?;
333            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
334            x = x2;
335        }
336
337        let mut hn = e.uninit(t * n_embd)?;
338        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
339        let logits = e.matmul(&self.output, &hn, t)?;
340        Ok(e.dtoh(&logits)?)
341    }
342
343    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
344    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
345    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
346    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
347    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
348    pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
349        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
350        let cfg = &self.cfg;
351        let n_embd = cfg.n_embd as usize;
352        let t = tokens.len();
353        let eps = cfg.rms_eps;
354        let pos: Vec<i32> = (0..t as i32).collect();
355        let pos_d = e.htod_i32(&pos)?;
356
357        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
358        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
359        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
360        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
361        for (il, layer) in self.layers.iter().enumerate() {
362            let mut h = e.uninit(t * n_embd)?;
363            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
364            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
365            let mixed = match &layer.mixer {
366                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t, il)?,
367                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
368                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
369            };
370            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
371            let mut x1 = e.uninit(t * n_embd)?;
372            e.add(&x, &mixed, &mut x1, t * n_embd)?;
373            let mut z = e.uninit(t * n_embd)?;
374            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
375            let ffn_out = match &layer.ffn {
376                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
377                    let n_ff = ffn_gate.out_features();
378                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
379                    let up = g2.pop().unwrap();
380                    let gate = g2.pop().unwrap();
381                    let mut act = e.uninit(t * n_ff)?;
382                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
383                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
384                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
385                    e.matmul(ffn_down, &act, t)?
386                }
387                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
388            };
389            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
390            let mut x2 = e.uninit(t * n_embd)?;
391            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
392            x = x2;
393        }
394        // norm over all T, then slice the LAST row and run lm_head on that single row.
395        let mut hn = e.uninit(t * n_embd)?;
396        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
397        let last = e.view(&hn, t * n_embd);            // [T, n_embd]
398        let last_row = last.slice((t - 1) * n_embd..t * n_embd);  // [1, n_embd]
399        let mut hlast = e.uninit(n_embd)?;
400        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
401        let logits = e.matmul(&self.output, &hlast, 1)?;   // [1, n_vocab] — lm_head on ONE row
402        Ok(e.dtoh(&logits)?)
403    }
404
405    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
406    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
407    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
408    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
409    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
410    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
411    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
412    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
413    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
414    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
415    ///       argmax gate is the accuracy authority, exactly as for forward_last);
416    ///   (c) `cache.pos`/KV len/len_d advance by T.
417    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
418    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
419    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
420    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
421    ///
422    /// `queued_after` (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME
423    /// REQUEST that the caller will prime in LATER calls — 0 when this call is the whole
424    /// request (every single-shot caller). Serve splits a long prompt across SEVERAL
425    /// prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the
426    /// request's absolute end position `seq_end = cache.pos + t + queued_after` steers step35's
427    /// SWA prefill arm — computing it per CALL made the arm a function of the tick budget
428    /// (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes
429    /// default to 256 AND cap by live SLO headroom, so identical judge requests primed
430    /// differently under load — research/tick-seg-20260807, receipt in
431    /// research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally
432    /// lacked: it cannot know from `tokens` and `cache` alone whether more of the request is
433    /// coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW
434    /// request — its arithmetic is keyed to its own extent, so those callers pass 0; only a
435    /// caller that SPLITS one request across calls passes the remainder.
436    pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize)
437                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
438        let n_embd = self.cfg.n_embd as usize;
439        let t = tokens.len();
440        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
441        // session cache — every chunk (including the first) takes the continuation arm
442        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
443        assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
444        assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
445
446        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
447        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
448        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
449        // each chunk runs the full layer stack with transients sized to the chunk, appending its
450        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
451        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
452        // exactly the state carry it was built for). Full-attn chunks after the first attend to
453        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
454        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
455        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
456        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
457        if self.is_gemma4_e4b() {
458            return self.gemma4_e4b_prime(e, tokens, cache);
459        }
460        if self.cfg.gemma4.is_some() {
461            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
462            return self.gemma4_prime(e, tokens, cache);
463        }
464        let chunk: usize = std::env::var("MEMRA_PRIME_CHUNK").ok()
465            .and_then(|v| v.parse().ok()).unwrap_or(4096);
466        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
467        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
468        // the prefill's ARITHMETIC, so two rigs with different values produced different
469        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
470        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
471        // (VERDICT.md) — and it is NOT what docs originally said:
472        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
473        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
474        //     output head), so growing a chunk cannot move an existing row's value.
475        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
476        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
477        //     not describe our leak.
478        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
479        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
480        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
481        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
482        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
483        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
484        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
485        // the source — every row is in one numeric class, so the chunk size no longer steers
486        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
487        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
488        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
489        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
490        // The REQUEST's absolute end position, computed ONCE before the loop: every chunk sees the
491        // same value, whatever the chunk size. step35's SWA arm selects on it so that kernel
492        // selection — and therefore the logits — cannot depend on MEMRA_PRIME_CHUNK
493        // (research/step35-chunkfix-20260807; see step35_attn_pre_wo's doc note).
494        // `+ queued_after` closes the SECOND axis (lane/tick-seg): when serve splits the request
495        // across calls, the request still ends at the same absolute position, whatever the tick
496        // budget or LCP split point. MEMRA_PRIME_CALLLOCAL=1 is the ROLLBACK SEAM to the FULL
497        // pre-fix arithmetic: the per-call value here AND the unaligned FA view offset in
498        // step35_attn_pre_wo. Both halves are required for tickinv35's canary teeth under the FA
499        // default. Read per call, not cached (the probe flips it in-process between arms). Never
500        // on in a measured default run.
501        let legacy_calllocal =
502            std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
503        let seq_end = if legacy_calllocal {
504            cache.pos + t
505        } else {
506            cache.pos + t + queued_after
507        };
508        if chunk == 0 || t <= chunk {
509            return self.prime_chunk(e, tokens, cache, seq_end);
510        }
511        let mut hiddens = e.uninit(t * n_embd)?;
512        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
513        let mut start = 0usize;
514        while start < t {
515            // keep the tail chunk >= PRIME_MIN_T (the stateful conv needs T >= d_conv-1).
516            let mut end = (start + chunk).min(t);
517            if t - end > 0 && t - end < PRIME_MIN_T { end = t; }
518            let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache, seq_end)?;
519            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
520            last = Some((l, hs));
521            start = end;
522        }
523        let (logits, h_seed) = last.unwrap();
524        Ok((logits, h_seed, hiddens))
525    }
526
527    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
528    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
529    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
530    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
531    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
532    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
533    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
534        if Engine::gdn_db_on()
535            && Engine::gdn_chunked_enabled() && t >= 16
536            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
537            && num_k * 2 == num_v
538        {
539            num_k
540        } else {
541            num_v
542        }
543    }
544
545    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
546    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
547    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
548    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
549    fn f16out_on(e: &Engine, t: usize) -> bool {
550        crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
551            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
552    }
553
554    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
555    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
556    /// PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating
557    /// engine's CUDA ordinal — under the prime stage split each stage's range walks through
558    /// its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be
559    /// a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs
560    /// see one entry, byte-identical behavior.
561    pub fn prime_slabs_get(&self, e: &Engine, t: usize, n_embd: usize, n_ff_max: usize)
562                           -> Result<std::sync::MutexGuard<'_, std::collections::HashMap<usize, PrimeSlabs>>, Box<dyn std::error::Error>> {
563        let mut g = self.prime_slabs.lock().unwrap();
564        let dev = e.ctx().ordinal();
565        let need_new = match g.get(&dev) { None => true, Some(sl) => sl.t_cap < t };
566        if need_new {
567            g.insert(dev, PrimeSlabs {
568                t_cap: t,
569                h: e.uninit(t * n_embd)?,
570                x1: e.uninit(t * n_embd)?,
571                z: e.uninit(t * n_embd)?,
572                act: e.uninit(t * n_ff_max)?,
573                xa: e.uninit(t * n_embd)?,
574                xb: e.uninit(t * n_embd)?,
575                h16: e.alloc_u8_uninit(t * n_embd * 2)?,
576                z16: e.alloc_u8_uninit(t * n_embd * 2)?,
577                gate: e.uninit(t * n_ff_max)?,
578                up: e.uninit(t * n_ff_max)?,
579                ffn_out: e.uninit(t * n_embd)?,
580                seg_glue: Vec::new(),
581                mixed: e.uninit(t * n_embd)?,
582                seg_mid: Vec::new(),
583                seg_t: 0,
584            });
585        }
586        Ok(g)
587    }
588
589    /// `seq_end` = the whole REQUEST's absolute end position (`cache.pos + prompt_len` at
590    /// `prime_cache` entry), NOT this chunk's end. Chunk-size-invariant by construction; step35's
591    /// SWA arm selects on it (see `step35_attn_pre_wo`). Every other arch ignores it.
592    fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize)
593                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
594        // LEVER B (lane/pp-leverb, 2026-08-08): the ppN door for the chunked prime. With the
595        // door open + per-stage streams, each chunk walks its layer ranges on the OWNING
596        // stage's engine/device (the anatomy receipt this kills: dev1 ran ZERO prefill
597        // kernels; stage-1 trunk weights were peer-read = 22% of the pp4096 wall).
598        // MEMRA_PRIME_PP=0 = the unsplit rollback (also the prime-split-gate's reference
599        // arm — prime deliberately keeps NO refuse_unsplit_if_remote, see pp.rs). The
600        // MEMRA_PP_STREAMS=0 seam keeps the unsplit walk too: in that regime the sharded
601        // loader is off and there is nothing remote to split for.
602        if self.cfg.gemma4.is_none()
603            && !crate::pp::pp2_streams_off()
604            && crate::pp::prime_pp_on()
605        {
606            if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
607                return self.prime_chunk_ppn(e, tokens, cache, seq_end, &fence);
608            }
609        }
610        let t = tokens.len();
611        let base = cache.pos;
612        debug_assert!(seq_end >= base + t, "prime_chunk: seq_end must cover this chunk");
613        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
614        let pos_d = e.htod_i32(&pos)?;
615
616        let x_embed = self.embed(e, tokens)?;   // [T, n_embd]
617        let x = self.prime_layers(e, x_embed, 0, self.layers.len(), &pos_d, t, cache, seq_end)?;
618        self.prime_chunk_epilogue(e, x, t, cache)
619    }
620
621    /// PRIME RANGE SUBGRAPH (lane/pp-leverb, 2026-08-08): layers `[lo, hi)` of the chunked
622    /// prime walk — `prime_chunk`'s trunk loop extracted verbatim to the
623    /// `decode_layers_eager(lo, hi)` / `verify_layers(lo, hi)` contract: enters with a
624    /// MATERIALIZED `[T, n_embd]` residual, exits with the range's final residual
625    /// materialized (cloned out of the slab). At `lo=0, hi=n_layers` — the unsplit call —
626    /// the launch sequence is byte-identical to the pre-extraction body. Range semantics:
627    ///   - the cross-layer [down-add + NEXT attn-norm] fusion is range-LOCAL (`il + 1 < hi`):
628    ///     layer `hi`'s attn_norm belongs to the NEXT stage's device, so the range ends with
629    ///     the plain add (materialize) and the next stage hoists its own first norm — the
630    ///     kernel-check-pinned `add_rms_norm == add then rms_norm` identity, the same law
631    ///     the decode split rests on (`prime-split-gate` arbitrates end-to-end);
632    ///   - prime slabs are PER-DEVICE (`prime_slabs_get` keys on the engine ordinal), so
633    ///     each stage walks through its own resident transients;
634    ///   - the S-glue/S-mid capture path requires the FULL range (its lookahead fuses
635    ///     `self.layers[il+1]` unconditionally) — `use_seg` gains `lo == 0 && hi == n_layers`.
636    #[allow(clippy::too_many_arguments)]
637    fn prime_layers(&self, e: &Engine, x_in: CudaSlice<f32>, lo: usize, hi: usize,
638                    pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, seq_end: usize)
639                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
640        let cfg = &self.cfg;
641        let n_embd = cfg.n_embd as usize;
642        let eps = cfg.rms_eps;
643        let base = cache.pos;
644        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
645        // standalone convert launches). Only when the f16 lane serves and T reaches the
646        // GEMM tier; bit-identical either way.
647        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
648        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
649        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
650        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
651        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
652        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
653        let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
654            crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
655            _ => n_embd,
656        }).max().unwrap_or(n_embd).max(n_embd);
657        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
658        let mut slab_guard = if use_slabs {
659            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
660        } else {
661            None
662        };
663        let mut x_own;   // fallback storage when slabs are off
664        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>);
665        let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
666        let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
667        let mut x_own2;
668        match slab_guard.as_mut() {
669            Some(g) => {
670                // per-device map (lane/pp-leverb): the getter above populated this engine's
671                // entry; a missing key here is a getter contract bug, fail loudly.
672                let slabs = g.get_mut(&e.ctx().ordinal()).expect("prime slabs for this device");
673                e.copy_into(&mut slabs.xa, 0, &x_in, t * n_embd)?;
674                let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
675                x_cur = xa;
676                x_nxt = xb;
677                seg = Some((seg_glue, seg_mid, mixed, seg_t));
678                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
679            }
680            None => {
681                x_own = x_in;
682                x_own2 = e.uninit(t * n_embd)?;
683                x_cur = &mut x_own;
684                x_nxt = &mut x_own2;
685                sl = None;
686            }
687        }
688        let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
689        let mut alloc_h16; let mut alloc_z16;
690        let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
691        let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
692        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
693        let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
694        match sl {
695            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
696                h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
697                sl_gate = g; sl_up = u; sl_fo = fo;
698            }
699            None => {
700                alloc_h = e.uninit(t * n_embd)?;
701                alloc_x1 = e.uninit(t * n_embd)?;
702                alloc_z = e.uninit(t * n_embd)?;
703                alloc_act = e.uninit(t * n_ff_max)?;
704                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
705                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
706                alloc_gate = e.uninit(t * n_ff_max)?;
707                alloc_up = e.uninit(t * n_ff_max)?;
708                alloc_fo = e.uninit(t * n_embd)?;
709                h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
710                h16 = &mut alloc_h16; z16 = &mut alloc_z16;
711                sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
712            }
713        }
714        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
715        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
716        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
717        // first prime at this t (capture does not execute -> launch right after).
718        let n_layers = self.layers.len();
719        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
720        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
721        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
722        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
723        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
724        // machinery stays (byte-identical) as their foundation.
725        // step35 is excluded: the core-split path calls `full_attn_prime_core_inner`, which is
726        // the GENERIC attn core (uniform n_head, rope_dim_count, no window, no head-wise gate).
727        // step35 rides its own mixer through the normal per-layer arm below.
728        let use_seg = f16fuse && seg.is_some() && self.cfg.step35.is_none()
729            && lo == 0 && hi == n_layers
730            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
731        if let Some((sg, sm, _, st)) = seg.as_mut() {
732            if **st != t {
733                sg.clear();
734                sg.extend((0..n_layers).map(|_| None));
735                sm.clear();
736                sm.extend((0..n_layers).map(|_| None));
737                **st = t;
738            }
739        }
740        {
741            let layer_lo = &self.layers[lo];
742            if f16fuse {
743                e.rms_norm_f16out(x_cur, layer_lo.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
744            } else {
745                e.rms_norm(x_cur, layer_lo.attn_norm.float_data(), h, n_embd, t, eps)?;
746            }
747        }
748        for il in lo..hi {
749            let layer = &self.layers[il];
750            let hx16 = if f16fuse { Some(&*h16) } else { None };
751            if use_seg {
752                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
753                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
754                let (pre, pre16, w_out) = match &layer.mixer {
755                    Mixer::Full(fa) => {
756                        let g3 = match hx16 {
757                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
758                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
759                        };
760                        let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
761                        (pre, pre16, &fa.wo)
762                    }
763                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
764                    Mixer::Linear(la) => {
765                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
766                        let g4 = match hx16 {
767                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
768                            None => e.matmul_group(&ws, h, t)?,
769                        };
770                        let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
771                        (pre, pre16, &la.ssm_out)
772                    }
773                };
774                {
775                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
776                    let pre_n = pre.len() / t;
777                    let xh_pre = match pre16 {
778                        Some(x) => x,
779                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
780                    };
781                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
782                        let y = e.matmul(w_out, &pre, t)?;
783                        e.copy_into(mslab, 0, &y, t * n_embd)?;
784                    }
785                    if sm[il].is_none() {
786                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
787                        let w_post = layer.post_attn_norm.float_data();
788                        e.stream().synchronize()?;
789                        e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
790                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
791                            e.add(x_cur, mslab, x1, t * n_embd)?;
792                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
793                            Ok(())
794                        })();
795                        let g = e.stream().end_capture(
796                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
797                        r?;
798                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
799                    }
800                    sm[il].as_ref().unwrap().launch()?;
801                }
802            } else {
803                let mixed = match &layer.mixer {
804                    Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il,
805                                                            seq_end)?,
806                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
807                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
808                };
809                if f16fuse {
810                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
811                    // bit-identical) — the standalone add pass disappears.
812                    e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
813                                          x1, z, z16, n_embd, t, eps)?;
814                } else {
815                    e.add(x_cur, &mixed, x1, t * n_embd)?;
816                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
817                }
818            }
819            let zx16 = if f16fuse { Some(&*z16) } else { None };
820            match &layer.ffn {
821                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
822                    let n_ff = ffn_gate.out_features();
823                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
824                    // the allocating group + copy when a mirror is missing.
825                    let mut into_ok = false;
826                    if let Some(xh) = zx16 {
827                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
828                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
829                    }
830                    if !into_ok {
831                        let mut g2 = match zx16 {
832                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
833                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
834                        };
835                        let up_y = g2.pop().unwrap();
836                        let gate_y = g2.pop().unwrap();
837                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
838                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
839                    }
840                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
841                    // operand in-epilogue; non-silu activations keep the standalone convert.
842                    // silu_mul_f16out is PLAIN silu(gate)*up — a clamped layer (step35 dense
843                    // blocks under a live swiglu_clamp_shexp) must take the ffn_act_lim arm.
844                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
845                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none()
846                        && d_lim.is_none() {
847                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
848                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
849                        Some(a16)
850                    } else {
851                        Self::ffn_act_lim(e, &self.cfg, sl_gate, sl_up, 1.0, 1.0, d_lim,
852                                          act, t * n_ff)?;
853                        None
854                    };
855                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
856                    let xh_act = match act16 {
857                        Some(x) => x,
858                        None => e.f16_act(act, t * n_ff, n_ff)?,
859                    };
860                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
861                        let y = e.matmul(ffn_down, &*act, t)?;
862                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
863                    }
864                }
865                crate::hybrid::Ffn::Moe(m) => {
866                    let y = self.moe_ffn_il(e, m, z, t, il as u16)?;
867                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
868                }
869            }
870            if use_seg && il + 1 < hi {
871                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
872                let w_next = self.layers[il + 1].attn_norm.float_data();
873                let (sg, _, _, _) = seg.as_mut().unwrap();
874                if sg[il].is_none() {
875                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
876                    e.stream().synchronize()?;
877                    e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
878                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
879                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
880                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
881                        Ok(())
882                    })();
883                    let g = e.stream().end_capture(
884                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
885                    r?;
886                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
887                }
888                sg[il].as_ref().unwrap().launch()?;
889            } else {
890                if il + 1 < hi {
891                    let w_next = self.layers[il + 1].attn_norm.float_data();
892                    if f16fuse {
893                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
894                    } else {
895                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
896                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
897                    }
898                } else {
899                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
900                }
901            }
902            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
903            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
904            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
905            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
906            // unset (the default) costs one OnceLock read per layer.
907            if let Some(path) = Self::prime_trace_path() {
908                let row = (base + t - 1) as usize;
909                let host = e.dtoh(x_nxt)?;
910                let last = &host[(t - 1) * n_embd..t * n_embd];
911                use std::io::Write as _;
912                let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
913                let mut h64: u64 = 0xcbf29ce484222325;
914                for v in last {
915                    h64 ^= v.to_bits() as u64;
916                    h64 = h64.wrapping_mul(0x100000001b3);
917                }
918                writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
919                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
920                         last[0], last[1], last[2])?;
921            }
922            std::mem::swap(&mut x_cur, &mut x_nxt);
923        }
924        // hidden-stack return: clone the final x out of the slab
925        let mut x = e.uninit(t * n_embd)?;
926        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
927        drop(slab_guard);
928        Ok(x)
929    }
930
931    /// The prime chunk's tail — h_seed + output_norm + one-row lm head + cache.pos advance —
932    /// shared verbatim by the unsplit walk and the last stage of the ppN walk (`e` = the
933    /// engine that produced `x`, i.e. the last stage's under the split; output_norm/output
934    /// were loaded through that engine by the sharded loader, hybrid.rs `e_head`).
935    fn prime_chunk_epilogue(&self, e: &Engine, x: CudaSlice<f32>, t: usize, cache: &mut Cache)
936                            -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
937        let n_embd = self.cfg.n_embd as usize;
938        let eps = self.cfg.rms_eps;
939        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
940        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
941        // the post-norm copy happens after hn exists).
942        let mut h_seed = e.uninit(n_embd)?;
943        if !crate::spec::spec_hpost() {
944            e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
945        }
946        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
947        let mut hn = e.uninit(t * n_embd)?;
948        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
949        if crate::spec::spec_hpost() {
950            e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
951        }
952        let last = e.view(&hn, t * n_embd);
953        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
954        let mut hlast = e.uninit(n_embd)?;
955        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
956        let logits = e.matmul(&self.output, &hlast, 1)?;
957        cache.pos += t;
958        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
959        // post-norm stack hn (MEMRA_SPEC_HPOST).
960        Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
961    }
962
963    /// THE PRIME STAGE SPLIT (lane/pp-leverb, 2026-08-08): one prime chunk as N stage
964    /// subgraphs, each range on ITS OWN engine/stream/device, with the `[T, n_embd]`
965    /// boundary handoff at every fence cut — the prime twin of `decode_step_h_ppn` /
966    /// `decode_step_t_core_ppn`, and the kill for the anatomy's two receipts: stage-1 trunk
967    /// weights stop being peer-read (22% of the pp4096 wall) and dev1 stops running zero
968    /// prefill kernels. Structure mirrors the verify split exactly:
969    ///   stage 0        `rt.enter(0)` → per-stage pos_d + embed (the table lives with
970    ///                  stage 0) → `prime_layers(fence[0], fence[1])` → `rt.tx`
971    ///   middle stages  `rt.rx` → per-stage pos_d → range → `rt.tx`
972    ///   last stage     `rt.rx` → range → the shared epilogue (output_norm + head live
973    ///                  there via the sharded loader) → `publish_to`
974    /// Laws inherited (not relearned): `fence_stages_behind` at entry (#87 — a previous
975    /// round's stage-freed buffers must not be reused under the caller's queued reads);
976    /// per-stage `pos_d` (allocated/consumed/freed on one stream); per-stage Engines from
977    /// `PpNRt` (shared-scratch race); EXIT PUBLICATION for the device-resident returns
978    /// (h_seed + hidden stack live on the last stage — the caller's stream must wait).
979    /// KV/MoE locality falls out: each range's KV appends run on the owning stage
980    /// (`pp::new_cache` placed the buffers there), each stage engine owns its own SLRU
981    /// pool (per-Engine `moe_cache`, sized on ITS device — the SGLang #33666 law), and the
982    /// slab-local MoE arm's `DevExps.dev` gate now matches on stage-1 layers too.
983    /// EXACTNESS: the split adds zero deviation by construction (same kernels, same bytes,
984    /// boundary = straight f32 copy); `prime-split-gate` (ppsplit) arbitrates bit-for-bit
985    /// and its liveness counter is bumped here — the gate goes green with this function.
986    fn prime_chunk_ppn(&self, e: &Engine, tokens: &[u32], cache: &mut Cache, seq_end: usize,
987                       fence: &[usize])
988                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
989        let rt = crate::pp::PpNRt::get(e)?;
990        let n_st = fence.len() - 1;
991        assert_eq!(
992            rt.n_stages(), n_st,
993            "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
994        );
995        let n_embd = self.cfg.n_embd as usize;
996        let t = tokens.len();
997        let base = cache.pos;
998        debug_assert!(seq_end >= base + t, "prime_chunk_ppn: seq_end must cover this chunk");
999        let payload = t * n_embd;
1000        // The caller's ambient stream, captured BEFORE any enter() pushes a stage stream
1001        // (inside a stage scope e.stream() IS the stage stream and the exit wait would
1002        // self-order into a no-op) — the decode_step_t_core_ppn pattern.
1003        let caller_stream = e.stream();
1004        rt.fence_stages_behind(&caller_stream)?;
1005        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
1006
1007        // ---- STAGE 0: embed + layers [0, fence[1]) + boundary-0 TX ----
1008        let mut slot = {
1009            let _st0 = rt.enter(0);
1010            let e0 = rt.engine(0, e);
1011            let pos_d = e0.htod_i32(&pos)?;
1012            let x = self.embed(e0, tokens)?;
1013            let x = self.prime_layers(e0, x, fence[0], fence[1], &pos_d, t, cache, seq_end)?;
1014            rt.tx(0, &x, payload)?
1015            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
1016        };
1017
1018        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1019        for s in 1..n_st - 1 {
1020            let _st = rt.enter(s);
1021            let es = rt.engine(s, e);
1022            let pos_d = es.htod_i32(&pos)?;
1023            let x = rt.rx(s - 1, slot, payload)?;
1024            let x = self.prime_layers(es, x, fence[s], fence[s + 1], &pos_d, t, cache, seq_end)?;
1025            slot = rt.tx(s, &x, payload)?;
1026        }
1027
1028        // ---- LAST STAGE: RX + final range + the shared epilogue ----
1029        let _stl = rt.enter(n_st - 1);
1030        let el = rt.engine(n_st - 1, e);
1031        let pos_d = el.htod_i32(&pos)?;
1032        let x = rt.rx(n_st - 2, slot, payload)?;
1033        let x = self.prime_layers(el, x, fence[n_st - 1], fence[n_st], &pos_d, t, cache, seq_end)?;
1034        let out = self.prime_chunk_epilogue(el, x, t, cache)?;
1035        // EXIT PUBLICATION: h_seed + the hidden stack are device-resident on the last
1036        // stage's stream; the caller resumes on its own stream (chunk-loop copy_into /
1037        // generate_spec's prompt_h consumer). The logits dtoh above already drained the
1038        // stage stream host-side, but the law is stated in events, not in a dtoh side
1039        // effect a later deferred form would remove.
1040        rt.publish_to(n_st - 1, &caller_stream)?;
1041        crate::pp::PRIME_SPLIT_CHUNKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1042        Ok(out)
1043    }
1044
1045    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
1046    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
1047    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
1048    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
1049    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
1050    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
1051    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
1052    /// bookkeeping still runs on the host per call — the real replay path moves the write
1053    /// slot to the len_d device counter (increment 3).
1054    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
1055    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
1056    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
1057    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
1058    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
1059    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
1060    pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
1061                                t: usize, cache: &mut Cache,
1062                                len_d: &CudaSlice<i32>,
1063                                logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
1064                                -> Result<(), Box<dyn std::error::Error>> {
1065        let cfg = &self.cfg;
1066        let n_embd = cfg.n_embd as usize;
1067        let eps = cfg.rms_eps;
1068        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
1069        let mut x = e.uninit(t * n_embd)?;
1070        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
1071        for (il, layer) in self.layers.iter().enumerate() {
1072            let mut h = e.uninit(t * n_embd)?;
1073            let mut hx16: Option<CudaSlice<u8>> = None;
1074            if f16fuse {
1075                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1076                e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
1077                hx16 = Some(b16);
1078            } else {
1079                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
1080            }
1081            let mixed = match &layer.mixer {
1082                // The captured prime is ONE unchunked bucket over a FRESH cache (pos == 0), so the
1083                // request ends at t. If bucketed capture ever composes with chunking, seq_end must
1084                // come from the caller (see step35_attn_pre_wo's doc note).
1085                Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache,
1086                                                        il, t)?,
1087                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1088                Mixer::Linear(la) => {
1089                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1090                    let g4 = match hx16.as_ref() {
1091                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
1092                        None => e.matmul_group(&ws, &h, t)?,
1093                    };
1094                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
1095                }
1096            };
1097            let mut x1 = e.uninit(t * n_embd)?;
1098            e.add(&x, &mixed, &mut x1, t * n_embd)?;
1099            let mut z = e.uninit(t * n_embd)?;
1100            let mut zx16: Option<CudaSlice<u8>> = None;
1101            if f16fuse {
1102                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
1103                e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
1104                zx16 = Some(b16);
1105            } else {
1106                e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
1107            }
1108            let ffn_out = match &layer.ffn {
1109                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1110                    let n_ff = ffn_gate.out_features();
1111                    let mut g2 = match &zx16 {
1112                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
1113                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
1114                    };
1115                    let up = g2.pop().unwrap();
1116                    let gate = g2.pop().unwrap();
1117                    let mut act = e.uninit(t * n_ff)?;
1118                    // dense FFN keys off the SHEXP clamp array — see forward()'s note.
1119                    Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0,
1120                                      self.cfg.clamp_shexp_at(il as u32), &mut act, t * n_ff)?;
1121                    e.matmul(ffn_down, &act, t)?
1122                }
1123                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
1124            };
1125            let mut x2 = e.uninit(t * n_embd)?;
1126            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
1127            x = x2;
1128        }
1129        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
1130        if !crate::spec::spec_hpost() {
1131            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
1132        }
1133        let mut hn = e.uninit(t * n_embd)?;
1134        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
1135        if crate::spec::spec_hpost() {
1136            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
1137        }
1138        let mut hlast = e.uninit(n_embd)?;
1139        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
1140        let logits = e.matmul(&self.output, &hlast, 1)?;
1141        let nv = logits.len();
1142        e.copy_into(logits_out, 0, &logits, nv)?;
1143        Ok(())
1144    }
1145
1146    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
1147    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
1148    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
1149    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
1150    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
1151    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
1152    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
1153    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
1154    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
1155    /// over the quantized past; Linear: the stateful pad_view twin — the same state
1156    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
1157    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
1158    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
1159    /// back to single-chunk serving).
1160    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
1161    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
1162    pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
1163                             -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
1164        let cfg = &self.cfg;
1165        let n_embd = cfg.n_embd as usize;
1166        let eps = cfg.rms_eps;
1167        let b = prompts.len();
1168        assert!(b >= 1 && b == caches.len());
1169        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
1170        let carried = pos0s.iter().any(|&p| p > 0);
1171        // gemma4: refuse UNCONDITIONALLY (2026-08-07, lane/gemma4-serve-gaps). The old guard
1172        // covered only `carried` — two concurrent FRESH gemma4 prompts batched into the
1173        // generic concat attn core below (uniform geometry, no per-layer swa window, no
1174        // softcapped head): compiles, runs, wrong logits. Same silent-wrong class as the
1175        // step35 refusal beneath. The per-sequence `gemma4_prime` is the supported prefill.
1176        if cfg.gemma4.is_some() {
1177            return Err("prime_cache_batch: gemma4 has no batched prime core (per-layer \
1178                        swa/global geometry, softcapped head) — use gemma4_prime per sequence".into());
1179        }
1180        // step35: the cross-request batch driver splits the concat projection outputs and feeds
1181        // them to `full_attn_prime_core_inner` (the GENERIC attn core — uniform n_head, 128-dim
1182        // rope on every layer, no window, no head-wise gate). Running step35 through it compiles
1183        // and produces plausible-but-wrong logits, so refuse until a step35 varlen core exists.
1184        // The single-sequence `prime_cache` path is the supported prefill.
1185        if cfg.step35.is_some() {
1186            return Err("prime_cache_batch: step35 has no batched prime core (per-layer n_head / \
1187                        partial rope / SWA / head-wise gate) — use prime_cache per sequence".into());
1188        }
1189        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
1190        for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
1191        for (s, c) in caches.iter().enumerate() {
1192            assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
1193        }
1194        let total: usize = ts.iter().sum();
1195        let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
1196        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
1197        let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
1198            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
1199            .collect::<Result<_, _>>()?;
1200        // split a concat [total, dim] buffer into per-seq copies
1201        let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
1202                     -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
1203            let mut out = Vec::with_capacity(b);
1204            for s in 0..b {
1205                let mut ys = e.uninit(ts[s] * dim)?;
1206                e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
1207                out.push(ys);
1208            }
1209            Ok(out)
1210        };
1211
1212        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
1213        let mut x = self.embed(e, &cat_tokens)?;   // [total, n_embd]
1214        for (il, layer) in self.layers.iter().enumerate() {
1215            let mut h = e.uninit(total * n_embd)?;
1216            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1217            e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
1218            // mixer: projection GROUP on the concat (m = total), stateful core per seq
1219            let mut mixed = e.uninit(total * n_embd)?;
1220            match &layer.mixer {
1221                Mixer::Full(fa) => {
1222                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
1223                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
1224                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
1225                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
1226                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
1227                    // back to the per-seq dispatch.
1228                    let (n_head, n_head_kv, head_dim) =
1229                        (self.cfg.n_head as usize, self.cfg.n_head_kv as usize, self.cfg.head_dim_k as usize);
1230                    let fa_scale = 1.0 / (head_dim as f32).sqrt();
1231                    let use_favl = !carried
1232                        && (2..=8).contains(&b)
1233                        && (head_dim == 256 || head_dim == 128)
1234                        && self.cfg.attn_out_gate()
1235                        && std::env::var("MEMRA_NOFA").is_err()
1236                        && std::env::var("MEMRA_FA_FLOOR").is_err()
1237                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
1238                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
1239                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
1240                    if use_favl {
1241                        let (qf_w, kf_w, vf_w) =
1242                            (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
1243                        struct APre {
1244                            q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
1245                            qn: CudaSlice<f32>, kn: CudaSlice<f32>,
1246                        }
1247                        let mut aps = Vec::with_capacity(b);
1248                        for &t in ts.iter().take(b) {
1249                            aps.push(APre {
1250                                q: e.uninit(t * n_head * head_dim)?,
1251                                gate: Some(e.uninit(t * n_head * head_dim)?),
1252                                qn: e.uninit(t * n_head * head_dim)?,
1253                                kn: e.uninit(t * n_head_kv * head_dim)?,
1254                            });
1255                        }
1256                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
1257                            let kvl = caches[0].kv[il].as_ref().unwrap();
1258                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
1259                        };
1260                        let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
1261                            let (o, t) = (offs[s], ts[s]);
1262                            let kvl = caches[s].kv[il].as_ref().unwrap();
1263                            assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
1264                                    "prime_cache_batch attn vl: fresh + capacity");
1265                            crate::AttnPreVl {
1266                                qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
1267                                kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
1268                                vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
1269                                q: e.addr_f32(&aps[s].q),
1270                                gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
1271                                qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
1272                                kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
1273                                t: t as i32, pad: 0,
1274                            }
1275                        }).collect();
1276                        e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
1277                                       head_dim, self.cfg.rope_dim_count as usize, n_head, n_head_kv,
1278                                       self.cfg.rms_eps, self.cfg.rope_freq_base, 1.0,
1279                                       kv_dim_k, kv_dim_v, ktb, vtb)?;
1280                        for s in 0..b {
1281                            let kvl = caches[s].kv[il].as_mut().unwrap();
1282                            kvl.len += ts[s];
1283                            let new_len = kvl.len as i32;
1284                            e.set_i32_one(&mut kvl.len_d, new_len)?;
1285                        }
1286                        let mut attns = Vec::with_capacity(b);
1287                        let mut mirrors = Vec::with_capacity(b);
1288                        for &t in ts.iter().take(b) {
1289                            attns.push(e.uninit(t * n_head * head_dim)?);
1290                            let n = t * n_head_kv * head_dim;
1291                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
1292                        }
1293                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
1294                        // promoted single-seq config is on; else the mma favl.
1295                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
1296                            Ok("0") => false,
1297                            Ok("1") => true,
1298                            _ => cfg!(memra_hopper_mma),
1299                        };
1300                        if fa3_on {
1301                            let mut q16s = Vec::with_capacity(b);
1302                            let mut v16s = Vec::with_capacity(b);
1303                            for s in 0..b {
1304                                let t = ts[s];
1305                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
1306                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
1307                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1308                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
1309                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1310                                e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
1311                                                &mut v16, t * n_head_kv * head_dim)?;
1312                                q16s.push(q16);
1313                                v16s.push((k16, v16));
1314                            }
1315                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
1316                            let mut kp = qp;
1317                            let mut vp = qp;
1318                            let mut op = [core::ptr::null_mut::<f32>(); 8];
1319                            let mut tsv = [0i32; 8];
1320                            for s in 0..b {
1321                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
1322                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
1323                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
1324                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
1325                                tsv[s] = ts[s] as i32;
1326                            }
1327                            let rc = unsafe {
1328                                crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
1329                                                  tsv.as_ptr(), b as i32, n_head as i32,
1330                                                  n_head_kv as i32, head_dim as i32, fa_scale,
1331                                                  e.stream().cu_stream() as *mut core::ffi::c_void)
1332                            };
1333                            if rc != 0 {
1334                                return Err(format!("memra_fa3_vl rc={rc}").into());
1335                            }
1336                        } else {
1337                            let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
1338                                q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
1339                                v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
1340                                kf: e.addr_f32(&aps[s].kn),
1341                                vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
1342                                t: ts[s] as i32, pad: 0,
1343                            }).collect();
1344                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
1345                        }
1346                        for (s, attn) in attns.into_iter().enumerate() {
1347                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
1348                                e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
1349                            let mut done = false;
1350                            if let Some(xh) = &ag16 {
1351                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1352                            }
1353                            if !done {
1354                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1355                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1356                            }
1357                        }
1358                    } else {
1359                        let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
1360                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
1361                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1362                                parts[s].push(ys);
1363                            }
1364                        }
1365                        for (s, g3s) in parts.into_iter().enumerate() {
1366                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
1367                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
1368                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
1369                            let mut done = false;
1370                            if let Some(xh) = &ag16 {
1371                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1372                            }
1373                            if !done {
1374                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1375                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1376                            }
1377                        }
1378                    }
1379                }
1380                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1381                Mixer::Linear(la) => {
1382                    // task #16: NO split copies (cores read row-offset views of the concat
1383                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
1384                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
1385                    // varlen K5 launch for all sequences.
1386                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1387                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
1388                    let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
1389                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
1390                        let (o, t) = (offs[s], ts[s]);
1391                        let mut done = false;
1392                        if let Some(xh) = &gn16 {
1393                            done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
1394                        }
1395                        if !done {
1396                            let m = e.matmul(&la.ssm_out, &gn, t)?;
1397                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
1398                        }
1399                    }
1400                }
1401            }
1402            let mut x1 = e.uninit(total * n_embd)?;
1403            let mut z = e.uninit(total * n_embd)?;
1404            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1405            e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
1406                                  &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
1407            let ffn_out = match &layer.ffn {
1408                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1409                    let n_ff = ffn_gate.out_features();
1410                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
1411                    let up = g2.pop().unwrap();
1412                    let gate = g2.pop().unwrap();
1413                    let mut act = e.uninit(total * n_ff)?;
1414                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
1415                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
1416                    // A clamped layer must skip the plain-SiLU twin (see prime_chunk's note).
1417                    let d_lim = self.cfg.clamp_shexp_at(il as u32);
1418                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() && d_lim.is_none() {
1419                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1420                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1421                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1422                            Some(y) => y,
1423                            None => e.matmul(ffn_down, &act, total)?,
1424                        }
1425                    } else {
1426                        Self::ffn_act_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, d_lim,
1427                                          &mut act, total * n_ff)?;
1428                        e.matmul(ffn_down, &act, total)?
1429                    }
1430                }
1431                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1432            };
1433            let mut x2 = e.uninit(total * n_embd)?;
1434            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1435            x = x2;
1436        }
1437        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
1438        let mut hn = e.uninit(total * n_embd)?;
1439        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
1440        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
1441        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
1442        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
1443        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
1444        // argmax battery arbitrates, same as every other prefill GEMM change.
1445        let mut hcat = e.uninit(b * n_embd)?;
1446        for s in 0..b {
1447            let last0 = (offs[s] + ts[s] - 1) * n_embd;
1448            e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
1449        }
1450        let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
1451        let logits_host: Option<Vec<f32>> = match &logits_cat {
1452            Some(lc) => Some(e.dtoh(lc)?),
1453            None => None,
1454        };
1455        let n_vocab = self.output.out_features();
1456        let mut hidden_all = if crate::spec::spec_hpost() {
1457            split(e, &hn, n_embd)?
1458        } else {
1459            split(e, &x, n_embd)?
1460        };
1461        let mut out = Vec::with_capacity(b);
1462        for s in 0..b {
1463            let last0 = (offs[s] + ts[s] - 1) * n_embd;
1464            let mut h_seed = e.uninit(n_embd)?;
1465            if !crate::spec::spec_hpost() {
1466                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
1467            } else {
1468                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1469            }
1470            let logits = match &logits_host {
1471                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
1472                None => {
1473                    let mut hlast = e.uninit(n_embd)?;
1474                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1475                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
1476                }
1477            };
1478            caches[s].pos += ts[s];
1479            out.push((logits, h_seed, hidden_all.remove(0)));
1480        }
1481        Ok(out)
1482    }
1483
1484    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
1485    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
1486    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
1487    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
1488    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
1489    ///
1490    /// `seq_end` = the ABSOLUTE end position of the WHOLE prime request (`cache.pos + prompt_len`
1491    /// measured BEFORE the chunk loop starts), i.e. a chunk-size-invariant property of the request.
1492    /// Only step35's SWA arm reads it (`step35_attn_pre_wo`'s doc note explains why keying on the
1493    /// chunk's own `t_kv` made the output depend on MEMRA_PRIME_CHUNK); every other arch ignores it.
1494    #[allow(clippy::too_many_arguments)]
1495    fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
1496                       hx: Option<&CudaSlice<u8>>,
1497                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize,
1498                       seq_end: usize)
1499                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1500        if self.cfg.step35.is_some() {
1501            return self.step35_attn_prime(e, fa, h, hx, pos_d, t, cache, il, seq_end);
1502        }
1503        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
1504        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
1505        // this single-seq path composes proj+core identically (byte-for-byte the old body).
1506        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
1507        let g3 = match hx {
1508            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1509            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1510        };
1511        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
1512    }
1513
1514    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
1515    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
1516    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
1517    fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1518                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1519                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1520        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
1521        if let Some(xh) = &ag16 {
1522            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
1523                return Ok(y);
1524            }
1525        }
1526        Ok(e.matmul(&fa.wo, &attn_g, t)?)
1527    }
1528
1529    fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1530                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1531                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1532        let cfg = &self.cfg;
1533        let n_head = cfg.n_head as usize;
1534        let n_head_kv = cfg.n_head_kv as usize;
1535        let head_dim = cfg.head_dim_k as usize;
1536        let scale = 1.0 / (head_dim as f32).sqrt();
1537        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
1538        let AttnPre { q, k, v, gate } = pre;
1539        let mut attn = e.uninit(t * n_head * head_dim)?;
1540        self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
1541                                         head_dim, n_head, n_head_kv, scale)?;
1542        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
1543    }
1544
1545    /// task #18 (attn side): projections tail through KV append — everything before the
1546    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
1547    /// present BEFORE this chunk's append (base_len; 0 == fresh).
1548    #[allow(clippy::type_complexity)]
1549    fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
1550                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1551                            -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
1552        let cfg = &self.cfg;
1553        let n_head = cfg.n_head as usize;
1554        let n_head_kv = cfg.n_head_kv as usize;
1555        let head_dim = cfg.head_dim_k as usize;
1556        let eps = cfg.rms_eps;
1557
1558        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
1559        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
1560        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
1561        let gated = cfg.attn_out_gate();
1562        let v = g3.pop().unwrap();
1563        let mut k = g3.pop().unwrap();
1564        let qf = g3.pop().unwrap();
1565        let (mut q, gate) = if gated {
1566            let mut q = e.uninit(t * n_head * head_dim)?;
1567            let mut gate = e.uninit(t * n_head * head_dim)?;
1568            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
1569            (q, Some(gate))
1570        } else {
1571            (qf, None)
1572        };
1573
1574        let mut qn = e.uninit(t * n_head * head_dim)?;
1575        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
1576        q = qn;
1577        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
1578        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
1579        k = kn;
1580        let rope_dims = cfg.rope_dim_count as usize;
1581        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
1582        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
1583
1584        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
1585        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
1586        {
1587            let kvl = cache.kv[il].as_mut().unwrap();
1588            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
1589            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
1590                                       kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
1591                                       crate::Engine::kv_fp8_on())?;
1592            kvl.len += t;
1593            let new_len = kvl.len as i32;
1594            e.set_i32_one(&mut kvl.len_d, new_len)?;
1595        }
1596
1597        let base_len = {
1598            let kvl = cache.kv[il].as_ref().unwrap();
1599            kvl.len - t   // KV rows present BEFORE this chunk's append above
1600        };
1601        Ok((AttnPre { q, k, v, gate }, base_len))
1602    }
1603
1604    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
1605    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
1606    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
1607    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
1608    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
1609    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
1610    #[allow(clippy::too_many_arguments)]
1611    fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
1612                            v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
1613                            t: usize, cache: &mut Cache, il: usize,
1614                            head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
1615                            -> Result<(), Box<dyn std::error::Error>> {
1616        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
1617        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
1618        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
1619        // attend through the quantized cache exactly like every later chunk (quantize-then-
1620        // attend). One numeric class for every row => the chunk size cannot decide where a
1621        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
1622        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
1623        // pin-the-boundary approach).
1624        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
1625        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
1626        // with the fix unconditional, only re-introducing the class edge can prove the gate
1627        // still detects the mechanism. Never on in a measured default run.
1628        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
1629            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1630                e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1631            } else {
1632                e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1633            }
1634            return Ok(());
1635        }
1636        let kvl = cache.kv[il].as_ref().unwrap();
1637        let t_kv = base_len + t;
1638        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
1639        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
1640        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
1641        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
1642        // same numeric class, so the uniform contract holds on the fallback too.
1643        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1644            e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
1645                                        n_head_kv, t, t_kv, scale, true,
1646                                        kvl.k_tok_bytes, kvl.v_tok_bytes)?;
1647            return Ok(());
1648        }
1649        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
1650        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
1651        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
1652        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
1653        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
1654        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
1655        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
1656        let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
1657        if deqw {
1658            e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1659                                 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1660                                 crate::Engine::kv_fp8_on())?;
1661        } else {
1662            e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1663                              t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1664                              crate::Engine::kv_fp8_on())?;
1665        }
1666        Ok(())
1667    }
1668
1669    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
1670    /// (bit-identical composition) and hands wo its fp16 operand directly.
1671    fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
1672                            gate: &Option<CudaSlice<f32>>, t: usize,
1673                            n_head: usize, head_dim: usize)
1674                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1675        let (attn_g, ag16) = match gate {
1676            Some(gate) => {
1677                let n = t * n_head * head_dim;
1678                let mut ag = e.uninit(n)?;
1679                if Self::f16out_on(e, t) {
1680                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
1681                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
1682                    (ag, Some(a16))
1683                } else {
1684                    let mut gsig = e.uninit(n)?;
1685                    e.sigmoid(gate, &mut gsig, n)?;
1686                    e.mul(&attn, &gsig, &mut ag, n)?;
1687                    (ag, None)
1688                }
1689            }
1690            None => (attn, None),
1691        };
1692        Ok((attn_g, ag16))
1693    }
1694
1695    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
1696    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
1697    /// carried THROUGH the cache like the spec verify does: carried-ring conv
1698    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
1699    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
1700    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
1701    fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
1702                         hx: Option<&CudaSlice<u8>>, t: usize,
1703                         cache: &mut Cache, il: usize)
1704                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1705        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
1706        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1707        let g4 = match hx {
1708            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1709            None => e.matmul_group(&ws, h, t)?,
1710        };
1711        self.linear_attn_prime_core(e, la, g4, t, cache, il)
1712    }
1713
1714    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
1715    fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1716                              t: usize, cache: &mut Cache, il: usize)
1717                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1718        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
1719    }
1720
1721    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
1722    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
1723    /// conv ring writes back from the true tail. None = classic path, byte-identical.
1724    #[allow(clippy::too_many_arguments)]
1725    fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1726                              t: usize, cache: &mut Cache, il: usize,
1727                              pad_len: Option<&CudaSlice<i32>>)
1728                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1729        // shim over the view twin (task #16): full-range views of the owned buffers.
1730        let ssm = self.cfg.ssm.as_ref().unwrap();
1731        let d_state = ssm.state_size as usize;
1732        let num_k = ssm.group_count as usize;
1733        let num_v = ssm.time_step_rank as usize;
1734        let key_dim = d_state * num_k;
1735        let value_dim = d_state * num_v;
1736        let conv_dim = key_dim * 2 + value_dim;
1737        let alpha = g4.pop().unwrap();                   // [T, num_v]
1738        let beta_raw = g4.pop().unwrap();                // [T, num_v]
1739        let z = g4.pop().unwrap();                       // [T, value_dim]
1740        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
1741        self.linear_attn_prime_core_pad_view(
1742            e, la,
1743            &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
1744            &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
1745            t, cache, il, pad_len)
1746    }
1747
1748    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
1749    /// shared verbatim by the per-seq scan path and the varlen batched path.
1750    #[allow(clippy::too_many_arguments)]
1751    fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
1752                            qkv_mixed: &cudarc::driver::CudaView<f32>,
1753                            beta_raw: &cudarc::driver::CudaView<f32>,
1754                            alpha: &cudarc::driver::CudaView<f32>,
1755                            t: usize, cache: &mut Cache, il: usize,
1756                            pad_len: Option<&CudaSlice<i32>>)
1757                            -> Result<GdnPrep, Box<dyn std::error::Error>> {
1758        let cfg = &self.cfg;
1759        let ssm = cfg.ssm.as_ref().unwrap();
1760        let d_state = ssm.state_size as usize;       // 128
1761        let num_k = ssm.group_count as usize;        // 16
1762        let num_v = ssm.time_step_rank as usize;     // 32
1763        let d_conv = ssm.conv_kernel as usize;       // 4
1764        let key_dim = d_state * num_k;               // 2048
1765        let value_dim = d_state * num_v;             // 4096
1766        let conv_dim = key_dim * 2 + value_dim;      // 8192
1767        let eps = cfg.rms_eps;
1768        debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
1769
1770        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
1771        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
1772        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
1773        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
1774        let rl = cache.recur[il].as_mut().unwrap();
1775        let hk = Self::gdn_hk(e, t, num_v, num_k);
1776        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
1777        let hk = if conv_fuse { hk } else { num_v };   // de-broadcast rides the fused conv
1778        let mut q_g = e.uninit(d_state * hk * t)?;
1779        let mut k_g = e.uninit(d_state * hk * t)?;
1780        let mut v_g = e.uninit(d_state * num_v * t)?;
1781        if conv_fuse {
1782            e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
1783                                  &mut q_g, &mut k_g, &mut v_g,
1784                                  conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
1785        } else {
1786            let mut conv_out = e.uninit(conv_dim * t)?;      // [conv_dim, T] channel-major, SiLU
1787            e.ssm_conv1d_tm_state_pad_v(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
1788                                  &mut conv_out, conv_dim, t, d_conv, pad_len)?;
1789            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)?;
1790        }
1791        let mut q_l2 = e.uninit(d_state * hk * t)?;
1792        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
1793        // Emitted only where a consumer exists (the wgmma config) — on other arches the
1794        // alloc + epilogue stores would be pure waste.
1795        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
1796            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1797            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
1798            Some(qb)
1799        } else {
1800            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
1801            None
1802        };
1803        let mut k_l2 = e.uninit(d_state * hk * t)?;
1804        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
1805        let kb16 = if Engine::l2_v2_on(d_state) {
1806            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1807            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
1808            Some(kb)
1809        } else {
1810            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
1811            None
1812        };
1813        let mut beta = e.uninit(t * num_v)?;
1814        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
1815        let mut g_log = e.uninit(t * num_v)?;
1816        e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
1817        if let Some(len_d) = pad_len {
1818            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
1819        }
1820        Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
1821    }
1822
1823    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
1824    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
1825    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
1826    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
1827    #[allow(clippy::too_many_arguments)]
1828    fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
1829                                    g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
1830                                    caches: &mut [&mut Cache], il: usize)
1831                                    -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
1832        let ssm = self.cfg.ssm.as_ref().unwrap();
1833        let d_state = ssm.state_size as usize;
1834        let num_k = ssm.group_count as usize;
1835        let num_v = ssm.time_step_rank as usize;
1836        let key_dim = d_state * num_k;
1837        let value_dim = d_state * num_v;
1838        let conv_dim = key_dim * 2 + value_dim;
1839        let eps = self.cfg.rms_eps;
1840        let scale = 1.0 / (d_state as f32).sqrt();
1841        let b = ts.len();
1842        let c = Engine::gdn_chunk_size();
1843        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
1844        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
1845        let carried = caches.iter().any(|c| c.pos > 0);
1846        let use_vl = !carried
1847            && (2..=8).contains(&b)
1848            && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
1849            && e.gdn_mma_enabled(c)
1850            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
1851        if !use_vl {
1852            return (0..b).map(|s| {
1853                let (o, t) = (offs[s], ts[s]);
1854                self.linear_attn_prime_core_pad_view(
1855                    e, la,
1856                    &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
1857                    &g4[1].slice(o * value_dim..(o + t) * value_dim),
1858                    &g4[2].slice(o * num_v..(o + t) * num_v),
1859                    &g4[3].slice(o * num_v..(o + t) * num_v),
1860                    t, caches[s], il, None)
1861            }).collect();
1862        }
1863        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
1864        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
1865        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
1866        struct SeqBufs {
1867            conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
1868            q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
1869            gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
1870        }
1871        let d_conv = ssm.conv_kernel as usize;
1872        let f16o = Self::f16out_on(e, 16);
1873        let hk = Self::gdn_hk(e, 16, num_v, num_k);   // vl path is always chunked+mma
1874        let mut sb = Vec::with_capacity(b);
1875        let mut pres = Vec::with_capacity(b);
1876        for &t in ts.iter().take(b) {
1877            sb.push(SeqBufs {
1878                conv_out: e.uninit(conv_dim * t)?,
1879                q_g: e.uninit(d_state * hk * t)?,
1880                k_g: e.uninit(d_state * hk * t)?,
1881                v_g: e.uninit(d_state * num_v * t)?,
1882                q_l2: e.uninit(d_state * hk * t)?,
1883                k_l2: e.uninit(d_state * hk * t)?,
1884                beta: e.uninit(t * num_v)?,
1885                g_log: e.uninit(t * num_v)?,
1886                gn: e.uninit(d_state * num_v * t)?,
1887                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
1888            });
1889            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
1890        }
1891        let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
1892            let (o, t) = (offs[s], ts[s]);
1893            let rl = caches[s].recur[il].as_ref().unwrap();
1894            crate::GdnPrepVl {
1895                qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
1896                conv_state: e.addr_f32(&rl.conv_state),
1897                conv_out: e.addr_f32(&sb[s].conv_out),
1898                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),
1899                q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
1900                beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
1901                alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
1902                beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
1903                o: e.addr_f32(&pres[s].o),
1904                z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
1905                gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
1906                kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
1907                qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
1908                t: t as i32, pad: 0,
1909            }
1910        }).collect();
1911        let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
1912            let rl = caches[s].recur[il].as_ref().unwrap();
1913            crate::GdnSeqVl {
1914                kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
1915                beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
1916                wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
1917                ssnap: e.addr_u8(&pres[s].ssnap16),
1918                state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
1919                q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
1920                o: e.addr_f32(&pres[s].o),
1921                k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
1922                g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
1923                w: e.addr_f32(&pres[s].w),
1924                t: ts[s] as i32, nc: pres[s].nc as i32,
1925            }
1926        }).collect();
1927        e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
1928                       la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
1929        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
1930        // both standalone mirror launches vanish on the default config.
1931        if !Engine::l2_v2_on(d_state) {
1932            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
1933        }
1934        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
1935        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
1936            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
1937            if !Engine::l2_v2_on(d_state) {
1938                for s in 0..b {
1939                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
1940                }
1941            }
1942            let mut wa = [crate::GdnWVl::default(); 8];
1943            for s in 0..b {
1944                wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
1945            }
1946            Some(crate::GdnWVl8(wa))
1947        } else { None };
1948        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
1949        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
1950        if f16o {
1951            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
1952        }
1953        // per-seq state swap (+ non-f16out tail fallback)
1954        let mut out = Vec::with_capacity(b);
1955        for (s, bufs) in sb.into_iter().enumerate() {
1956            let rl = caches[s].recur[il].as_mut().unwrap();
1957            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1958            let (o, t) = (offs[s], ts[s]);
1959            let SeqBufs { mut gn, gn16, .. } = bufs;
1960            if f16o {
1961                out.push((gn, Some(gn16)));
1962            } else {
1963                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
1964                e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
1965                                   d_state, num_v * t, eps)?;
1966                out.push((gn, None));
1967            }
1968        }
1969        Ok(out)
1970    }
1971
1972    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
1973    /// views of the CONCAT projection outputs directly (no per-seq split copies).
1974    /// Same kernels, same values, byte-identical to the Vec shim above.
1975    #[allow(clippy::too_many_arguments)]
1976    fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
1977                              qkv_mixed: &cudarc::driver::CudaView<f32>,
1978                              z: &cudarc::driver::CudaView<f32>,
1979                              beta_raw: &cudarc::driver::CudaView<f32>,
1980                              alpha: &cudarc::driver::CudaView<f32>,
1981                              t: usize, cache: &mut Cache, il: usize,
1982                              pad_len: Option<&CudaSlice<i32>>)
1983                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1984        let cfg = &self.cfg;
1985        let ssm = cfg.ssm.as_ref().unwrap();
1986        let d_state = ssm.state_size as usize;       // 128
1987        let num_v = ssm.time_step_rank as usize;     // 32
1988        let eps = cfg.rms_eps;
1989        let scale = 1.0 / (d_state as f32).sqrt();
1990
1991        let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
1992
1993        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
1994        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
1995        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
1996        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
1997        // verify keep the sequential kernel).
1998        let mut o = e.uninit(d_state * num_v * t)?;
1999        let rl = cache.recur[il].as_mut().unwrap();
2000        {
2001            let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
2002            e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
2003                               prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
2004                               prep.hk)?;
2005        }
2006        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2007
2008        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
2009        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
2010        let mut gn = e.uninit(d_state * num_v * t)?;
2011        let gn16 = if Self::f16out_on(e, t) {
2012            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
2013            e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
2014                                      d_state, num_v * t, eps)?;
2015            Some(g16)
2016        } else {
2017            e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
2018            None
2019        };
2020        Ok((gn, gn16))
2021    }
2022
2023    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
2024    #[allow(clippy::too_many_arguments)]
2025    fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
2026                              t: usize, cache: &mut Cache, il: usize,
2027                              pad_len: Option<&CudaSlice<i32>>)
2028                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2029        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
2030        if let Some(xh) = &gn16 {
2031            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
2032                return Ok(y);
2033            }
2034        }
2035        Ok(e.matmul(&la.ssm_out, &gn, t)?)
2036    }
2037
2038    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
2039    ///
2040    /// `il` = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and
2041    /// routes to its own mixer. Every other arch ignores it (uniform geometry).
2042    pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize)
2043                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2044        if self.cfg.step35.is_some() {
2045            return self.step35_attn(e, fa, h, pos_d, t, il);
2046        }
2047        let cfg = &self.cfg;
2048        let _n_embd = cfg.n_embd as usize;
2049        let n_head = cfg.n_head as usize;
2050        let n_head_kv = cfg.n_head_kv as usize;
2051        let head_dim = cfg.head_dim_k as usize;
2052        let eps = cfg.rms_eps;
2053        let scale = 1.0 / (head_dim as f32).sqrt();
2054
2055        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
2056        // gate — wq out = n_head*head_dim, no split (see prime-path note).
2057        let gated = cfg.attn_out_gate();
2058        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
2059        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
2060        let v = g3.pop().unwrap();
2061        let mut k = g3.pop().unwrap();
2062        let qf = g3.pop().unwrap();
2063        let (mut q, gate) = if gated {
2064            let mut q = e.uninit(t * n_head * head_dim)?;
2065            let mut gate = e.uninit(t * n_head * head_dim)?;
2066            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
2067            (q, Some(gate))
2068        } else {
2069            (qf, None)
2070        };
2071
2072        // QK-norm (per head_dim row), then partial RoPE.
2073        let mut qn = e.uninit(t * n_head * head_dim)?;
2074        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
2075        q = qn;
2076        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
2077        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
2078        k = kn;
2079        let rope_dims = cfg.rope_dim_count as usize;
2080        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
2081        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
2082
2083        // SDPA
2084        let mut attn = e.uninit(t * n_head * head_dim)?;
2085        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
2086        // falls back to naive sdpa.
2087        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
2088            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
2089            e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2090        } else {
2091            e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
2092        }
2093
2094        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
2095        let attn_g = match &gate {
2096            Some(gate) => {
2097                let mut gsig = e.uninit(t * n_head * head_dim)?;
2098                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
2099                let mut ag = e.uninit(t * n_head * head_dim)?;
2100                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
2101                ag
2102            }
2103            None => attn,
2104        };
2105
2106        // o projection
2107        let o = e.matmul(&fa.wo, &attn_g, t)?;
2108        Ok(o)
2109    }
2110
2111    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
2112    pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
2113                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2114        let cfg = &self.cfg;
2115        let _n_embd = cfg.n_embd as usize;
2116        let ssm = cfg.ssm.as_ref().unwrap();
2117        let d_state = ssm.state_size as usize;       // 128
2118        let num_k = ssm.group_count as usize;        // 16
2119        let num_v = ssm.time_step_rank as usize;     // 32
2120        let d_conv = ssm.conv_kernel as usize;       // 4
2121        let head_k = d_state; let head_v = d_state;
2122        let key_dim = head_k * num_k;                // 2048
2123        let value_dim = head_v * num_v;              // 4096
2124        let conv_dim = key_dim * 2 + value_dim;      // 8192
2125        let eps = cfg.rms_eps;
2126        let scale = 1.0 / (d_state as f32).sqrt();
2127
2128        // projections
2129        // grouped: one f16 activation convert feeds all four projections (matmul_group)
2130        let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
2131        let alpha = g4.pop().unwrap();                   // [T, num_v]
2132        let beta_raw = g4.pop().unwrap();                // [T, num_v]
2133        let z = g4.pop().unwrap();                       // [T, value_dim]
2134        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
2135
2136        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
2137        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
2138        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
2139        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
2140        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
2141        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
2142        let _ = (head_k, head_v);
2143        let mut q_g = e.uninit(d_state * num_v * t)?;
2144        let mut k_g = e.uninit(d_state * num_v * t)?;
2145        let mut v_g = e.uninit(d_state * num_v * t)?;
2146        e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
2147                         conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
2148        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
2149        let mut q_l2 = e.uninit(d_state * num_v * t)?;
2150        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
2151        let mut k_l2 = e.uninit(d_state * num_v * t)?;
2152        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
2153        let v_gd = v_g;
2154
2155        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
2156        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
2157        let mut beta = e.uninit(t * num_v)?;
2158        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
2159        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
2160        let mut g_log = e.uninit(t * num_v)?;
2161        e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
2162
2163        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
2164        let state_in = e.zeros(d_state * d_state * num_v)?;  // zero state (prefill)
2165        let mut state_out = e.zeros(d_state * d_state * num_v)?;
2166        let mut o = e.uninit(d_state * num_v * t)?;
2167        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)?;
2168
2169        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
2170        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
2171        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
2172        // o rows are (t*num_v+vh) too. Good.
2173        let mut gn = e.uninit(d_state * num_v * t)?;
2174        e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
2175
2176        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
2177        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
2178        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
2179        let out = e.matmul(&la.ssm_out, &gn, t)?;
2180        Ok(out)
2181    }
2182}
2183
2184impl HybridModel {
2185    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
2186    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
2187    ///
2188    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
2189    /// different 860160-byte block than the same expert of layer 7).
2190    ///
2191    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
2192    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
2193    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
2194    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
2195    pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
2196               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2197        Self::moe_ffn(e, m, z, t, &self.cfg, il, self.max_moe_block())
2198    }
2199
2200    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
2201    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
2202    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
2203    pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
2204                          zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
2205               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2206        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block())
2207    }
2208
2209    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
2210    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
2211    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
2212    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
2213    ///
2214    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
2215    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
2216    pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
2217                          cfg: &ModelConfig, il: u16, max_block: usize)
2218               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2219        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block)
2220    }
2221
2222    #[allow(clippy::too_many_arguments)]
2223    pub(crate) fn moe_ffn_inner(
2224        e: &Engine,
2225        m: &MoeWeights,
2226        z: &CudaSlice<f32>,
2227        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
2228        t: usize,
2229        cfg: &ModelConfig,
2230        il: u16,
2231        max_block: usize,
2232    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2233        let worker_io = crate::spill_pread::worker_enabled();
2234        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
2235        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
2236            e.with_moe_cache(max_block, |cache, _| {
2237                cache.begin_forward_epoch(il, t);
2238                if worker_io {
2239                    cache.begin_worker_scope();
2240                }
2241                Ok(())
2242            })?;
2243        }
2244        // A2: Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED=1 routes here.
2245        if t > 1 && std::env::var("MEMRA_MOE_GROUPED").is_ok() {
2246            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
2247            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
2248            // KNOWN t>1 MISMATCH maxdiff ~3.4e-4 (deterministic, 5x bit-identical 2026-07-05): the
2249            // sequential arm routes resident experts through the dev_q8 dp4a path (q8_1-quantized z
2250            // and act rows) while grouped stays f32-dequant qmatvec — a quantize-path difference,
2251            // not a bug (per-stage: act q8-vs-f32 ~4-9e-3 abs on |act|<=3, down-only ~1-3e-4; the
2252            // q8_1 activation-quantize error class). MEMRA_MOE_Q8=0 restores BYTE-IDENTICAL.
2253            if std::env::var("MEMRA_MOE_GATE").is_ok() {
2254                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
2255                let g_host = e.dtoh(&grouped_out)?;
2256                let s_host = e.dtoh(&seq_out)?;
2257                let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
2258                let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
2259                if g_bytes == s_bytes {
2260                    if il == 0 { println!("moe-gate il={il} t={t} BYTE-IDENTICAL (first layer only printed)"); }
2261                } else {
2262                    let diffs = g_host.iter().zip(s_host.iter()).enumerate()
2263                        .filter(|(_, (a, b))| a != b).count();
2264                    let maxdiff = g_host.iter().zip(s_host.iter())
2265                        .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
2266                    panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
2267                }
2268            }
2269            return Ok(grouped_out);
2270        }
2271        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
2272    }
2273
2274    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
2275    pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
2276                          cfg: &ModelConfig, il: u16, max_block: usize)
2277               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2278        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
2279    }
2280
2281    /// Append the host-visible router selection for one layer/forward when calibration tracing is
2282    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
2283    /// trace is independent of the dispatch optimization selected for the forward.
2284    fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
2285                        -> Result<(), Box<dyn std::error::Error>> {
2286        use std::io::Write as _;
2287        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
2288            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2289            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
2290            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
2291        }
2292        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
2293            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2294            let pairs: Vec<String> = sel_all.iter().zip(weights)
2295                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
2296                .collect();
2297            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
2298        }
2299        Ok(())
2300    }
2301
2302    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
2303    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
2304    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
2305    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
2306    fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
2307                       -> Result<(), Box<dyn std::error::Error>> {
2308        use std::io::Write as _;
2309        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
2310        let host = e.dtoh(z)?;
2311        if host.len() != t * n_embd {
2312            return Err(format!(
2313                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
2314                host.len(), t, n_embd
2315            ).into());
2316        }
2317        let bytes = unsafe {
2318            std::slice::from_raw_parts(
2319                host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
2320            )
2321        };
2322        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
2323        let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
2324        if state.is_none() {
2325            let dir = std::path::PathBuf::from(&dir);
2326            std::fs::create_dir_all(&dir)?;
2327            let index = std::fs::OpenOptions::new().create(true).append(true)
2328                .open(dir.join("index.jsonl"))?;
2329            *state = Some(MoeInputTraceWriter {
2330                dir,
2331                index,
2332                payloads: std::collections::HashMap::new(),
2333            });
2334        }
2335        let writer = state.as_mut().unwrap();
2336        if writer.dir != std::path::Path::new(&dir) {
2337            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
2338        }
2339        let file_name = format!("layer-{il:03}.f32");
2340        if !writer.payloads.contains_key(&il) {
2341            let payload = std::fs::OpenOptions::new().create(true).append(true)
2342                .open(writer.dir.join(&file_name))?;
2343            let offset = payload.metadata()?.len();
2344            writer.payloads.insert(il, (payload, offset));
2345        }
2346        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
2347        let row_offset = *offset;
2348        payload.write_all(bytes)?;
2349        *offset += bytes.len() as u64;
2350        writeln!(
2351            writer.index,
2352            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
2353             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
2354             \"payload_bytes\":{}}}",
2355            bytes.len()
2356        )?;
2357        Ok(())
2358    }
2359
2360    #[allow(clippy::too_many_arguments)]
2361    pub(crate) fn moe_ffn_sequential_zq8(
2362        e: &Engine,
2363        m: &MoeWeights,
2364        z: &CudaSlice<f32>,
2365        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
2366        t: usize,
2367        cfg: &ModelConfig,
2368        il: u16,
2369        max_block: usize,
2370    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2371        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
2372        let moe = cfg.moe.as_ref().unwrap();
2373        let n_embd = cfg.n_embd as usize;          // 2048 (gate/up in_f, down out_f)
2374        let n_expert = moe.expert_count as usize;  // 256
2375        let n_used = moe.expert_used_count as usize; // 8
2376        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
2377
2378        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
2379        debug_assert_eq!(m.gate_exps.in_f, n_embd);
2380        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
2381        debug_assert_eq!(m.down_exps.in_f, n_ff_exp);  // down is TRANSPOSED: in=512
2382        debug_assert_eq!(m.down_exps.out_f, n_embd);   //                     out=2048
2383        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
2384
2385        // step35 PER-LAYER SwiGLU clamp (None on every other arch and on every unclamped layer).
2386        // Routed experts and the shared expert read SEPARATE arrays — never share one value.
2387        let lim_exp = cfg.clamp_exp_at(il as u32);
2388        let lim_shexp = cfg.clamp_shexp_at(il as u32);
2389        let use_cache = Engine::moe_cache_enabled();
2390        let uniform_experts = m.has_uniform_expert_layout();
2391        let moe_q8 = uniform_experts && moe_q8_enabled()
2392            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2393            && q8_expert_supported(m.down_exps.qtype);
2394        // Experimental secondary backend: complete experts already resident in the SLRU stay on
2395        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
2396        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
2397        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
2398        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
2399        // commands and CI have no llama.cpp or OpenMP dependency.
2400        let cpu_expert_requested = crate::cpu_experts::configured();
2401        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
2402            return Err(std::io::Error::other(
2403                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
2404            )
2405            .into());
2406        }
2407        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
2408        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
2409        // Those backends are each deterministic but are different numeric configurations, so a
2410        // later prefill eviction can change greedy output. Freeze after the first real prefill;
2411        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
2412        // staging below and cannot change backend assignment.
2413        let freeze_cpu_residency = cpu_expert_requested
2414            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
2415        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
2416            .ok()
2417            .and_then(|value| value.parse::<usize>().ok())
2418            .is_some_and(|tokens| tokens > 0);
2419        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
2420            e.freeze_moe_cache();
2421        }
2422        let cache_frozen = use_cache && e.moe_cache_frozen();
2423        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
2424
2425        // 1. ROUTER: logits = ffn_gate_inp @ z  -> [T, 256]. gate_inp is F32 -> cuBLASLt, whose
2426        // reductions are n-DEPENDENT (lt_ndep probe: m=1 vs m=2 col0 differs every bit). At
2427        // small t (spec verify, 2..15) that shifts router logits vs the T=1 decode chain ->
2428        // top-k WEIGHTS (and at tie margins the SELECTION) differ -> verify != decode. Route
2429        // small-t through per-column m=1 calls (decode-exact contract); real prefill keeps the
2430        // batched GEMM.
2431        let logits = if t < PRIME_MIN_T {
2432            // t == 1 included since 2026-07-10 (was cuBLAS gemvx, 3.1% + adjacent of the depth
2433            // decode map): decode and verify now route through the SAME kernel — the
2434            // verify==decode router parity holds by construction instead of by FP-order luck.
2435            if crate::router_kernel_on() {
2436                // MEMRA_ROUTER_KERNEL=1: in-house router GEMV (battery-gated numeric config —
2437                // top-k discontinuity means FP-order changes can flip routing; oracle arbitrates).
2438                e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2439                              m.gate_exps.n_expert, t)?
2440            } else {
2441                e.matmul_decode_exact(&m.gate_inp, z, t)?
2442            }
2443        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
2444            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): the cuBLASLt router GEMM
2445            // is m-DEPENDENT — probed on the Ornith-35B router weight, rows [0,19) of an m=65
2446            // call differ from the m=19 call by 3.9e-3 while the same probe on the lm_head /
2447            // wq MMQ+f16 weights is BIT-IDENTICAL across m (research/concat-prime-exact-20260802,
2448            // gemm-razor-router-o35b.log vs gemm-razor-o35b.log). Because the router feeds a
2449            // top-k DISCONTINUITY, that perturbation reorders ties and at ~16% of (layer,token)
2450            // pairs changes the selected expert SET — so a request's own prefill routing depended
2451            // on how many OTHER requests' tokens shared its concat batch (cross-request prime
2452            // batching, worker.rs task #13). The in-house router GEMV computes one row per
2453            // (expert, token) block with a fixed per-row reduction order and is m-INVARIANT
2454            // (same probe: BIT-IDENTICAL, gemm-razor-router-gemv-o35b.log), so routing prefill
2455            // through it makes a session's routing a function of its OWN tokens alone — the
2456            // serving isolation contract at the prime level. MEMRA_ROUTER_PREFILL_EXACT=0 reverts
2457            // to the batched cuBLASLt GEMM (numeric-config rollback seam).
2458            e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2459                          m.gate_exps.n_expert, t)?
2460        } else {
2461            e.matmul(&m.gate_inp, z, t)?
2462        };
2463
2464        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
2465        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
2466        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
2467        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
2468        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
2469        // per-token host stall that dominated the 35B decode wall after stages 1+2.
2470        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
2471        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
2472        // only difference is where sel/w/pointers are READ from (device instead of params).
2473        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
2474        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
2475        // Any non-resident layer falls through to host routing + the gdec/sequential path.
2476        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
2477        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
2478        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
2479        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
2480        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
2481        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
2482        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
2483        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
2484        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
2485        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
2486        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
2487        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
2488        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
2489        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
2490        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
2491        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
2492        // now rides the dev loop below (same kernels per token as decode); pairs serves real
2493        // prefill (t >= 16, where spec never verifies).
2494        // sigmoid-router archs (M3, Hy3) must NOT enter the pairs/dev arms: those route via the
2495        // fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the M3
2496        // gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Host sigmoid routing below is correct.
2497        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
2498        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
2499        // ride the macro-aware sequential/staged paths below or every expert output is off by
2500        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
2501        let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2502            && m.down_exps.macros.is_none();
2503        // A clamped-SwiGLU layer (step35 43/44) also cannot ride pairs: moe_pairs_silu_mul's
2504        // epilogue is plain silu(gate)*up with no clamp form, and moe_ffn_pairs takes no `il`
2505        // so it cannot even see the per-layer limit.
2506        if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
2507            && !cfg.swiglu_clamped_at(il as u32)
2508            && no_exp_macros
2509            && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
2510            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2511            && q8_expert_supported(m.down_exps.qtype)
2512            && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
2513            && std::env::var("MEMRA_MOE_STATS").is_err() {
2514            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
2515        }
2516
2517        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
2518        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
2519        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
2520        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk) — sigmoid
2521        // routing (M3, Hy3: +expert bias) has no device kernel yet, so those arches must NOT
2522        // enter the dev arms: with MOE_CACHE=1 M3 silently routed softmax = wrong experts
2523        // (gate MISMATCH 74602 vs 92, caught 2026-07-07). Host sigmoid path below is correct.
2524        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
2525        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
2526        // step35 (2026-08-06) is the third sigmoid-router arch and the deny above was written as
2527        // two arch names, not as the mechanism — so `dev_ok` let it into moe_ffn_dev's
2528        // moe_router_topk (softmax, no exp_probs_b bias, no expert_weights_scale) = silently
2529        // wrong experts, the same failure the pairs gate at :2254 already blocks by predicate.
2530        // Keyed off sigmoid_router() so arch #4 is denied by construction.
2531        // The clamped-SwiGLU layers must also fall through: every moe_ffn_dev kernel's fused
2532        // epilogue is plain silu(gate)*up (moe_gate_up_silu8_dev_*), which has no clamp form.
2533        let dev_ok = uniform_experts && cfg.sigmoid_router().is_none()
2534            && cfg.m3.is_none() && cfg.hy3.is_none()
2535            && !cfg.swiglu_clamped_at(il as u32);
2536        // Observation modes must route through the host-visible selection below. Otherwise a fully
2537        // resident layer returns through device dispatch before its trace/stats row is recorded,
2538        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
2539        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
2540            || std::env::var("MEMRA_MOE_TRACE").is_ok()
2541            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
2542            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
2543        if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
2544            && !observe_routes {
2545            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2546        }
2547        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
2548            && !observe_routes {
2549            let row_ok = e.with_moe_cache(max_block, |c, eng| {
2550                if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
2551                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
2552            })?;
2553            if row_ok {
2554                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2555            }
2556        }
2557
2558        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
2559        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
2560            if cpu_hybrid {
2561                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
2562                    e,
2563                    &logits,
2564                    z,
2565                    t,
2566                    n_expert,
2567                    n_used,
2568                    m.exp_probs_b.as_deref(),
2569                    sig,
2570                    m.active_experts.as_deref(),
2571                )?;
2572                (sel, w, Some(input))
2573            } else {
2574                let (sel, w) = Self::moe_route_cfg(
2575                    e,
2576                    &logits,
2577                    t,
2578                    n_expert,
2579                    n_used,
2580                    m.exp_probs_b.as_deref(),
2581                    Some(sig),
2582                    m.active_experts.as_deref(),
2583                )?;
2584                (sel, w, None)
2585            }
2586        } else {
2587            let (sel, w) = Self::moe_route_cfg(
2588                e,
2589                &logits,
2590                t,
2591                n_expert,
2592                n_used,
2593                None,
2594                None,
2595                m.active_experts.as_deref(),
2596            )?;
2597            (sel, w, None)
2598        };
2599
2600        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
2601        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
2602        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
2603        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
2604        Self::trace_moe_input(e, il, t, n_embd, z)?;
2605
2606        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
2607        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
2608        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
2609        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
2610        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
2611        // wait for each pending block, so later copies can overlap the earlier expert kernels while
2612        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
2613        // T=1; batched forwards can have token-local consumers still in flight between selections.
2614        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
2615        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
2616        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
2617        let worker_disk_prefetch =
2618            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
2619        let promote_worker_h2d =
2620            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
2621        if promote_worker_h2d {
2622            let mut selected_blocks = Vec::with_capacity(n_used * 3);
2623            for &ex in sel_all.iter().take(n_used) {
2624                let ex = ex as u16;
2625                selected_blocks.extend([
2626                    BlockId::new(il, PROJ_GATE, ex),
2627                    BlockId::new(il, PROJ_UP, ex),
2628                    BlockId::new(il, PROJ_DOWN, ex),
2629                ]);
2630            }
2631            for &ex in sel_all.iter().take(n_used) {
2632                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
2633            }
2634            e.with_moe_cache(max_block, |cache, eng| {
2635                cache.promote_worker_reads_at_safe_boundary(
2636                    &selected_blocks,
2637                    &selected_blocks,
2638                    eng,
2639                )?;
2640                Ok(())
2641            })?;
2642        }
2643
2644        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
2645        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
2646        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
2647            let mut cnt = vec![0u32; n_expert];
2648            for &s in sel_all.iter() { cnt[s as usize] += 1; }
2649            let total = sel_all.len() as f64;
2650            let mut h = 0.0f64;
2651            let mut active = 0usize;
2652            for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
2653            let maxc = cnt.iter().copied().max().unwrap_or(0);
2654            println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
2655                     il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
2656        }
2657
2658        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
2659        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
2660        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
2661        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
2662        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
2663        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
2664        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
2665        // zeroed-then-accumulated exactly as before (fallback).
2666        // `!swiglu_clamped_at`: the grouped-decode kernels' fused epilogue is plain
2667        // silu(gate)*up (see the m3 note at the call sites below) — a clamped layer must fall
2668        // through to the sequential loop's ffn_act_lim. Hoisted into the fire predicate so the
2669        // uninit/memset invariant at :2418 stays consistent with what actually dispatches.
2670        let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled()
2671            && !cfg.swiglu_clamped_at(il as u32);
2672        // SLAB-LOCAL RESIDENT ARM (lane/pp-leverb 2026-08-08): the fits-VRAM resident slabs
2673        // (`dev_exps`, sized per PP device by cx-503b) were consumed ONLY by the pairs/dev
2674        // arms — every one of which DENIES sigmoid-router archs (step35/M3/Hy3). So on those
2675        // archs the slabs were uploaded but never read, and every expert went through the
2676        // SLRU (37 GB H2D staging per pp4096 prime on the Step SKU; and post-cx-503b the
2677        // dead slabs additionally STARVE the SLRU, which sizes itself on free-VRAM-after-
2678        // residents). This arm reads the slabs through the SAME kernels the SLRU arm runs —
2679        // gdec's fused pair when eligible, the per-expert qmatvec twins otherwise — with the
2680        // slab base + ex*stride as the pointer provenance (the exact bit-identity class
2681        // `moe_ffn_dev` documents for its resident-vs-SLRU arms; bytes are the same HostExps
2682        // bytes either way). LOCALITY GATE: `d.dev == e.ctx().ordinal()` — a slab on another
2683        // device must NOT be dereferenced (m=1 peer reads are the measured 34-150x class,
2684        // strictly worse than staging); under PP-2 without the prime walker this admits
2685        // stage-0 layers on dev0 and leaves stage-1 layers on the SLRU, which is the honest
2686        // interim shape. MEMRA_MOE_SLAB=0 = rollback/A-B seam (read per call).
2687        let slab_local = m.dev_exps.as_ref()
2688            .filter(|d| !d.gu_il && moe_slab_enabled() && d.dev == e.ctx().ordinal());
2689        let slab_bases = slab_local.map(|d| {
2690            use cudarc::driver::DevicePtr;
2691            let s = e.stream();
2692            let (pg, _g0) = d.gate.device_ptr(&s);
2693            let (pu, _g1) = d.up.device_ptr(&s);
2694            let (pd, _g2) = d.down.device_ptr(&s);
2695            (pg as u64, pu as u64, pd as u64)
2696        });
2697        // Fused-pair eligibility mirrors the gdec call sites exactly (plain-SiLU epilogue,
2698        // no macros, m3 clamp excluded, q8-supported qtypes, <=8 experts) — INCLUDING
2699        // `gdec_enabled()`: the pair IS the gdec kernel pair, so MEMRA_MOE_GDEC=0 must
2700        // disable it here too. That seam is also the exactness localizer: with GDEC=0 both
2701        // provenances run the SAME per-expert qmatvec kernels (slab base+stride vs SLRU
2702        // slot) and must be BIT-IDENTICAL — the true provenance-only pair — while GDEC=1
2703        // compares the fused-pair class against the SLRU's hit/miss MIX (gdec for
2704        // all-resident tokens, staged loop for misses), which is a dispatch-class
2705        // comparison, not a provenance one.
2706        let slab_fused_may_fire = slab_bases.is_some() && n_used <= 8 && gdec_enabled()
2707            && !cfg.swiglu_clamped_at(il as u32) && cfg.m3.is_none()
2708            && no_exp_macros && moe_q8;
2709        // moe_out memset elision: BOTH full-row-overwrite arms (gdec + slab fused) allocate
2710        // uninit; a token that falls through to any accumulating loop zeroes its own row.
2711        let mut moe_out = if gdec_may_fire || slab_fused_may_fire {
2712            e.uninit(t * n_embd)?
2713        } else {
2714            e.zeros(t * n_embd)?
2715        };
2716        // The router readback above already established a host boundary. Copy each small-t hidden
2717        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
2718        let cpu_input = if cpu_hybrid {
2719            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
2720        } else {
2721            None
2722        };
2723
2724        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
2725        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
2726        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
2727        // measured ~123 memsets/token of the decode wall).
2728        let g_len = m.gate_exps.max_expert_bytes();  // 860160 for the uniform 35B gate
2729        let u_len = m.up_exps.max_expert_bytes();    // 860160 for the uniform 35B up
2730        let d_len = m.down_exps.max_expert_bytes();  // 1114112 for the uniform 35B down
2731        let mut scratch_g: Option<CudaSlice<u8>> = None;
2732        let mut scratch_u: Option<CudaSlice<u8>> = None;
2733        let mut scratch_d: Option<CudaSlice<u8>> = None;
2734        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
2735        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
2736
2737        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
2738        // the copy stream before launching the current expert's compute. Pending slots stay invisible
2739        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
2740        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
2741        let page_window = moe_page_prefetch_window();
2742
2743        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
2744        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
2745        for tok in 0..t {
2746            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
2747            let w = &w_all[tok * n_used..(tok + 1) * n_used];
2748            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);  // CudaView<f32>
2749            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
2750
2751            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
2752            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
2753            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
2754            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
2755            // memcpy, zero admission, so no slot can move under the collected pointers) — any
2756            // miss falls through to the sequential loop below, which admits as before. In steady
2757            // state on a fully-resident rig every token-layer takes the grouped path.
2758            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
2759            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
2760            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
2761            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
2762            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
2763            // per-expert macro-scales the fused kernels don't fold — those fall through too.
2764            let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2765                && m.down_exps.macros.is_none();
2766            // SLAB-LOCAL FUSED PAIR (lane/pp-leverb 2026-08-08): the gdec launch pair
2767            // (moe_gate_up_silu8_q8 + moe_down8_fma_q8 — the EXACT kernels, same FP chains)
2768            // with pointers computed from the resident slab base + ex*stride instead of
2769            // collected SLRU slot addresses. No cache lock, no residency predicate — the
2770            // slab holds every expert by construction, so this arm never falls through
2771            // (gdec's P(all 24 resident) ≈ 0.37 coin-flip and the miss path's ~49-launch
2772            // staging both die). Bit-identity class: pointer provenance only, the same
2773            // slab-vs-SLRU equivalence moe_ffn_dev documents. Ordered BEFORE gdec: when a
2774            // slab exists it is strictly better (no lock, no miss).
2775            if slab_fused_may_fire {
2776                let (pg, pu, pd) = slab_bases.unwrap();
2777                let mut gp = [0u64; 8];
2778                let mut up = [0u64; 8];
2779                let mut dp = [0u64; 8];
2780                for (j, &ex) in sel.iter().enumerate() {
2781                    let ex = ex as usize;
2782                    gp[j] = pg + (ex * m.gate_exps.expert_stride) as u64;
2783                    up[j] = pu + (ex * m.up_exps.expert_stride) as u64;
2784                    dp[j] = pd + (ex * m.down_exps.expert_stride) as u64;
2785                }
2786                let mut wv = [0f32; 8];
2787                wv[..n_used].copy_from_slice(w);
2788                if tok_q8.is_none() {
2789                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2790                }
2791                let (zq, zd) = tok_q8.as_ref().unwrap();
2792                let act = e.moe_gate_up_silu8_q8(crate::WPtr8(gp), crate::WPtr8(up), zq, zd,
2793                                                 n_embd, n_ff_exp, n_used,
2794                                                 m.gate_exps.qtype, m.up_exps.qtype,
2795                                                 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
2796                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
2797                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2798                e.moe_down8_fma_q8(crate::WPtr8(dp), crate::F32x8(wv), &aq2, &ad2, &mut dst,
2799                                   n_ff_exp, n_embd, n_used,
2800                                   m.down_exps.qtype, m.down_exps.row_bytes)?;
2801                continue;
2802            }
2803            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
2804                if tok_q8.is_none() {
2805                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2806                }
2807                let (zq, zd) = tok_q8.as_ref().unwrap();
2808                if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
2809                                           &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2810                    continue;
2811                }
2812            } else if gdec_may_fire && cfg.m3.is_none() && no_macros
2813                && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
2814                                        &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2815                continue;
2816            }
2817
2818            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec or the
2819            // slab pair could fire. This token fell through to a sequential axpy loop, which
2820            // ACCUMULATES — zero its row first (row-sized memset; other rows are owned by the
2821            // full-overwrite arms). slab_fused_may_fire never actually reaches here (its arm
2822            // has no fallible predicate), included for the allocation invariant's symmetry.
2823            if gdec_may_fire || slab_fused_may_fire {
2824                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2825                e.memset_zeros_view(&mut row)?;
2826            }
2827
2828            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
2829            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
2830            // stall this path exists to remove, while mixing projections would require another
2831            // activation round-trip. Weight addresses remain valid until this worker is joined at
2832            // the bottom of the token scope.
2833            let mut cpu_mask = vec![false; sel.len()];
2834            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
2835                let gpu_resident = if use_cache {
2836                    e.with_moe_cache(max_block, |cache, _| {
2837                        Ok(sel
2838                            .iter()
2839                            .map(|&expert| {
2840                                let expert = expert as u16;
2841                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
2842                                    .into_iter()
2843                                    .filter(|&projection| {
2844                                        cache
2845                                            .resident(BlockId::new(il, projection, expert))
2846                                            .is_some()
2847                                    })
2848                                    .count()
2849                            })
2850                            .collect::<Vec<_>>())
2851                    })?
2852                } else {
2853                    vec![0; sel.len()]
2854                };
2855                let mut cpu_selected = Vec::new();
2856                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
2857                    if gpu_resident[index] != 3 {
2858                        cpu_mask[index] = true;
2859                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
2860                        let expert = expert as usize;
2861                        cpu_selected.push((expert, route_weight));
2862                    }
2863                }
2864                if crate::cpu_experts::predictor_enabled() {
2865                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
2866                    // from this layer's MoE input and prefetches predicted-and-missing
2867                    // experts into the companion RAM cache. Never blocks this thread.
2868                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2869                    crate::cpu_experts::predictor_submit(il, row);
2870                }
2871                if cpu_selected.is_empty() {
2872                    None
2873                } else {
2874                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2875                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
2876                        .map_err(std::io::Error::other)?;
2877                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
2878                }
2879            } else {
2880                None
2881            };
2882
2883            let worker_window = worker_disk_prefetch
2884                .then(worker_prefetch_window)
2885                .unwrap_or(0);
2886            for (j, &ex) in sel.iter().enumerate() {
2887                if cpu_mask[j] {
2888                    continue;
2889                }
2890                let ex = ex as usize;
2891                // PER-EXPERT SLAB READ (lane/pp-leverb 2026-08-08): the layers the fused
2892                // pair above excludes — step35's CLAMPED layers 43/44 (ffn_act_lim has no
2893                // fused form) and macro-carrying artifacts — still have their bytes in the
2894                // local resident slab. Same kernels as the SLRU arms (`qmatvec_expert_q8` /
2895                // `qmatvec_view`), same ffn_act_lim/macro folds; provenance = slab base +
2896                // ex*stride. No dispatch lock, no admission, no prefetch — nothing to miss.
2897                if let Some(d) = slab_local {
2898                    let gl = m.gate_exps.expert_layout(ex);
2899                    let ul = m.up_exps.expert_layout(ex);
2900                    let dl = m.down_exps.expert_layout(ex);
2901                    let (g0, u0, d0) = (ex * m.gate_exps.expert_stride,
2902                                        ex * m.up_exps.expert_stride,
2903                                        ex * m.down_exps.expert_stride);
2904                    let (gate, up) = if moe_q8 {
2905                        if tok_q8.is_none() {
2906                            tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2907                        }
2908                        let (zq, zd) = tok_q8.as_ref().unwrap();
2909                        (e.qmatvec_expert_q8(&d.gate, g0..g0 + gl.len, zq, zd, 1,
2910                                             m.gate_exps.in_f, m.gate_exps.out_f,
2911                                             gl.qtype, gl.row_bytes)?,
2912                         e.qmatvec_expert_q8(&d.up, u0..u0 + ul.len, zq, zd, 1,
2913                                             m.up_exps.in_f, m.up_exps.out_f,
2914                                             ul.qtype, ul.row_bytes)?)
2915                    } else {
2916                        (e.qmatvec_view(&d.gate, g0..g0 + gl.len, &zt, 1,
2917                                        m.gate_exps.in_f, m.gate_exps.out_f,
2918                                        gl.qtype, gl.row_bytes)?,
2919                         e.qmatvec_view(&d.up, u0..u0 + ul.len, &zt, 1,
2920                                        m.up_exps.in_f, m.up_exps.out_f,
2921                                        ul.qtype, ul.row_bytes)?)
2922                    };
2923                    let mut act = e.uninit(n_ff_exp)?;
2924                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
2925                                      m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
2926                    let y = if moe_q8 {
2927                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
2928                        e.qmatvec_expert_q8(&d.down, d0..d0 + dl.len, &aq2, &ad2, 1,
2929                                            m.down_exps.in_f, m.down_exps.out_f,
2930                                            dl.qtype, dl.row_bytes)?
2931                    } else {
2932                        let actv = act.slice(0..n_ff_exp);
2933                        e.qmatvec_view(&d.down, d0..d0 + dl.len, &actv, 1,
2934                                       m.down_exps.in_f, m.down_exps.out_f,
2935                                       dl.qtype, dl.row_bytes)?
2936                    };
2937                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2938                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2939                    continue;
2940                }
2941                for next in page_prefetch_positions(j, sel.len(), page_window) {
2942                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
2943                }
2944                let keep = [
2945                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
2946                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
2947                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
2948                ];
2949                if worker_disk_prefetch && worker_window > 0 {
2950                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
2951                        Self::moe_prefetch_disk_expert(
2952                            e,
2953                            il,
2954                            sel[next] as usize,
2955                            m,
2956                            max_block,
2957                            &keep,
2958                        )?;
2959                    }
2960                } else if cache_dispatch
2961                    && !cpu_hybrid
2962                    && moe_prefetch_enabled()
2963                    && j + 1 < sel.len()
2964                {
2965                    let next = sel[j + 1] as usize;
2966                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
2967                }
2968                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
2969                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
2970                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
2971                    // layouts stay on the metadata-aware f32 path.
2972                    if (gate_q8 || up_q8) && tok_q8.is_none() {
2973                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2974                    }
2975                    let gate = if gate_q8 {
2976                        let (zq, zd) = tok_q8.as_ref().unwrap();
2977                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
2978                    } else {
2979                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
2980                    };
2981                    let up = if up_q8 {
2982                        let (zq, zd) = tok_q8.as_ref().unwrap();
2983                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
2984                    } else {
2985                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
2986                    };
2987                    let mut act = e.uninit(n_ff_exp)?;
2988                    Self::ffn_act_lim(
2989                        e,
2990                        cfg,
2991                        &gate,
2992                        &up,
2993                        m.gate_exps.macro_scale(ex),
2994                        m.up_exps.macro_scale(ex),
2995                        lim_exp,
2996                        &mut act,
2997                        n_ff_exp,
2998                    )?;
2999                    let y = if down_q8 {
3000                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
3001                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
3002                    } else {
3003                        let actv = act.slice(0..n_ff_exp);
3004                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
3005                    };
3006                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3007                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
3008                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3009                } else if cache_dispatch {
3010                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
3011                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
3012                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
3013                    // only difference between HIT and MISS is whether the memcpy_htod ran.
3014                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
3015                    let up   = Self::moe_cached_gemm(e, il, PROJ_UP,   ex, m, max_block, &zt)?;
3016                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
3017                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3018                        m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3019                    let actv = act.slice(0..n_ff_exp);
3020                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
3021                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3022                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
3023                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3024                } else if cache_frozen {
3025                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
3026                    // first prime. Reuse every fixed resident projection directly and stage only a
3027                    // true miss through the ordinary scratch slot. This preserves the established
3028                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
3029                    let gate = Self::moe_frozen_gemm(
3030                        e,
3031                        il,
3032                        PROJ_GATE,
3033                        ex,
3034                        m,
3035                        max_block,
3036                        &zt,
3037                        &mut scratch_g,
3038                        g_len,
3039                    )?;
3040                    let up = Self::moe_frozen_gemm(
3041                        e,
3042                        il,
3043                        PROJ_UP,
3044                        ex,
3045                        m,
3046                        max_block,
3047                        &zt,
3048                        &mut scratch_u,
3049                        u_len,
3050                    )?;
3051                    let mut act = e.uninit(n_ff_exp)?;
3052                    Self::ffn_act_lim(
3053                        e,
3054                        cfg,
3055                        &gate,
3056                        &up,
3057                        m.gate_exps.macro_scale(ex),
3058                        m.up_exps.macro_scale(ex),
3059                        lim_exp,
3060                        &mut act,
3061                        n_ff_exp,
3062                    )?;
3063                    let actv = act.slice(0..n_ff_exp);
3064                    let y = Self::moe_frozen_gemm(
3065                        e,
3066                        il,
3067                        PROJ_DOWN,
3068                        ex,
3069                        m,
3070                        max_block,
3071                        &actv,
3072                        &mut scratch_d,
3073                        d_len,
3074                    )?;
3075                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3076                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3077                } else {
3078                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
3079                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
3080                    // fully overwrites the byte range the GEMM reads).
3081                    if scratch_g.is_none() {
3082                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
3083                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
3084                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
3085                    }
3086                    let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
3087                                        scratch_d.as_mut().unwrap());
3088                    let gl = m.gate_exps.expert_layout(ex);
3089                    let ul = m.up_exps.expert_layout(ex);
3090                    let dl = m.down_exps.expert_layout(ex);
3091                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
3092                    let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
3093                        m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
3094
3095                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
3096                    let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
3097                        m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
3098
3099                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
3100                    Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
3101                        m.up_exps.macro_scale(ex), lim_exp, &mut act, n_ff_exp)?;
3102
3103                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
3104                    let actv = act.slice(0..n_ff_exp);
3105                    let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
3106                        m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
3107
3108                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3109                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
3110                }
3111            }
3112            if let Some(worker) = cpu_worker {
3113                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
3114                let cpu_output = e.htod(&cpu_output)?;
3115                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3116                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
3117            }
3118            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
3119                for (j, &ex) in sel.iter().enumerate() {
3120                    if cpu_mask[j] {
3121                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
3122                    }
3123                }
3124            }
3125        }
3126
3127        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
3128        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
3129        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
3130        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
3131        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3132            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3133        {
3134            let n_ff_sh = gate_shexp.out_features();  // 512
3135            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
3136            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
3137            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
3138            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
3139            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
3140            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
3141            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
3142            let verify_t = t > 1 && t < PRIME_MIN_T;
3143            let (sg_gate, sg_up) = if t == 1 {
3144                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
3145                    Some(pair) => pair,
3146                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
3147                }
3148            } else if verify_t {
3149                (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
3150            } else {
3151                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)   // [T, 512] each
3152            };
3153            let mut sa = e.uninit(t * n_ff_sh)?;  // activation fully overwrites
3154            Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp, &mut sa, t * n_ff_sh)?;
3155            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
3156                     else { e.matmul(down_shexp, &sa, t)? };     // [T, n_embd]
3157
3158            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
3159            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
3160            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
3161            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
3162            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
3163            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
3164            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
3165            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
3166            // expert's contribution into every token's residual, so under cross-request
3167            // concat prefill a session's hidden state depended on its co-arrivals' token
3168            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
3169            let g = match &m.gate_inp_shexp {
3170                Some(gate_inp_shexp) => {
3171                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
3172                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3173                    } else {
3174                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3175                        let mut g = e.uninit(t)?;  // sigmoid fully overwrites
3176                        e.sigmoid(&gs, &mut g, t)?;
3177                        g
3178                    }
3179                }
3180                None => e.htod(&vec![1.0f32; t])?,
3181            };
3182            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
3183            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3184        }
3185
3186        Ok(moe_out)
3187    }
3188
3189    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
3190    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
3191    pub fn stage1_h2d_per_token(&self) -> u64 {
3192        use crate::hybrid::Ffn;
3193        let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
3194        let mut bytes = 0u64;
3195        for l in self.layers.iter() {
3196            if let Ffn::Moe(m) = &l.ffn {
3197                bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
3198                                   + m.down_exps.max_expert_bytes()) as u64;
3199            }
3200        }
3201        bytes
3202    }
3203
3204    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
3205    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
3206    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
3207    pub(crate) fn max_moe_block(&self) -> usize {
3208        use crate::hybrid::Ffn;
3209        let mut mx = 0usize;
3210        let mut scan = |ffn: &Ffn| {
3211            if let Ffn::Moe(m) = ffn {
3212                mx = mx.max(m.gate_exps.max_expert_bytes())
3213                       .max(m.up_exps.max_expert_bytes())
3214                       .max(m.down_exps.max_expert_bytes());
3215            }
3216        };
3217        for l in self.layers.iter() { scan(&l.ffn); }
3218        if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
3219        mx
3220    }
3221
3222    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
3223    /// but have no bytes and therefore consume no residency slot.
3224    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
3225        use crate::hybrid::Ffn;
3226        let mut sizes = Vec::new();
3227        let mut scan = |ffn: &Ffn| {
3228            let Ffn::Moe(m) = ffn else { return };
3229            for ex in 0..m.gate_exps.n_expert {
3230                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
3231                    continue;
3232                }
3233                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
3234                    let len = exps.expert_layout(ex).len;
3235                    if len > 0 {
3236                        sizes.push(len);
3237                    }
3238                }
3239            }
3240        };
3241        for layer in &self.layers {
3242            scan(&layer.ffn);
3243        }
3244        if let Some(mtp) = &self.mtp {
3245            scan(&mtp.ffn);
3246        }
3247        sizes
3248    }
3249
3250    /// Persist the frozen residency set so a later process can restage it directly and skip
3251    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
3252    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
3253    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
3254    /// post-freeze argmax gate still validates the serving assignment.
3255    pub fn save_cpu_expert_residency_profile(
3256        &self,
3257        e: &Engine,
3258        path: &std::path::Path,
3259    ) -> Result<(), Box<dyn std::error::Error>> {
3260        let Some(ids) = e.export_moe_residency() else {
3261            return Err("no MoE residency cache to persist".into());
3262        };
3263        let mut body = format!(
3264            "memra-freeze-profile v1 max_block={} blocks={}\n",
3265            self.max_moe_block(),
3266            ids.len()
3267        );
3268        for (layer, proj, ex) in &ids {
3269            body.push_str(&format!("{layer} {proj} {ex}\n"));
3270        }
3271        let tmp = path.with_extension("tmp");
3272        std::fs::write(&tmp, body)?;
3273        std::fs::rename(&tmp, path)?;
3274        println!(
3275            "[moe-cache] freeze profile saved: {} blocks -> {}",
3276            ids.len(),
3277            path.display()
3278        );
3279        Ok(())
3280    }
3281
3282    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
3283    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
3284    /// missing or its header does not match this model's slot geometry.
3285    pub fn restore_cpu_expert_residency_profile(
3286        &self,
3287        e: &Engine,
3288        path: &std::path::Path,
3289    ) -> Result<bool, Box<dyn std::error::Error>> {
3290        use crate::hybrid::Ffn;
3291        use crate::moe_cache::BlockId;
3292        let Ok(content) = std::fs::read_to_string(path) else {
3293            return Ok(false);
3294        };
3295        let mut lines = content.lines();
3296        let Some(header) = lines.next() else { return Ok(false) };
3297        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
3298        if !header.starts_with(&expected) {
3299            println!(
3300                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
3301                path.display()
3302            );
3303            return Ok(false);
3304        }
3305        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
3306            std::collections::HashMap::new();
3307        for line in lines {
3308            let mut fields = line.split_whitespace();
3309            let (Some(layer), Some(proj), Some(ex)) =
3310                (fields.next(), fields.next(), fields.next())
3311            else {
3312                continue;
3313            };
3314            let (Ok(layer), Ok(proj), Ok(ex)) =
3315                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
3316            else {
3317                continue;
3318            };
3319            by_layer
3320                .entry(layer)
3321                .or_default()
3322                .push(BlockId::new(layer, proj, ex));
3323        }
3324        let requested: usize = by_layer.values().map(Vec::len).sum();
3325        if requested == 0 {
3326            return Ok(false);
3327        }
3328        let max_block = self.max_moe_block();
3329        let mut restaged = 0usize;
3330        let mut stage_layer = |layer_index: u16,
3331                               ffn: &Ffn|
3332         -> Result<(), Box<dyn std::error::Error>> {
3333            let Ffn::Moe(m) = ffn else { return Ok(()) };
3334            let Some(ids) = by_layer.get(&layer_index) else {
3335                return Ok(());
3336            };
3337            e.with_moe_cache(max_block, |cache, eng| {
3338                for id in ids {
3339                    if cache.restage_block(*id, m, eng)? {
3340                        restaged += 1;
3341                    }
3342                }
3343                Ok(())
3344            })
3345        };
3346        for (index, layer) in self.layers.iter().enumerate() {
3347            stage_layer(index as u16, &layer.ffn)?;
3348        }
3349        if let Some(mtp) = self.mtp.as_ref() {
3350            stage_layer(u16::MAX, &mtp.ffn)?;
3351        }
3352        e.freeze_moe_cache();
3353        println!(
3354            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
3355            path.display()
3356        );
3357        Ok(true)
3358    }
3359
3360    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
3361    pub fn freeze_cpu_expert_residency(
3362        &self,
3363        e: &Engine,
3364    ) -> Result<(), Box<dyn std::error::Error>> {
3365        e.freeze_moe_cache();
3366        Ok(())
3367    }
3368
3369    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
3370    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
3371    /// the model's activation exactly.
3372    ///
3373    /// NO-`il` FORM: cannot apply step35's PER-LAYER SwiGLU clamp. Only call it from a site whose
3374    /// layer provably has no live limit (dense-FFN layers, MTP blocks) — `ffn_act_lim` is the
3375    /// form for anything that can land on a clamped layer.
3376    pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3377               act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
3378        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
3379    }
3380
3381    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
3382    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
3383    /// carries a `weight_scale_2`). Same no-`il` contract as `ffn_act`.
3384    #[allow(clippy::too_many_arguments)]
3385    pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3386               gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
3387               -> Result<(), Box<dyn std::error::Error>> {
3388        Self::ffn_act_lim(e, cfg, gate, up, gs, us, None, act, n)
3389    }
3390
3391    /// ffn_act_scaled + step35's PER-LAYER clamped SwiGLU. `limit`:
3392    ///   * `None`   -> the unclamped dispatch (every arch except step35's layers 43-44).
3393    ///   * `Some(l)`-> `min(silu(gate*gs), l) * clamp(up*us, +-l)` (llama-graph.cpp:2146/1751,
3394    ///                 non-DEEPSEEK4 branch). Callers source it from `cfg.clamp_exp_at(il)`
3395    ///                 (routed experts) or `cfg.clamp_shexp_at(il)` (shared expert) — the two
3396    ///                 arrays are SEPARATE and a layer can have one without the other.
3397    /// The `> 1e-6` eps gate lives in `clamp_exp_at`/`clamp_shexp_at`, so a `Some` here is
3398    /// already known live.
3399    #[allow(clippy::too_many_arguments)]
3400    pub(crate) fn ffn_act_lim(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
3401               gs: f32, us: f32, limit: Option<f32>, act: &mut CudaSlice<f32>, n: usize)
3402               -> Result<(), Box<dyn std::error::Error>> {
3403        if let Some(m3) = cfg.m3.as_ref() {
3404            debug_assert!(limit.is_none(), "m3 swigluoai and step35 clamp are different archs");
3405            return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
3406        }
3407        if let Some(l) = limit {
3408            return e.swiglu_clamped_mul_scaled(gate, up, gs, us, l, act, n);
3409        }
3410        if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
3411        e.silu_mul_scaled(gate, up, gs, us, act, n)
3412    }
3413
3414    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
3415    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
3416    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
3417    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
3418    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
3419    fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
3420                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3421        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None, None, None)
3422    }
3423
3424    /// DeepSeek-V3-class sigmoid routing (MiniMax-M3, Hy3), host oracle. Reference:
3425    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
3426    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
3427    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
3428    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
3429    /// `sig` = (scaling_factor, route_norm) from `cfg.sigmoid_router()`; softmax archs pass
3430    /// None -> the qwen35moe/OLMoE path below.
3431    fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
3432                     bias: Option<&[f32]>, sig: Option<(f32, bool)>, active: Option<&[bool]>)
3433                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3434        if let Some((sf, route_norm)) = sig {
3435            // sigmoid routing. Host path only for now (fused-router kernel is softmax-top-k).
3436            let lg = e.dtoh(logits)?;
3437            return Self::moe_route_sigmoid_host(
3438                &lg, t, n_expert, n_used, bias, sf, route_norm, active,
3439            );
3440        }
3441        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
3442        // rollback) via the single-sync pinned readback — softmax arch only; the M3 sigmoid arm
3443        // above returns before this (host path until a sigmoid fused-router kernel exists).
3444        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
3445            return e.moe_router_topk_host(logits, t, n_expert, n_used);
3446        }
3447        // Host oracle (the §D bit-identity reference).
3448        let lg = e.dtoh(logits)?;   // [T*n_expert] host
3449        let mut sel = vec![0u32; t * n_used];
3450        let mut w_out = vec![0f32; t * n_used];
3451        for tok in 0..t {
3452            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3453            // softmax over ALL n_expert (stable: subtract max)
3454            let maxl = row.iter().enumerate()
3455                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
3456                .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
3457            let mut probs = vec![0f32; n_expert];
3458            let mut den = 0f32;
3459            for i in 0..n_expert {
3460                if active.is_some_and(|mask| !mask[i]) { continue; }
3461                let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
3462            }
3463            for p in probs.iter_mut() { *p /= den; }
3464            // stable DESC sort: prob DESC, ascending-index tiebreak.
3465            let mut idx: Vec<usize> = (0..n_expert)
3466                .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
3467            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
3468            let sl = &idx[..n_used];
3469            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
3470            let mut ws: f32 = wv.iter().sum();
3471            ws = ws.max(6.103515625e-5_f32);  // F16 smallest normal, clamp BEFORE divide
3472            for x in wv.iter_mut() { *x /= ws; }
3473            for j in 0..n_used {
3474                sel[tok * n_used + j] = sl[j] as u32;
3475                w_out[tok * n_used + j] = wv[j];
3476            }
3477        }
3478        Ok((sel, w_out))
3479    }
3480
3481    #[allow(clippy::too_many_arguments)]
3482    fn moe_route_sigmoid_with_input(
3483        e: &Engine,
3484        logits: &CudaSlice<f32>,
3485        input: &CudaSlice<f32>,
3486        t: usize,
3487        n_expert: usize,
3488        n_used: usize,
3489        bias: Option<&[f32]>,
3490        (sf, route_norm): (f32, bool),
3491        active: Option<&[bool]>,
3492    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
3493        let (lg, input) = e.dtoh_pair(logits, input)?;
3494        let (sel, w) =
3495            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
3496        Ok((sel, w, input))
3497    }
3498
3499    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
3500    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
3501    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
3502    /// active mask, prebuilt projection descriptors) so no model reference escapes.
3503    pub fn start_moe_prefetch_predictor(
3504        &self,
3505        e: &Engine,
3506        cfg: &ModelConfig,
3507    ) -> Result<(), Box<dyn std::error::Error>> {
3508        use crate::hybrid::Ffn;
3509        let Some(sig) = cfg.sigmoid_router() else {
3510            return Err("prefetch predictor requires a sigmoid-router arch".into());
3511        };
3512        let resident: std::collections::HashSet<(u16, u8, u16)> = e
3513            .export_moe_residency()
3514            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
3515            .into_iter()
3516            .collect();
3517        let mut layers = Vec::new();
3518        for (index, layer) in self.layers.iter().enumerate() {
3519            let Ffn::Moe(m) = &layer.ffn else { continue };
3520            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
3521            let router = e.dtoh(data)?;
3522            let n_expert = m.gate_exps.n_expert;
3523            let n_embd = m.gate_exps.in_f;
3524            if router.len() != n_embd * n_expert {
3525                continue;
3526            }
3527            let build = |exps: &crate::model::HostExps| {
3528                (0..n_expert)
3529                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
3530                    .collect::<Vec<_>>()
3531            };
3532            layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
3533                router,
3534                bias: m.exp_probs_b.clone(),
3535                active: m.active_experts.clone(),
3536                n_embd,
3537                n_used: cfg
3538                    .moe
3539                    .as_ref()
3540                    .map(|moe| moe.expert_used_count as usize)
3541                    .ok_or("prefetch predictor requires MoE config")?,
3542                sig,
3543                weights_n_expert: n_expert,
3544                gate: build(&m.gate_exps),
3545                up: build(&m.up_exps),
3546                down: build(&m.down_exps),
3547            }));
3548        }
3549        crate::cpu_experts::start_prefetch_predictor(layers, resident)
3550            .map_err(|error| error.into())
3551    }
3552
3553    /// Crate-visible sigmoid-routing oracle for the prefetch predictor: identical selection
3554    /// math to the runtime router, applied to host-computed lookahead logits.
3555    #[allow(clippy::too_many_arguments)]
3556    pub(crate) fn moe_route_sigmoid_host_public(
3557        logits: &[f32],
3558        t: usize,
3559        n_expert: usize,
3560        n_used: usize,
3561        bias: Option<&[f32]>,
3562        sf: f32,
3563        route_norm: bool,
3564        active: Option<&[bool]>,
3565    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3566        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
3567    }
3568
3569    #[allow(clippy::too_many_arguments)]
3570    fn moe_route_sigmoid_host(
3571        lg: &[f32],
3572        t: usize,
3573        n_expert: usize,
3574        n_used: usize,
3575        bias: Option<&[f32]>,
3576        sf: f32,
3577        route_norm: bool,
3578        active: Option<&[bool]>,
3579    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3580        if lg.len() != t * n_expert {
3581            return Err(format!(
3582                "sigmoid router logits length mismatch: got {}, expected {}",
3583                lg.len(),
3584                t * n_expert,
3585            )
3586            .into());
3587        }
3588        let mut sel = vec![0u32; t * n_used];
3589        let mut w_out = vec![0f32; t * n_used];
3590        for tok in 0..t {
3591            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3592            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
3593            // selection score = sigmoid + bias; weight = plain sigmoid.
3594            let selsc: Vec<f32> = match bias {
3595                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
3596                None => scores.clone(),
3597            };
3598            let mut idx: Vec<usize> = (0..n_expert)
3599                .filter(|&i| active.is_none_or(|mask| mask[i]))
3600                .collect();
3601            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
3602            let sl = &idx[..n_used];
3603            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
3604            if route_norm {
3605                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
3606                for x in wv.iter_mut() {
3607                    *x = *x / ws * sf;
3608                }
3609            } else {
3610                for x in wv.iter_mut() {
3611                    *x *= sf;
3612                }
3613            }
3614            for j in 0..n_used {
3615                sel[tok * n_used + j] = sl[j] as u32;
3616                w_out[tok * n_used + j] = wv[j];
3617            }
3618        }
3619        Ok((sel, w_out))
3620    }
3621
3622    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
3623    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
3624    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
3625    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
3626    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
3627    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
3628    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
3629    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
3630    fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
3631                     t: usize, cfg: &ModelConfig)
3632                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3633        let moe = cfg.moe.as_ref().unwrap();
3634        let n_embd = cfg.n_embd as usize;
3635        let n_expert = moe.expert_count as usize;
3636        let n_used = moe.expert_used_count as usize;
3637        let n_ff_exp = moe.expert_ff_length as usize;
3638        // This arm has no `il` and every kernel in it fuses PLAIN silu(gate)*up, so a clamped
3639        // model must never reach it. The caller's gate at `moe_ffn_sequential_zq8` denies per
3640        // layer via `swiglu_clamped_at(il)`; assert the whole-model form here so a future caller
3641        // that forgets the gate fails loudly in debug instead of returning wrong logits.
3642        debug_assert!(!cfg.swiglu_clamped_anywhere(),
3643                      "moe_ffn_pairs has no per-layer clamp: fused epilogues are plain SiLU");
3644        let dev = m.dev_exps.as_ref().unwrap();
3645        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
3646        let (rbg_d, rbu_d) = if dev.gu_il {
3647            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3648        } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3649
3650        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
3651        let n_pairs = t * n_used;
3652        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
3653        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
3654        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
3655        let pair_ex:  Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
3656        let pair_w:   Vec<f32> = w_all.clone();
3657        let tok_off:  Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
3658        let tok_ids:  Vec<i32> = (0..n_pairs as i32).collect();
3659        let pt = e.htod_i32(&pair_tok)?;
3660        let px = e.htod_i32(&pair_ex)?;
3661        let pw = e.htod(&pair_w)?;
3662        let toff = e.htod_i32(&tok_off)?;
3663        let tids = e.htod_i32(&tok_ids)?;
3664
3665        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
3666        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
3667        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
3668        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
3669        for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
3670        let mut ex_ids: Vec<i32> = Vec::new();
3671        let mut ex_off: Vec<i32> = vec![0];
3672        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
3673        for (ex, list) in by_ex.iter().enumerate() {
3674            if list.is_empty() { continue; }
3675            ex_ids.push(ex as i32);
3676            ex_pairs.extend_from_slice(list);
3677            ex_off.push(ex_pairs.len() as i32);
3678        }
3679        let n_active = ex_ids.len();
3680        let exi = e.htod_i32(&ex_ids)?;
3681        let exo = e.htod_i32(&ex_off)?;
3682        let exp_d = e.htod_i32(&ex_pairs)?;
3683        let _ = &px;   // pair-major twin keeps it; em path uses CSR
3684
3685        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
3686        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
3687        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
3688        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
3689        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
3690        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
3691        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
3692        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
3693        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
3694        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
3695        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
3696        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
3697        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
3698        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
3699        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
3700        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
3701        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
3702        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
3703        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
3704        let mma_t = *MMA_T.get_or_init(|| {
3705            std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
3706        });
3707        let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
3708            && t >= mma_t
3709            && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
3710            && q8_expert_dec_supported(m.down_exps.qtype)
3711            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3712        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
3713        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
3714        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
3715        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
3716        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
3717        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
3718        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
3719        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
3720        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
3721        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
3722        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
3723        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
3724        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
3725        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
3726        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
3727        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
3728            && q8_expert_dec_supported(m.up_exps.qtype)
3729            && q8_expert_dec_supported(m.down_exps.qtype)
3730            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3731        let f16g_mode = crate::moe_f16g_mode();
3732        let f16g = f16g_mode != 0 && t >= mma_t
3733            && (f16g_mode != 3 || !mma_capable)
3734            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
3735            && f16g_proj_ok(m.up_exps.qtype, n_embd)
3736            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
3737        if use_mma || f16g {
3738            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
3739            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
3740            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
3741            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
3742            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
3743            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
3744            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
3745            let y_down = if f16g {
3746                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
3747                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
3748                // permute at the very end back to pair-id order for the scatter.
3749                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
3750                let csr_tok_d = e.htod_i32(&csr_tok)?;
3751                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
3752                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
3753                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3754                                              m.gate_exps.qtype, rbg_d)?;
3755                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
3756                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3757                                              m.up_exps.qtype, rbu_d)?;
3758                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
3759                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
3760                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
3761                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
3762                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
3763                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
3764            } else {
3765            // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
3766            let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
3767            let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3768                                        n_embd, n_ff_exp, n_active, n_pairs, t,
3769                                        m.gate_exps.qtype, rbg_d)?;
3770            let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3771                                      n_embd, n_ff_exp, n_active, n_pairs, t,
3772                                      m.up_exps.qtype, rbu_d)?;
3773            // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
3774            // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
3775            // registers and writes ONLY the quantized scratch — the two-pass chain
3776            // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
3777            // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
3778            let a_scr = if crate::moe_fuse_actq_on() {
3779                e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
3780            } else {
3781                let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3782                e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
3783            };
3784            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3785            let pself = e.htod_i32(&pair_self)?;
3786            e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
3787                             n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
3788                             m.down_exps.qtype, m.down_exps.row_bytes)?
3789            };
3790            let mut moe_out = e.uninit(t * n_embd)?;
3791            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3792            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3793                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3794            {
3795                let n_ff_sh = gate_shexp.out_features();
3796                let sg_gate = e.matmul(gate_shexp, z, t)?;
3797                let sg_up = e.matmul(up_shexp, z, t)?;
3798                let mut sa = e.uninit(t * n_ff_sh)?;
3799                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3800                let sh = e.matmul(down_shexp, &sa, t)?;
3801                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
3802                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
3803                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
3804                // i.e. the one real prefill actually takes on a resident-expert MoE model,
3805                // so the concat-prime isolation fix has to land here as well.
3806                let g = match &m.gate_inp_shexp {
3807                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3808                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3809                    }
3810                    Some(gate_inp_shexp) => {
3811                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3812                        let mut g = e.uninit(t)?;
3813                        e.sigmoid(&gs, &mut g, t)?;
3814                        g
3815                    }
3816                    None => e.htod(&vec![1.0f32; t])?,
3817                };
3818                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3819            }
3820            return Ok(moe_out);
3821        }
3822
3823        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
3824        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
3825        let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
3826        let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
3827                      inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3828            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
3829            let dec = dec && q8_expert_dec_supported(qtype);
3830            if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3831                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3832            else   { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3833                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3834        };
3835        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3836        let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
3837                          n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
3838        let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
3839                        n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
3840        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3841        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3842        // down consumes PAIR-major activation rows: pair_tok = identity.
3843        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3844        let pself = e.htod_i32(&pair_self)?;
3845        let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
3846                            n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
3847        let mut moe_out = e.uninit(t * n_embd)?;   // scatter fully overwrites per (token,col)
3848        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3849
3850        // SHARED EXPERT epilogue — same as the other paths.
3851        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
3852        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
3853        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3854            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3855        {
3856            let n_ff_sh = gate_shexp.out_features();
3857            let sg_gate = e.matmul(gate_shexp, z, t)?;
3858            let sg_up = e.matmul(up_shexp, z, t)?;
3859            let mut sa = e.uninit(t * n_ff_sh)?;
3860            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3861            let sh = e.matmul(down_shexp, &sa, t)?;
3862            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
3863            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
3864            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
3865            // dispatch choice cannot change bits.
3866            let g = match &m.gate_inp_shexp {
3867                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3868                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3869                }
3870                Some(gate_inp_shexp) => {
3871                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3872                    let mut g = e.uninit(t)?;
3873                    e.sigmoid(&gs, &mut g, t)?;
3874                    g
3875                }
3876                None => e.htod(&vec![1.0f32; t])?,
3877            };
3878            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3879        }
3880        Ok(moe_out)
3881    }
3882
3883    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
3884    #[allow(clippy::too_many_arguments)]
3885    #[allow(clippy::too_many_arguments)]
3886    fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3887                   zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
3888                   t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
3889                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3890        let moe = cfg.moe.as_ref().unwrap();
3891        let n_embd = cfg.n_embd as usize;
3892        let n_expert = moe.expert_count as usize;
3893        let n_used = moe.expert_used_count as usize;
3894        let n_ff_exp = moe.expert_ff_length as usize;
3895        // moe_router_topk is SOFTMAX-only (no exp_probs_b bias, no expert_weights_scale) and every
3896        // gate_up kernel here fuses PLAIN silu(gate)*up. `dev_ok` denies sigmoid-router archs and
3897        // clamped layers; assert both so a future caller that skips the gate fails loudly.
3898        debug_assert!(cfg.sigmoid_router().is_none(),
3899                      "moe_ffn_dev routes SOFTMAX: a sigmoid-router arch would pick wrong experts");
3900        debug_assert!(!cfg.swiglu_clamped_at(il as u32),
3901                      "moe_ffn_dev's fused epilogue is plain SiLU: no clamped form");
3902
3903        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
3904        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
3905        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
3906        // skipped entirely for macro-free experts (every k-quant GGUF).
3907        if m.has_macros {
3908            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
3909        }
3910
3911        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
3912        let mut moe_out = e.uninit(t * n_embd)?;
3913
3914        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
3915        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
3916        if let Some(dev) = m.dev_exps.as_ref() {
3917            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
3918            // the combined stride; up's base is offset in the ptr table. Down unchanged.
3919            let (rbg_d, rbu_d) = if dev.gu_il {
3920                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3921            } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3922            let q8 = moe_q8_enabled()
3923                && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3924                && q8_expert_supported(m.down_exps.qtype);
3925            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
3926            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
3927            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
3928            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
3929            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
3930            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
3931            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
3932            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
3933            let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
3934                && n_ff_exp == 512 && n_used <= 8
3935                && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
3936                && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
3937            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
3938            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
3939            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
3940            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
3941            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
3942            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
3943            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
3944            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
3945            let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
3946                .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
3947            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
3948            let csr_arm = rows_arm && csr_mode > 0 && t <= 10
3949                && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
3950                && csr_qt(m.down_exps.qtype);
3951            if csr_arm {
3952                if csr_mode == 2 {
3953                    static ENGAGED: std::sync::Once = std::sync::Once::new();
3954                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
3955                }
3956                let n_pairs = t * n_used;
3957                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3958                let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
3959                                                         n_embd, n_ff_exp, n_used, n_expert,
3960                                                         m.gate_exps.qtype, m.up_exps.qtype,
3961                                                         rbg_d, rbu_d)?;
3962                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3963                // down stays on the _rows twin — BOTH CSR down variants measured negative
3964                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
3965                // 16-group rows have too little decode to amortize any dedup structure.
3966                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
3967                                            t, n_ff_exp, n_embd, n_used, n_expert,
3968                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
3969                if csr_mode == 2 {
3970                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
3971                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
3972                                                                n_embd, n_ff_exp, n_used, n_expert,
3973                                                                m.gate_exps.qtype, m.up_exps.qtype,
3974                                                                rbg_d, rbu_d, &m.dev_macros)?;
3975                    let mut out_r = e.uninit(t * n_embd)?;
3976                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
3977                    e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
3978                                                t, n_ff_exp, n_embd, n_used, n_expert,
3979                                                m.down_exps.qtype, m.down_exps.row_bytes)?;
3980                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
3981                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
3982                    let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3983                    let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3984                    if ba + bo > 0 {
3985                        eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
3986                                  a1.len(), o1.len());
3987                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
3988                        let sel_h = e.dtoh_i32(&sel_d)?;
3989                        let mut shown = 0;
3990                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
3991                            if x.to_bits() != y.to_bits() && shown < 4 {
3992                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
3993                                let ex = sel_h[p];
3994                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
3995                                eprintln!("  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
3996                                shown += 1;
3997                            }
3998                        }
3999                        std::process::exit(3);
4000                    }
4001                }
4002            } else if rows_arm {
4003                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
4004                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
4005                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
4006                    use std::sync::atomic::{AtomicU64, Ordering};
4007                    static PAIRS: AtomicU64 = AtomicU64::new(0);
4008                    static UNIQ: AtomicU64 = AtomicU64::new(0);
4009                    static CALLS: AtomicU64 = AtomicU64::new(0);
4010                    let sel_h = e.dtoh_i32(&sel_d)?;
4011                    let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
4012                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
4013                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
4014                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
4015                    if c % 480 == 0 {
4016                        let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
4017                        eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
4018                                  q as f64 / p as f64);
4019                    }
4020                }
4021                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4022                let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
4023                                                          n_embd, n_ff_exp, n_used, n_expert,
4024                                                          m.gate_exps.qtype, m.up_exps.qtype,
4025                                                          rbg_d, rbu_d, &m.dev_macros)?;
4026                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
4027                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
4028                                            t, n_ff_exp, n_embd, n_used, n_expert,
4029                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
4030            } else {
4031            for tok in 0..t {
4032                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4033                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4034                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4035                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4036                if q8 {
4037                    let (zq, zd) = match (t, zq8) {
4038                        (1, Some((q, d))) => (q.clone(), d.clone()),
4039                        _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
4040                    };
4041                    let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
4042                                                         n_embd, n_ff_exp, n_used, n_expert,
4043                                                         m.gate_exps.qtype, m.up_exps.qtype,
4044                                                         rbg_d, rbu_d, &m.dev_macros)?;
4045                    let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4046                    e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
4047                                           n_ff_exp, n_embd, n_used, n_expert,
4048                                           m.down_exps.qtype, m.down_exps.row_bytes)?;
4049                } else {
4050                    let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
4051                                                      n_used, n_expert,
4052                                                      m.gate_exps.qtype, m.up_exps.qtype,
4053                                                      rbg_d, rbu_d, &m.dev_macros)?;
4054                    e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
4055                                        n_ff_exp, n_embd, n_used, n_expert,
4056                                        m.down_exps.qtype, m.down_exps.row_bytes)?;
4057                }
4058            }
4059            }
4060        } else {
4061        // Launch under the cache lock: the row borrow lives as long as the closure, and the
4062        // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
4063        // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
4064        // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
4065        // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
4066        // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
4067        let q8 = moe_q8_enabled()
4068            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
4069            && q8_expert_supported(m.down_exps.qtype);
4070        e.with_moe_cache(max_block, |c, eng| {
4071            let row = c.layer_dev_row(il, n_expert, eng)?
4072                .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
4073            for tok in 0..t {
4074                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
4075                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
4076                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
4077                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4078                if q8 {
4079                    let (zq, zd) = match (t, zq8) {
4080                        (1, Some((q, d))) => (q.clone(), d.clone()),
4081                        _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
4082                    };
4083                    let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
4084                                                           n_embd, n_ff_exp, n_used, n_expert,
4085                                                           m.gate_exps.qtype, m.up_exps.qtype,
4086                                                           m.gate_exps.row_bytes, m.up_exps.row_bytes,
4087                                                           &m.dev_macros)?;
4088                    let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
4089                    eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
4090                                             n_ff_exp, n_embd, n_used, n_expert,
4091                                             m.down_exps.qtype, m.down_exps.row_bytes)?;
4092                } else {
4093                    let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
4094                                                        n_used, n_expert,
4095                                                        m.gate_exps.qtype, m.up_exps.qtype,
4096                                                        m.gate_exps.row_bytes, m.up_exps.row_bytes,
4097                                                        &m.dev_macros)?;
4098                    eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
4099                                          n_ff_exp, n_embd, n_used, n_expert,
4100                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
4101                }
4102            }
4103            // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
4104            c.hits += (t * 3 * n_used) as u64;
4105            Ok(())
4106        })?;
4107        }
4108
4109        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
4110        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
4111        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
4112        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
4113        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4114            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4115        {
4116            let n_ff_sh = gate_shexp.out_features();
4117            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
4118            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
4119            let verify_t = t > 1 && t < PRIME_MIN_T;
4120            let (sg_gate, sg_up) = if t == 1 {
4121                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
4122                    Some(pair) => pair,
4123                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
4124                }
4125            } else if verify_t {
4126                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
4127                // rides one shared quantize + one fused2 batched launch instead of two
4128                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
4129                let mut fused = None;
4130                if crate::spec::spec_fused_t() && (2..=4).contains(&t)
4131                    && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
4132                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
4133                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
4134                }
4135                match fused {
4136                    Some(pair) => pair,
4137                    None => (e.matmul_decode_exact(gate_shexp, z, t)?,
4138                             e.matmul_decode_exact(up_shexp, z, t)?),
4139                }
4140            } else {
4141                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
4142            };
4143            let mut sa = e.uninit(t * n_ff_sh)?;  // silu_mul fully overwrites
4144            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
4145            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
4146                     else { e.matmul(down_shexp, &sa, t)? };
4147            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
4148            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
4149            // between the two arms; prefill keeps the batched cuBLASLt linear).
4150            let g = match &m.gate_inp_shexp {
4151                Some(gate_inp_shexp) => {
4152                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
4153                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
4154                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4155                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4156                    } else {
4157                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4158                        let mut g = e.uninit(t)?;
4159                        e.sigmoid(&gs, &mut g, t)?;
4160                        g
4161                    }
4162                }
4163                None => e.htod(&vec![1.0f32; t])?,
4164            };
4165            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4166        }
4167
4168        Ok(moe_out)
4169    }
4170
4171    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
4172    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
4173    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
4174    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
4175    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
4176    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
4177    /// the collected raw pointers cannot move between collection and launch (single-threaded
4178    /// decode; the lock is held only for collection, launches are stream-ordered after any
4179    /// prior same-stream staging writes).
4180    #[allow(clippy::too_many_arguments)]
4181    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
4182    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
4183    #[allow(clippy::too_many_arguments)]
4184    fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
4185                      zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
4186                      moe_out: &mut CudaSlice<f32>, tok: usize,
4187                      n_embd: usize, n_ff_exp: usize, n_used: usize)
4188                      -> Result<bool, Box<dyn std::error::Error>> {
4189        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4190        use cudarc::driver::DevicePtr;
4191        let ptrs = e.with_moe_cache(max_block, |c, eng| {
4192            let mut g = [0u64; 8];
4193            let mut u = [0u64; 8];
4194            let mut d = [0u64; 8];
4195            for (j, &ex) in sel.iter().enumerate() {
4196                let ex = ex as u16;
4197                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
4198                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
4199                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
4200                else { return Ok(None); };
4201                let __s = eng.stream();
4202                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
4203                let (pu, _e1) = c.slot(su).device_ptr(&__s);
4204                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
4205                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
4206            }
4207            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
4208                for &ex in sel {
4209                    let ex = ex as u16;
4210                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
4211                        c.note_profile_hit(BlockId::new(il, proj, ex));
4212                    }
4213                }
4214            }
4215            c.hits += (3 * n_used) as u64;
4216            Ok(Some((g, u, d)))
4217        })?;
4218        let Some((g, u, d)) = ptrs else { return Ok(false) };
4219        let mut wv = [0f32; 8];
4220        wv[..n_used].copy_from_slice(w);
4221        let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
4222                                         n_embd, n_ff_exp, n_used,
4223                                         m.gate_exps.qtype, m.up_exps.qtype,
4224                                         m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4225        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
4226        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4227        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4228        e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
4229                           n_ff_exp, n_embd, n_used,
4230                           m.down_exps.qtype, m.down_exps.row_bytes)?;
4231        Ok(true)
4232    }
4233
4234    fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
4235                      zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
4236                      moe_out: &mut CudaSlice<f32>, tok: usize,
4237                      n_embd: usize, n_ff_exp: usize, n_used: usize)
4238                      -> Result<bool, Box<dyn std::error::Error>> {
4239        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4240        use cudarc::driver::DevicePtr;
4241        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
4242        let ptrs = e.with_moe_cache(max_block, |c, eng| {
4243            let mut g = [0u64; 8];
4244            let mut u = [0u64; 8];
4245            let mut d = [0u64; 8];
4246            for (j, &ex) in sel.iter().enumerate() {
4247                let ex = ex as u16;
4248                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
4249                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
4250                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
4251                else { return Ok(None); };
4252                let __s = eng.stream();
4253                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
4254                let (pu, _e1) = c.slot(su).device_ptr(&__s);
4255                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
4256                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
4257            }
4258            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
4259                for &ex in sel {
4260                    let ex = ex as u16;
4261                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
4262                        c.note_profile_hit(BlockId::new(il, proj, ex));
4263                    }
4264                }
4265            }
4266            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
4267            Ok(Some((g, u, d)))
4268        })?;
4269        let Some((g, u, d)) = ptrs else { return Ok(false) };
4270        let mut wv = [0f32; 8];
4271        wv[..n_used].copy_from_slice(w);
4272        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
4273        let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
4274                                      n_embd, n_ff_exp, n_used,
4275                                      m.gate_exps.qtype, m.up_exps.qtype,
4276                                      m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4277        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4278        e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
4279                             n_ff_exp, n_embd, n_used,
4280                             m.down_exps.qtype, m.down_exps.row_bytes)?;
4281        Ok(true)
4282    }
4283
4284    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
4285    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
4286    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
4287    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
4288    fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
4289                          max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
4290                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4291        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
4292        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
4293        let layout = exps.expert_layout(ex);
4294        let id = BlockId::new(il, proj, ex as u16);
4295        let source = exps.expert_source(ex);
4296        e.with_moe_cache(max_block, |c, eng| {
4297            let slot = c.dispatch_source(id, source, eng)?;
4298            let DispatchSlot::Resident(sl) = slot;
4299            let buf = c.slot(sl);
4300            eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
4301                                  layout.qtype, layout.row_bytes)
4302        })
4303    }
4304
4305    fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
4306                       max_block: usize, x: &cudarc::driver::CudaView<f32>)
4307                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4308        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
4309        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
4310        let layout = exps.expert_layout(ex);
4311        let id = BlockId::new(il, proj, ex as u16);
4312        let source = exps.expert_source(ex);
4313        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
4314        e.with_moe_cache(max_block, |c, eng| {
4315            let slot = c.dispatch_source(id, source, eng)?;
4316            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
4317            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
4318            let DispatchSlot::Resident(sl) = slot;
4319            let buf = c.slot(sl);
4320            eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
4321                             layout.qtype, layout.row_bytes)
4322        })
4323    }
4324
4325    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
4326    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
4327    /// so the current forward's backend assignment and output remain unchanged.
4328    fn moe_profile_admit_expert(
4329        e: &Engine,
4330        il: u16,
4331        ex: usize,
4332        m: &MoeWeights,
4333        max_block: usize,
4334    ) -> Result<(), Box<dyn std::error::Error>> {
4335        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4336        e.with_moe_cache(max_block, |cache, eng| {
4337            for (proj, exps) in [
4338                (PROJ_GATE, &m.gate_exps),
4339                (PROJ_UP, &m.up_exps),
4340                (PROJ_DOWN, &m.down_exps),
4341            ] {
4342                let id = BlockId::new(il, proj, ex as u16);
4343                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
4344            }
4345            Ok(())
4346        })
4347    }
4348
4349    /// Read a projection from the immutable residency set when present; otherwise use one
4350    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
4351    #[allow(clippy::too_many_arguments)]
4352    fn moe_frozen_gemm(
4353        e: &Engine,
4354        il: u16,
4355        proj: u8,
4356        ex: usize,
4357        m: &MoeWeights,
4358        max_block: usize,
4359        x: &cudarc::driver::CudaView<f32>,
4360        scratch: &mut Option<CudaSlice<u8>>,
4361        scratch_len: usize,
4362    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4363        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
4364        let exps = match proj {
4365            PROJ_GATE => &m.gate_exps,
4366            PROJ_UP => &m.up_exps,
4367            _ => &m.down_exps,
4368        };
4369        let layout = exps.expert_layout(ex);
4370        let id = BlockId::new(il, proj, ex as u16);
4371        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
4372            let Some(slot) = cache.resident(id) else {
4373                return Ok(None);
4374            };
4375            let buf = cache.slot(slot);
4376            Ok(Some(eng.qmatvec_view(
4377                buf,
4378                0..layout.len,
4379                x,
4380                1,
4381                exps.in_f,
4382                exps.out_f,
4383                layout.qtype,
4384                layout.row_bytes,
4385            )?))
4386        })? {
4387            return Ok(output);
4388        }
4389        if scratch.is_none() {
4390            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
4391        }
4392        let scratch = scratch.as_mut().unwrap();
4393        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
4394        e.qmatvec_view(
4395            scratch,
4396            0..layout.len,
4397            x,
4398            1,
4399            exps.in_f,
4400            exps.out_f,
4401            layout.qtype,
4402            layout.row_bytes,
4403        )
4404    }
4405
4406    fn moe_prefetch_expert(
4407        e: &Engine,
4408        il: u16,
4409        ex: usize,
4410        m: &MoeWeights,
4411        max_block: usize,
4412        keep: &[crate::moe_cache::BlockId],
4413    ) -> Result<(), Box<dyn std::error::Error>> {
4414        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4415        e.with_moe_cache(max_block, |c, eng| {
4416            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
4417                                 (PROJ_DOWN, &m.down_exps)] {
4418                let id = BlockId::new(il, proj, ex as u16);
4419                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
4420            }
4421            Ok(())
4422        })
4423    }
4424
4425    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
4426    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
4427    fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
4428                                max_block: usize, keep: &[crate::moe_cache::BlockId])
4429                                -> Result<(), Box<dyn std::error::Error>> {
4430        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4431        e.with_moe_cache(max_block, |c, eng| {
4432            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
4433                                 (PROJ_DOWN, &m.down_exps)] {
4434                let source = exps.expert_source(ex);
4435                if let crate::model::ExpertSource::Disk { .. } = &source {
4436                    let id = BlockId::new(il, proj, ex as u16);
4437                    let _ = c.prefetch_source(id, source, keep, eng)?;
4438                }
4439            }
4440            Ok(())
4441        })
4442    }
4443
4444    #[inline]
4445    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
4446        let _ = m.gate_exps.prefetch_expert_pages(ex);
4447        let _ = m.up_exps.prefetch_expert_pages(ex);
4448        let _ = m.down_exps.prefetch_expert_pages(ex);
4449    }
4450}
4451
4452// ================================================================================================
4453// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
4454//
4455// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
4456// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
4457// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
4458//
4459// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
4460// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
4461// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
4462// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
4463// identical to the per-token loop regardless of expert processing order.
4464//
4465// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
4466// ================================================================================================
4467
4468impl HybridModel {
4469    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
4470    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
4471    pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
4472                                  cfg: &ModelConfig, il: u16, _max_block: usize)
4473                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4474        let moe = cfg.moe.as_ref().unwrap();
4475        let n_embd = cfg.n_embd as usize;
4476        let n_expert = moe.expert_count as usize;
4477        let n_used = moe.expert_used_count as usize;
4478        let n_ff_exp = moe.expert_ff_length as usize;
4479        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
4480        let lim_exp = cfg.clamp_exp_at(il as u32);
4481        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4482
4483        // 1. ROUTER (identical to moe_ffn).
4484        let logits = e.matmul(&m.gate_inp, z, t)?;
4485        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4486            Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4487                                m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4488        } else {
4489            Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4490                                None, None, m.active_experts.as_deref())?
4491        };
4492        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4493
4494        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
4495        // For each expert e, we need: which tokens use it, their positions in z, their top-k
4496        // slot index (for bit-identical accumulation), and their weights.
4497        struct ExpertGroup {
4498            tok_indices: Vec<i32>,   // indices into z rows (0..T-1)
4499            slot_indices: Vec<i32>,  // top-k slot (0..n_used-1) for that token-expert pair
4500            weights: Vec<f32>,       // renormalized weight for that token-expert pair
4501        }
4502        let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
4503            tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
4504        }).collect();
4505
4506        for tok in 0..t {
4507            for j in 0..n_used {
4508                let ex = sel_all[tok * n_used + j] as usize;
4509                let w = w_all[tok * n_used + j];
4510                groups[ex].tok_indices.push(tok as i32);
4511                groups[ex].slot_indices.push(j as i32);
4512                groups[ex].weights.push(w);
4513            }
4514        }
4515
4516        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
4517        // Each token's 8 expert contributions land in their respective slots.
4518        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
4519        let mut wbuf = e.zeros(t * n_used)?;  // [T, n_used] weight buffer for FMA reduce
4520
4521        // Expert weight dimensions (used in both cache and staging paths).
4522        let g_len = m.gate_exps.max_expert_bytes();
4523        let u_len = m.up_exps.max_expert_bytes();
4524        let d_len = m.down_exps.max_expert_bytes();
4525        let use_cache = Engine::moe_cache_enabled();
4526        let max_block = _max_block;
4527
4528        // GPU scratch for staging (only allocated when NOT using cache).
4529        let (mut scratch_g, mut scratch_u, mut scratch_d) = if !use_cache {
4530            (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
4531        } else {
4532            (None, None, None)
4533        };
4534
4535        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
4536        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
4537        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
4538        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
4539        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
4540        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
4541        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
4542        // at long prompts where every expert stages regardless. Order is FREE to change without
4543        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
4544        // regardless of expert processing order (the whole point of the slots).
4545        let mut order: Vec<usize> =
4546            (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
4547        order.sort_by(|&a, &b| groups[b].tok_indices.len()
4548            .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
4549        let mut m_dist: Vec<usize> = Vec::new();  // for stats
4550        let page_window = moe_page_prefetch_window();
4551        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
4552        if worker_disk_prefetch {
4553            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
4554                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
4555            }
4556        }
4557        for (order_pos, &ex) in order.iter().enumerate() {
4558            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
4559                Self::moe_prefetch_host_expert(order[next], m);
4560            }
4561            if worker_disk_prefetch {
4562                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
4563                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4564                    let keep = [
4565                        BlockId::new(il, PROJ_GATE, ex as u16),
4566                        BlockId::new(il, PROJ_UP, ex as u16),
4567                        BlockId::new(il, PROJ_DOWN, ex as u16),
4568                    ];
4569                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
4570                }
4571            }
4572            let grp = &groups[ex];
4573            let m_e = grp.tok_indices.len();
4574            m_dist.push(m_e);
4575            let gl = m.gate_exps.expert_layout(ex);
4576            let ul = m.up_exps.expert_layout(ex);
4577            let dl = m.down_exps.expert_layout(ex);
4578
4579            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
4580            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
4581            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
4582            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
4583            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
4584            let dmac = m.down_exps.macro_scale(ex);
4585            let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
4586                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
4587                e.htod(&scaled)?
4588            };
4589
4590            // GATHER: collect m_e activation rows from z into a contiguous buffer.
4591            let mut gathered = e.zeros(m_e * n_embd)?;
4592            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
4593            let gv = gathered.slice(0..m_e * n_embd);
4594
4595            // Compute gate/up/down matmuls -- two paths: cache-resident or host-staged.
4596            let y = if use_cache {
4597                use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4598                // CACHE PATH: dispatch through MOE cache, get device-resident buffer, GEMM at m=m_e.
4599                let gate = e.with_moe_cache(max_block, |c, eng| {
4600                    let id = BlockId::new(il, PROJ_GATE, ex as u16);
4601                    let slot = c.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
4602                    let buf = c.buf(slot);
4603                    eng.qmatvec_view(buf, 0..gl.len, &gv, m_e,
4604                        m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4605                })?;
4606                let up = e.with_moe_cache(max_block, |c, eng| {
4607                    let id = BlockId::new(il, PROJ_UP, ex as u16);
4608                    let slot = c.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
4609                    let buf = c.buf(slot);
4610                    eng.qmatvec_view(buf, 0..ul.len, &gv, m_e,
4611                        m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4612                })?;
4613                // SiLU-MUL activation (per-expert macro-scales folded; step35's per-layer clamp).
4614                let mut act = e.zeros(m_e * n_ff_exp)?;
4615                Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4616                    m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4617                let actv = act.slice(0..m_e * n_ff_exp);
4618                e.with_moe_cache(max_block, |c, eng| {
4619                    let id = BlockId::new(il, PROJ_DOWN, ex as u16);
4620                    let slot = c.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
4621                    let buf = c.buf(slot);
4622                    eng.qmatvec_view(buf, 0..dl.len, &actv, m_e,
4623                        m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4624                })?
4625            } else {
4626                // STAGING PATH: H2D the expert blocks into scratch buffers, then GEMM.
4627                let sg = scratch_g.as_mut().unwrap();
4628                let su = scratch_u.as_mut().unwrap();
4629                let sd = scratch_d.as_mut().unwrap();
4630                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4631                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4632                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4633                let gate = e.qmatvec_view(sg, 0..gl.len, &gv, m_e,
4634                    m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
4635                let up = e.qmatvec_view(su, 0..ul.len, &gv, m_e,
4636                    m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
4637                // SiLU-MUL activation (per-expert macro-scales folded; step35's per-layer clamp).
4638                let mut act = e.zeros(m_e * n_ff_exp)?;
4639                Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4640                    m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4641                let actv = act.slice(0..m_e * n_ff_exp);
4642                e.qmatvec_view(sd, 0..dl.len, &actv, m_e,
4643                    m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?
4644            };
4645
4646            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
4647            e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
4648                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4649        }
4650
4651        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
4652        let mut moe_out = e.zeros(t * n_embd)?;
4653        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
4654
4655        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
4656        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
4657            m_dist.sort_unstable();
4658            let active = m_dist.len();
4659            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
4660            let median = m_dist[active / 2];
4661            let max_m = *m_dist.last().unwrap();
4662            let min_m = m_dist[0];
4663            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
4664            println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
4665                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
4666                      above_gemm_threshold(>=16)={above16}/{active}");
4667        }
4668
4669        // 6. SHARED EXPERT (same as moe_ffn — untouched).
4670        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
4671        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
4672        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4673            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4674        {
4675            let n_ff_sh = gate_shexp.out_features();
4676            let sg_gate = e.matmul(gate_shexp, z, t)?;
4677            let sg_up = e.matmul(up_shexp, z, t)?;
4678            let mut sa = e.zeros(t * n_ff_sh)?;
4679            Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp, &mut sa, t * n_ff_sh)?;
4680            let sh = e.matmul(down_shexp, &sa, t)?;
4681            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
4682            // Fused sigmoid-dot below PRIME_MIN_T — one fold order with the sequential and
4683            // dev decode arms (dispatch choice must not change bits).
4684            let g = match &m.gate_inp_shexp {
4685                Some(gate_inp_shexp) => {
4686                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
4687                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
4688                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4689                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4690                    } else {
4691                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4692                        let mut g = e.uninit(t)?;
4693                        e.sigmoid(&gs, &mut g, t)?;
4694                        g
4695                    }
4696                }
4697                None => e.htod(&vec![1.0f32; t])?,
4698            };
4699            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4700        }
4701
4702        Ok(moe_out)
4703    }
4704
4705    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
4706    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
4707    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
4708    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
4709    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
4710    /// expert-sum order identical to the sequential path.
4711    pub(crate) fn moe_ffn_lockstep(
4712        &self,
4713        e: &Engine,
4714        m: &MoeWeights,
4715        zbatch: &CudaSlice<f32>,
4716        mrows: usize,
4717        il: u16,
4718        max_block: usize,
4719    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4720        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4721        let cfg = &self.cfg;
4722        let moe = cfg.moe.as_ref().unwrap();
4723        let n_embd = cfg.n_embd as usize;
4724        let n_expert = moe.expert_count as usize;
4725        let n_used = moe.expert_used_count as usize;
4726        let n_ff_exp = moe.expert_ff_length as usize;
4727        // step35 per-layer SwiGLU clamp; None on every other arch / unclamped layer.
4728        let lim_exp = cfg.clamp_exp_at(il as u32);
4729        let lim_shexp = cfg.clamp_shexp_at(il as u32);
4730
4731        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
4732        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4733            Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4734                                m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4735        } else {
4736            Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4737                                None, None, m.active_experts.as_deref())?
4738        };
4739        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
4740
4741        // Residency split at whole-expert granularity against the (frozen) cache.
4742        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
4743            Ok((0..n_expert)
4744                .map(|ex| {
4745                    [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
4746                        c.resident(BlockId::new(il, p, ex as u16)).is_some()
4747                    })
4748                })
4749                .collect())
4750        })?;
4751
4752        struct Group {
4753            rows: Vec<i32>,
4754            slots: Vec<i32>,
4755            weights: Vec<f32>,
4756        }
4757        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
4758        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
4759        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
4760            Default::default();
4761        for row in 0..mrows {
4762            for j in 0..n_used {
4763                let ex = sel_all[row * n_used + j] as usize;
4764                let w = w_all[row * n_used + j];
4765                if resident_expert[ex] {
4766                    let group = groups.entry(ex).or_insert_with(|| Group {
4767                        rows: Vec::new(),
4768                        slots: Vec::new(),
4769                        weights: Vec::new(),
4770                    });
4771                    group.rows.push(row as i32);
4772                    group.slots.push(j as i32);
4773                    group.weights.push(w);
4774                } else {
4775                    crate::cpu_experts::record_incomplete_gpu_residency(0);
4776                    cpu_rows[row].push((ex, w));
4777                    cpu_by_expert.entry(ex).or_default().push((row, w));
4778                }
4779            }
4780        }
4781
4782        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
4783        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
4784        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
4785        // order per row differs from the sequential single-call chunk — part of the
4786        // documented lockstep numeric class.
4787        let host_rows = e.dtoh(zbatch)?;
4788        let rows_ok = crate::cpu_experts::rows_supported();
4789        enum CpuPart {
4790            Single { row: usize },
4791            Rows { rows: Vec<usize> },
4792        }
4793        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
4794        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
4795        if rows_ok {
4796            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
4797                .into_iter()
4798                .filter(|(_, rows)| rows.len() >= 2)
4799                .collect();
4800            shared.sort_by_key(|(ex, _)| *ex);
4801            for (ex, mut row_weights) in shared {
4802                row_weights.sort_by_key(|(row, _)| *row);
4803                let inputs: Vec<(&[f32], f32)> = row_weights
4804                    .iter()
4805                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
4806                    .collect();
4807                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
4808                    .map_err(std::io::Error::other)?;
4809                for &(row, _) in &row_weights {
4810                    rows_served.insert((row, ex));
4811                }
4812                tickets.push((
4813                    CpuPart::Rows {
4814                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
4815                    },
4816                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
4817                ));
4818            }
4819        }
4820        for (row, selected) in cpu_rows.iter().enumerate() {
4821            let leftover: Vec<(usize, f32)> = selected
4822                .iter()
4823                .copied()
4824                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
4825                .collect();
4826            if leftover.is_empty() {
4827                continue;
4828            }
4829            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
4830            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
4831                .map_err(std::io::Error::other)?;
4832            tickets.push((
4833                CpuPart::Single { row },
4834                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
4835            ));
4836        }
4837
4838        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
4839        let mut wbuf = e.zeros(mrows * n_used)?;
4840        let mut order: Vec<usize> = groups.keys().copied().collect();
4841        order.sort_by(|&a, &b| {
4842            groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
4843        });
4844        for &ex in &order {
4845            let group = &groups[&ex];
4846            let m_e = group.rows.len();
4847            let gl = m.gate_exps.expert_layout(ex);
4848            let ul = m.up_exps.expert_layout(ex);
4849            let dl = m.down_exps.expert_layout(ex);
4850            let row_idx_d = e.htod_i32(&group.rows)?;
4851            let slot_idx_d = e.htod_i32(&group.slots)?;
4852            let dmac = m.down_exps.macro_scale(ex);
4853            let weight_d = if dmac == 1.0 {
4854                e.htod(&group.weights)?
4855            } else {
4856                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
4857                e.htod(&scaled)?
4858            };
4859            let mut gathered = e.zeros(m_e * n_embd)?;
4860            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
4861            let gv = gathered.slice(0..m_e * n_embd);
4862            let gate = e.with_moe_cache(max_block, |c, eng| {
4863                let slot = c
4864                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
4865                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4866                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
4867                    m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4868            })?;
4869            let up = e.with_moe_cache(max_block, |c, eng| {
4870                let slot = c
4871                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
4872                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4873                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
4874                    m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4875            })?;
4876            let mut act = e.zeros(m_e * n_ff_exp)?;
4877            Self::ffn_act_lim(e, cfg, &gate, &up, m.gate_exps.macro_scale(ex),
4878                m.up_exps.macro_scale(ex), lim_exp, &mut act, m_e * n_ff_exp)?;
4879            let actv = act.slice(0..m_e * n_ff_exp);
4880            let y = e.with_moe_cache(max_block, |c, eng| {
4881                let slot = c
4882                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
4883                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4884                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
4885                    m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4886            })?;
4887            e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
4888                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4889        }
4890        let mut moe_out = e.zeros(mrows * n_embd)?;
4891        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
4892
4893        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
4894        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
4895        for (part, ticket) in tickets {
4896            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
4897            let mut add_row = |row: usize, chunk: &[f32]| {
4898                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
4899                for (accumulator, value) in sum.iter_mut().zip(chunk) {
4900                    *accumulator += value;
4901                }
4902            };
4903            match part {
4904                CpuPart::Single { row } => add_row(row, &cpu_output),
4905                CpuPart::Rows { rows } => {
4906                    for (slot, row) in rows.into_iter().enumerate() {
4907                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
4908                    }
4909                }
4910            }
4911        }
4912        for (row, sum) in row_sums.into_iter().enumerate() {
4913            let Some(sum) = sum else { continue };
4914            let cpu_output = e.htod(&sum)?;
4915            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
4916            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4917        }
4918
4919        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4920            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4921        {
4922            let n_ff_sh = gate_shexp.out_features();
4923            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
4924            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
4925            let mut sa = e.zeros(mrows * n_ff_sh)?;
4926            Self::ffn_act_lim(e, cfg, &sg_gate, &sg_up, 1.0, 1.0, lim_shexp,
4927                              &mut sa, mrows * n_ff_sh)?;
4928            let sh = e.matmul(down_shexp, &sa, mrows)?;
4929            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
4930            // decode matches the single-sequence decode chain bit-for-bit.
4931            let g = match &m.gate_inp_shexp {
4932                Some(gate_inp_shexp) => {
4933                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
4934                }
4935                None => e.htod(&vec![1.0f32; mrows])?,
4936            };
4937            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
4938        }
4939
4940        Ok(moe_out)
4941    }
4942}
4943
4944// ============================ gemma4 (R8 verified wiring) ==================================
4945// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
4946// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
4947// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
4948// gemma variants after the correctness gate).
4949impl HybridModel {
4950    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
4951    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
4952        let g = self.cfg.gemma4.as_ref().unwrap();
4953        let swa = g.swa_pattern[il];
4954        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
4955        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
4956        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
4957        // rows exact (softmax over one element) while every later position drifted).
4958        (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
4959         if swa { g.rope_base_swa } else { g.rope_base_global },
4960         1.0, swa)
4961    }
4962
4963    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
4964    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
4965    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
4966    fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
4967                       -> Result<(), Box<dyn std::error::Error>> {
4968        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
4969            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
4970        }
4971        Ok(())
4972    }
4973
4974    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
4975    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
4976    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
4977    /// only (v0): attends within `tokens` via the f32 sdpa.
4978    fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
4979                         h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
4980                         cache: Option<&mut Cache>)
4981                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4982        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
4983        let eps = self.cfg.rms_eps;
4984        let aux = self.gemma4_aux.as_ref().unwrap();
4985
4986        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
4987        // (h stays borrowed across the triple, so the cache key can't go stale).
4988        e.mmq_act_begin();
4989        let q0 = e.matmul(&fa.wq, h, t)?;   // [t, nh*hd]
4990        let k0 = e.matmul(&fa.wk, h, t)?;   // [t, nkv*hd]
4991        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
4992        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
4993        let v0 = if swa { e.matmul(&fa.wv, h, t)? } else { e.clone_dtod(&k0)? };
4994
4995        let mut q = e.uninit(t * nh * hd)?;
4996        let mut k = e.uninit(t * nkv * hd)?;
4997        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
4998        let mut v = e.uninit(t * nkv * hd)?;
4999        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
5000        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
5001        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
5002        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5003        let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
5004            && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
5005        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
5006        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
5007        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
5008        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
5009        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
5010        let v_f16 = emit && crate::fa_f16pv_on() && match hd {
5011            512 => true,
5012            256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
5013            _ => false,
5014        };
5015        if emit {
5016            e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5017                               &aux.ones, &mut q, &mut k, &mut v, &mut vb,
5018                               hd, nh * t, nkv * t, eps, v_f16)?;
5019        } else {
5020            e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5021                           &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
5022        }
5023
5024        let ff = if swa { None } else {
5025            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5026        };
5027        if emit {
5028            e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
5029                               base, 1.0, ff)?;
5030        } else {
5031            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
5032        }
5033
5034        if let Some(cache) = cache {
5035            let kvl = cache.kv[il].as_mut().unwrap();
5036            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
5037            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
5038                                       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()))?;
5039            kvl.len += t;
5040        }
5041        let mut attn = e.zeros(t * nh * hd)?;
5042        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
5043        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
5044        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
5045        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5046        if swa && t > win {
5047            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
5048                if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
5049                                             scale, true, win, v_f16)?; }
5050                else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
5051                                      win)?; }
5052            } else {
5053                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
5054            }
5055        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
5056            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
5057        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
5058            if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
5059                                             scale, true, v_f16)?; }
5060            else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
5061        } else {
5062            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
5063        }
5064        Ok(e.matmul(&fa.wo, &attn, t)?)
5065    }
5066
5067    /// Back-compat wrapper (pure prefill, no cache).
5068    fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5069                   h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
5070                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5071        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
5072    }
5073
5074    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
5075    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
5076    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
5077    /// the q8z epilogue is quantize_q8_1 verbatim).
5078    fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
5079                     bits: &crate::hybrid::Gemma4MoeBits,
5080                     mq: &(CudaSlice<i8>, CudaSlice<f32>),
5081                     router_in: &CudaSlice<f32>, t: usize)
5082                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5083        let cfg = &self.cfg;
5084        let moe = cfg.moe.as_ref().unwrap();
5085        let n_embd = cfg.n_embd as usize;
5086        let n_expert = moe.expert_count as usize;
5087        let n_used = moe.expert_used_count as usize;
5088        let n_ff_exp = moe.expert_ff_length as usize;
5089        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
5090        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
5091        // the pair's 12us is kernel time, not launch gaps.
5092        let logits = if crate::router_kernel_on() {
5093            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
5094        } else {
5095            e.matmul(&m.gate_inp, router_in, t)?
5096        };
5097        let dev = m.dev_exps.as_ref().unwrap();
5098        let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
5099                                                    &bits.per_expert_scale_d)?;
5100        let (zq, zd) = mq;
5101        if t == 1 {
5102            let selv = sel_d.slice(0..n_used);
5103            let wv = w_d.slice(0..n_used);
5104            let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
5105                                                 n_embd, n_ff_exp, n_used, n_expert,
5106                                                 m.gate_exps.qtype, m.up_exps.qtype,
5107                                                 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5108            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5109            let mut moe_out = e.uninit(n_embd)?;
5110            e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
5111                                   &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
5112                                   n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
5113            return Ok(moe_out);
5114        }
5115        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
5116        let act = if csr {
5117            e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
5118                                           n_embd, n_ff_exp, n_used, n_expert,
5119                                           m.gate_exps.qtype, m.up_exps.qtype,
5120                                           m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5121        } else {
5122            e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
5123                                            n_embd, n_ff_exp, n_used, n_expert,
5124                                            m.gate_exps.qtype, m.up_exps.qtype,
5125                                            m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5126        };
5127        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5128        let mut moe_out = e.uninit(t * n_embd)?;
5129        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
5130        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
5131        e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
5132                                      n_ff_exp, n_embd, n_used, n_expert,
5133                                      m.down_exps.qtype, m.down_exps.row_bytes)?;
5134        Ok(moe_out)
5135    }
5136
5137    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
5138    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
5139    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
5140    fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
5141                  bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
5142                  router_in: &CudaSlice<f32>, t: usize)
5143                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5144        let cfg = &self.cfg;
5145        let moe = cfg.moe.as_ref().unwrap();
5146        let n_embd = cfg.n_embd as usize;
5147        let n_expert = moe.expert_count as usize;
5148        let n_used = moe.expert_used_count as usize;
5149        let n_ff_exp = moe.expert_ff_length as usize;
5150
5151        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
5152        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
5153        // batched matmul only at real prefill.
5154        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
5155            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
5156        } else {
5157            e.matmul(&m.gate_inp, router_in, t)?
5158        };
5159
5160        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
5161        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
5162        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
5163        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
5164        if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5165            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
5166            && expert_dp4a_supported(m.down_exps.qtype)
5167            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
5168            let dev = m.dev_exps.as_ref().unwrap();
5169            let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
5170                                                        &bits.per_expert_scale_d)?;
5171            if t == 1 {
5172                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
5173                let selv = sel_d.slice(0..n_used);
5174                let wv = w_d.slice(0..n_used);
5175                let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
5176                                                     n_embd, n_ff_exp, n_used, n_expert,
5177                                                     m.gate_exps.qtype, m.up_exps.qtype,
5178                                                     m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
5179                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
5180                let mut moe_out = e.uninit(n_embd)?;
5181                e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
5182                                       &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
5183                                       n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
5184                return Ok(moe_out);
5185            }
5186            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
5187            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
5188            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
5189            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
5190            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
5191            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
5192            let act = if csr {
5193                e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
5194                                               n_embd, n_ff_exp, n_used, n_expert,
5195                                               m.gate_exps.qtype, m.up_exps.qtype,
5196                                               m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5197            } else {
5198                e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
5199                                                n_embd, n_ff_exp, n_used, n_expert,
5200                                                m.gate_exps.qtype, m.up_exps.qtype,
5201                                                m.gate_exps.row_bytes, m.up_exps.row_bytes)?
5202            };
5203            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
5204            let mut moe_out = e.uninit(t * n_embd)?;
5205            e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
5206                                          n_ff_exp, n_embd, n_used, n_expert,
5207                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
5208            return Ok(moe_out);
5209        }
5210
5211        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
5212        for (i, &sx) in sel_all.iter().enumerate() {
5213            w_all[i] *= bits.per_expert_scale[sx as usize];
5214        }
5215
5216        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
5217        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
5218        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
5219        if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5220            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
5221            && expert_dp4a_supported(m.down_exps.qtype)
5222            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
5223            let dev = m.dev_exps.as_ref().unwrap();
5224            let n_pairs = t * n_used;
5225            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
5226            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
5227            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
5228            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
5229            let pt = e.htod_i32(&pair_tok)?;
5230            let pw = e.htod(&w_all)?;
5231            let toff = e.htod_i32(&tok_off)?;
5232            let tids = e.htod_i32(&tok_ids)?;
5233            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
5234            for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
5235            let mut ex_ids: Vec<i32> = Vec::new();
5236            let mut ex_off: Vec<i32> = vec![0];
5237            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
5238            for (ex, list) in by_ex.iter().enumerate() {
5239                if list.is_empty() { continue; }
5240                ex_ids.push(ex as i32);
5241                ex_pairs.extend_from_slice(list);
5242                ex_off.push(ex_pairs.len() as i32);
5243            }
5244            let n_active = ex_ids.len();
5245            let exi = e.htod_i32(&ex_ids)?;
5246            let exo = e.htod_i32(&ex_off)?;
5247            let exp_d = e.htod_i32(&ex_pairs)?;
5248            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
5249            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
5250            // end-to-end (gelu is elementwise), one row permute before the scatter. The
5251            // ragged down k (704) needs no padding here — cublas takes any k.
5252            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
5253            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
5254            // Hopper default — see moe_f16g_gemma_on.
5255            if crate::moe_f16g_gemma_on()
5256                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
5257                && f16g_proj_ok(m.up_exps.qtype, n_embd)
5258                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
5259                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
5260                let csr_tok_d = e.htod_i32(&csr_tok)?;
5261                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
5262                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
5263                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
5264                                              m.gate_exps.qtype, m.gate_exps.row_bytes)?;
5265                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
5266                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
5267                                              m.up_exps.qtype, m.up_exps.row_bytes)?;
5268                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
5269                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
5270                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
5271                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
5272                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
5273                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
5274                let mut moe_out = e.uninit(t * n_embd)?;
5275                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5276                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
5277                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
5278                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
5279                    eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
5280                              scan(&yd), scan(&mo));
5281                }
5282                return Ok(moe_out);
5283            }
5284            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
5285            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
5286            let mma = n_embd % 256 == 0
5287                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
5288            let (gate, up) = if mma {
5289                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
5290                (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
5291                                  n_embd, n_ff_exp, n_active, n_pairs, t,
5292                                  m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5293                 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
5294                                  n_embd, n_ff_exp, n_active, n_pairs, t,
5295                                  m.up_exps.qtype, m.up_exps.row_bytes)?)
5296            } else {
5297                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
5298                (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
5299                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
5300                                           m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5301                 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
5302                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
5303                                           m.up_exps.qtype, m.up_exps.row_bytes)?)
5304            };
5305            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
5306            let pself = e.htod_i32(&pair_self)?;
5307            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
5308            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
5309            // to the 256-val superblock (768) while the act quantizer's zero padding
5310            // makes every padded-k product exactly zero (weight overread bytes multiply
5311            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
5312            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
5313            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
5314            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
5315            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
5316            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
5317            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
5318            let y_down = if mma {
5319                let in_pad = n_ff_exp.div_ceil(256) * 256;
5320                let a_scr = if crate::moe_fuse_actq_on() {
5321                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
5322                } else {
5323                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5324                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
5325                };
5326                e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
5327                                 in_pad, n_embd, n_active, n_pairs, n_pairs,
5328                                 m.down_exps.qtype, m.down_exps.row_bytes)?
5329            } else {
5330                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
5331                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
5332                e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
5333                                          n_ff_exp, n_embd, n_expert, n_active, n_pairs,
5334                                          m.down_exps.qtype, m.down_exps.row_bytes)?
5335            };
5336            let mut moe_out = e.uninit(t * n_embd)?;
5337            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
5338            return Ok(moe_out);
5339        }
5340
5341        let g_len = m.gate_exps.expert_stride;
5342        let u_len = m.up_exps.expert_stride;
5343        let d_len = m.down_exps.expert_stride;
5344        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
5345        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
5346        // the spill fallback.
5347        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
5348        let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
5349            (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
5350        };
5351        let mut moe_out = e.zeros(t * n_embd)?;
5352        for tok in 0..t {
5353            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
5354            let w = &w_all[tok * n_used..(tok + 1) * n_used];
5355            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
5356            for (j, &ex) in sel.iter().enumerate() {
5357                let ex = ex as usize;
5358                let gate = match dev {
5359                    Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
5360                        m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
5361                    None => {
5362                        let sg = sg.as_mut().unwrap();
5363                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
5364                        e.qmatvec_view(sg, 0..g_len, &zt, 1,
5365                            m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
5366                    }
5367                };
5368                let up = match dev {
5369                    Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
5370                        m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
5371                    None => {
5372                        let su = su.as_mut().unwrap();
5373                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
5374                        e.qmatvec_view(su, 0..u_len, &zt, 1,
5375                            m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
5376                    }
5377                };
5378                let mut act = e.uninit(n_ff_exp)?;
5379                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
5380                let actv = act.slice(0..n_ff_exp);
5381                let y = match dev {
5382                    Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
5383                        m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
5384                    None => {
5385                        let sd = sd.as_mut().unwrap();
5386                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
5387                        e.qmatvec_view(sd, 0..d_len, &actv, 1,
5388                            m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
5389                    }
5390                };
5391                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
5392                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
5393            }
5394        }
5395        Ok(moe_out)
5396    }
5397
5398    /// One gemma4 trunk layer (R8): x -> x_next.
5399    fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
5400                    x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
5401                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5402        let n_embd = self.cfg.n_embd as usize;
5403        let eps = self.cfg.rms_eps;
5404
5405        let mut h = e.zeros(t * n_embd)?;
5406        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5407        let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5408        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
5409        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
5410        let mut cur = e.zeros(t * n_embd)?;
5411        e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
5412        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
5413    }
5414
5415    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
5416    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
5417    /// layer scale — shared verbatim by the prefill, decode and verify paths.
5418    fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5419                             cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
5420                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5421        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
5422    }
5423
5424    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
5425    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
5426    fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5427                               cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5428                               next_norm: Option<&CudaSlice<f32>>)
5429                               -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
5430        let n_embd = self.cfg.n_embd as usize;
5431        let bits = layer.gemma4.as_ref().unwrap();
5432        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
5433        let mut xn = e.uninit(t * n_embd)?;
5434        match next_norm {
5435            Some(w) => {
5436                let mut hn = e.uninit(t * n_embd)?;
5437                e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
5438                                     n_embd, t, self.cfg.rms_eps)?;
5439                Ok((xn, Some(hn)))
5440            }
5441            None => {
5442                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
5443                Ok((xn, None))
5444            }
5445        }
5446    }
5447
5448    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
5449    /// norm — returns (sn, attn_out) for the closing add+scale variants.
5450    fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5451                              cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
5452                              -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5453        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
5454    }
5455
5456    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
5457    /// means `cur` is the RAW attention output and the dense entry runs
5458    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
5459    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
5460    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
5461    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
5462    fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5463                                 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5464                                 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
5465                                 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5466        let n_embd = self.cfg.n_embd as usize;
5467        let eps = self.cfg.rms_eps;
5468        let bits = layer.gemma4.as_ref().unwrap();
5469
5470        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
5471        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
5472        let Some(mbits) = bits.moe_bits.as_ref() else {
5473            let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
5474            else { panic!("gemma4 dense layer without Dense ffn") };
5475            let mut attn_out = e.uninit(t * n_embd)?;
5476            let mut zsh = e.uninit(t * n_embd)?;
5477            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
5478            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
5479            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5480            match pre_norm {
5481                Some(wa) if t == 1 => {
5482                    zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
5483                                                            bits.ffn_norm.float_data(),
5484                                                            &mut attn_out, &mut zsh,
5485                                                            n_embd, t, eps)?);
5486                }
5487                Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
5488                                                   &mut attn_out, &mut zsh, n_embd, t, eps)?,
5489                None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
5490                                       &mut zsh, n_embd, t, eps)?,
5491            }
5492            let n_ff = ffn_gate.out_features();
5493            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
5494            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
5495            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
5496            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
5497            // rescue segment C — the megakernel front is closed for the dense tail.
5498            let (gate, up) = if t == 1 {
5499                let (zq, zd) = match zpair {
5500                    Some(p) => p,
5501                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
5502                };
5503                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
5504                    Some(p) => p,
5505                    None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
5506                             e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
5507                }
5508            } else {
5509                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
5510                // launch for the verify's gate+up — the up segment's blocks fill SMs as
5511                // the gate segment drains (the launch-tail mechanism behind the b-tier
5512                // plateau; first positive after six falsified in-kernel variants).
5513                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5514                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
5515                let fused = if f2b {
5516                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
5517                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
5518                } else { None };
5519                match fused {
5520                    Some(p) => p,
5521                    None => {
5522                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
5523                        e.mmq_act_begin();
5524                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
5525                    }
5526                }
5527            };
5528            let mut act = e.uninit(t * n_ff)?;
5529            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
5530            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
5531            let f0 = if e.uses_q8_1_fast(ffn_down) {
5532                let upv = e.view(&up, t * n_ff);
5533                let up_all = upv.slice(0..t * n_ff);
5534                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
5535                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5536            } else {
5537                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5538                e.matmul(ffn_down, &act, t)?
5539            };
5540            if defer_post_norm { return Ok((f0, attn_out)); }
5541            let mut sn = e.uninit(t * n_embd)?;
5542            e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
5543            return Ok((sn, attn_out));
5544        };
5545
5546        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
5547        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
5548        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
5549        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
5550        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
5551        let mut attn_out = e.uninit(t * n_embd)?;
5552        let mut router_in = e.uninit(t * n_embd)?;
5553        let fast_moe = match &layer.ffn {
5554            crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5555                && expert_dp4a_supported(m.gate_exps.qtype)
5556                && expert_dp4a_supported(m.up_exps.qtype)
5557                && expert_dp4a_supported(m.down_exps.qtype)
5558                && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
5559            _ => false,
5560        };
5561        let q8z = t < PRIME_MIN_T && fast_moe;
5562        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
5563            let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
5564                                               &mbits.router_scale_pre,
5565                                               mbits.pre_ffw_norm_2.float_data(),
5566                                               &mut attn_out, &mut router_in, n_embd, t, eps)?;
5567            (None, Some(z0), Some(m2))
5568        } else {
5569            let mut zsh = e.uninit(t * n_embd)?;
5570            let mut moe_in = e.uninit(t * n_embd)?;
5571            e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
5572                            mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
5573                            &mut router_in, &mut moe_in, n_embd, t, eps)?;
5574            (Some((zsh, moe_in)), None, None)
5575        };
5576        let attn_out2 = attn_out;
5577        #[allow(unused_variables)]
5578        let attn_out = &attn_out2;
5579        let n_ff = mbits.shared_gate.out_features();
5580        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
5581            if t == 1 {
5582                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
5583                    Some(p) => p,
5584                    None => {
5585                        let h0 = e.zeros(0)?;
5586                        (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
5587                         e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
5588                    }
5589                }
5590            } else {
5591                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
5592                let h0 = e.zeros(0)?;
5593                (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
5594                 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
5595            }
5596        } else {
5597            let (zsh, _) = zsh_f32.as_ref().unwrap();
5598            (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
5599        };
5600        let mut act = e.uninit(t * n_ff)?;
5601        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5602        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
5603        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
5604        let moe0 = match (&moe_q8, &zsh_f32) {
5605            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
5606            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
5607            _ => unreachable!(),
5608        };
5609        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
5610        let mut mlp = e.uninit(t * n_embd)?;
5611        let mut moe = e.uninit(t * n_embd)?;
5612        e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
5613                     mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
5614
5615        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
5616        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
5617        let mut sum = e.uninit(t * n_embd)?;
5618        let mut sn = e.uninit(t * n_embd)?;
5619        e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
5620                       n_embd, t, eps)?;
5621        Ok((sn, attn_out2))
5622    }
5623
5624    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
5625    fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5626                                cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5627                                next_norm: Option<&CudaSlice<f32>>)
5628                                -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
5629        let n_embd = self.cfg.n_embd as usize;
5630        let bits = layer.gemma4.as_ref().unwrap();
5631        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
5632        let mut xn = e.uninit(t * n_embd)?;
5633        match next_norm {
5634            Some(w) => {
5635                let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
5636                                                     n_embd, t, self.cfg.rms_eps)?;
5637                Ok((xn, Some(pair)))
5638            }
5639            None => {
5640                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
5641                Ok((xn, None))
5642            }
5643        }
5644    }
5645
5646    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
5647    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
5648    fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
5649                      -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5650        // E4B routes to its own forward regardless of the caller's entry point (forward /
5651        // forward_last / prime paths all funnel here for gemma4).
5652        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
5653        let n_embd = self.cfg.n_embd as usize;
5654        let t = tokens.len();
5655        let pos: Vec<i32> = (0..t as i32).collect();
5656        let pos_d = e.htod_i32(&pos)?;
5657
5658        let mut x = self.embed(e, tokens)?;
5659        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5660        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
5661        // the bring-up bisect vs llama-eval-callback node stats.
5662        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
5663        let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
5664            let h = e.dtoh(x)?;
5665            let bad = h.iter().filter(|v| !v.is_finite()).count();
5666            let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
5667            eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
5668            Ok(())
5669        };
5670        if probe { stat(e, &x, "embed")?; }
5671        for (il, layer) in self.layers.iter().enumerate() {
5672            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
5673            if probe { stat(e, &x, &format!("L{il}"))?; }
5674        }
5675        let mut hn = e.zeros(t * n_embd)?;
5676        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
5677        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5678        let n_vocab = self.output.out_features();
5679        let logits = if last_only {
5680            let hv = e.view(&hn, t * n_embd);
5681            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
5682            let mut hlast = e.zeros(n_embd)?;
5683            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
5684            let mut ld = e.matmul(&self.output, &hlast, 1)?;
5685            e.softcap(&mut ld, cap, n_vocab)?;
5686            self.gemma4_suppress(e, &mut ld, 1)?;
5687            e.dtoh(&ld)?
5688        } else {
5689            let mut ld = e.matmul(&self.output, &hn, t)?;
5690            e.softcap(&mut ld, cap, t * n_vocab)?;
5691            self.gemma4_suppress(e, &mut ld, t)?;
5692            e.dtoh(&ld)?
5693        };
5694        Ok(logits)
5695    }
5696
5697    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
5698    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
5699    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
5700    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
5701    pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
5702                               -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5703        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): a served gemma4 prompt longer
5704        // than the worker's prefill tick used to chunk here, and chunk 2 (pos > 0) killed the
5705        // whole worker process on this line. The worker now primes gemma4 monolithically and
5706        // routes continuation suffixes tokenwise; this is the per-request backstop.
5707        if cache.pos != 0 {
5708            return Err("gemma4 prime v0 is fresh-prompt only (no continuation/chunked prime) \
5709                        — prime the full prompt in one call or decode tokenwise".into());
5710        }
5711        let n_embd = self.cfg.n_embd as usize;
5712        let eps = self.cfg.rms_eps;
5713        let t = tokens.len();
5714        let pos: Vec<i32> = (0..t as i32).collect();
5715        let pos_d = e.htod_i32(&pos)?;
5716        let mut x = self.embed(e, tokens)?;
5717        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5718        for (il, layer) in self.layers.iter().enumerate() {
5719            let mut h = e.zeros(t * n_embd)?;
5720            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5721            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
5722            let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
5723            let mut cur = e.zeros(t * n_embd)?;
5724            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
5725            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
5726            self.dflash_tap(e, cache, il, &x, t)?;
5727        }
5728        cache.pos += t;
5729        let hiddens = e.clone_dtod(&x)?;
5730        let xv = e.view(&x, t * n_embd);
5731        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
5732        let mut h_seed = e.zeros(n_embd)?;
5733        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
5734        let mut hn = e.uninit(n_embd)?;
5735        e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5736        let mut ld = e.matmul(&self.output, &hn, 1)?;
5737        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5738        e.softcap(&mut ld, cap, self.output.out_features())?;
5739        self.gemma4_suppress(e, &mut ld, 1)?;
5740        let logits = e.dtoh(&ld)?;
5741        Ok((logits, h_seed, hiddens))
5742    }
5743
5744    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
5745    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
5746    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
5747    /// fused norm emits q8 directly — the f32 h never materializes).
5748    fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5749                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
5750                          pos_d: &CudaSlice<i32>, cache: &mut Cache)
5751                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5752        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
5753        let eps = self.cfg.rms_eps;
5754        let aux = self.gemma4_aux.as_ref().unwrap();
5755        let (hq, hdq) = (hq, hdq);
5756        let h0 = e.zeros(0)?;
5757        let h = &h0;
5758        let (q0, k0, v0) = if swa {
5759            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5760                Some(t3) => t3,
5761                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5762                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5763                         e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
5764            }
5765        } else {
5766            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
5767                Some(p) => p,
5768                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5769                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
5770            };
5771            let v0 = e.clone_dtod(&k0)?;
5772            (q0, k0, v0)
5773        };
5774        let mut q = e.uninit(nh * hd)?;
5775        let mut k = e.uninit(nkv * hd)?;
5776        let mut v = e.uninit(nkv * hd)?;
5777        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
5778        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
5779        let ff = if swa { None } else {
5780            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5781        };
5782        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5783                            &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
5784                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
5785        let kvl = cache.kv[il].as_mut().unwrap();
5786        e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
5787                              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()))?;
5788        kvl.len += 1;
5789        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
5790        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
5791        // positional). Globals attend the full history.
5792        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5793        let mut attn = e.uninit(nh * hd)?;
5794        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
5795        if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
5796            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5797            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5798            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5799            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
5800            let base = kvl.len as i32;
5801            e.i32_set_k(&mut kvl.len_d, base)?;
5802            e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
5803                             kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
5804                             false, None)?;
5805            return Ok(e.matmul(&fa.wo, &attn, 1)?);
5806        }
5807        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
5808        if swa && kvl.len > win && hd == 256
5809            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5810            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5811            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5812            let base = kvl.len as i32;
5813            e.i32_set_k(&mut kvl.len_d, base)?;
5814            e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
5815                               win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
5816            return Ok(e.matmul(&fa.wo, &attn, 1)?);
5817        }
5818        let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
5819        let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
5820                                     (off_tok + t_kv) * kvl.k_tok_bytes);
5821        let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
5822                                     (off_tok + t_kv) * kvl.v_tok_bytes);
5823        e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
5824                    kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
5825        Ok(e.matmul(&fa.wo, &attn, 1)?)
5826    }
5827
5828    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
5829    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
5830    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
5831    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
5832    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
5833    /// in-graph; the driver gates).
5834    #[allow(clippy::too_many_arguments)]
5835    pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
5836                                 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5837                                 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5838                                 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
5839                                 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
5840        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
5841        self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
5842                                        n_vocab, cap_bucket_max, &mut tok_out)?;
5843        Ok(tok_out)
5844    }
5845
5846    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
5847    /// every replay; pass `token_d` itself for the self-feeding graph loop).
5848    #[allow(clippy::too_many_arguments)]
5849    pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
5850                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5851                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5852                                      n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5853                                      tok_out: &mut CudaSlice<u32>)
5854                                      -> Result<(), Box<dyn std::error::Error>> {
5855        let n_embd = self.cfg.n_embd as usize;
5856        let eps = self.cfg.rms_eps;
5857        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
5858        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
5859        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5860        let n_layers = self.layers.len();
5861        for (il, layer) in self.layers.iter().enumerate() {
5862            let (hq, hdq) = match h_carry.take() {
5863                Some(p) => p,
5864                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
5865            };
5866            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5867            let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
5868            let mut cur = e.uninit(n_embd)?;
5869            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
5870            let next_norm = if il + 1 < n_layers {
5871                Some(self.layers[il + 1].attn_norm.float_data())
5872            } else { None };
5873            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
5874            x = xn;
5875            h_carry = hn;
5876        }
5877        let mut hn = e.uninit(n_embd)?;
5878        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5879        let mut logits = e.matmul(&self.output, &hn, 1)?;
5880        self.gemma4_suppress(e, &mut logits, 1)?;   // cap skipped (monotonic); the mask is not
5881        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
5882        e.inc_seqlen(pos_d)?;
5883        if cap_bucket_max.is_none() { cache.pos += 1; }
5884        Ok(())
5885    }
5886
5887    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
5888    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
5889    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
5890    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
5891
5892    /// Build the slot set (call OUTSIDE any capture).
5893    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
5894        let n_embd = self.cfg.n_embd as usize;
5895        let n_vocab = self.output.out_features();
5896        let n_layers = self.layers.len();
5897        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
5898        for il in 0..n_layers {
5899            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
5900            qmax = qmax.max(nh * hd);
5901            kvmax = kvmax.max(nkv * hd);
5902            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
5903                ffmax = ffmax.max(ffn_gate.out_features());
5904            }
5905        }
5906        Ok(G4DcSlots {
5907            x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
5908            hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
5909            q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
5910            q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
5911            attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
5912            attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
5913            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
5914            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
5915            zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
5916            gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
5917            act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
5918            f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
5919            hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
5920        })
5921    }
5922
5923    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
5924    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
5925    fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
5926                         aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
5927                         -> Result<(), Box<dyn std::error::Error>> {
5928        use crate::model::GpuTensor;
5929        let (bytes, qtype, row_bytes, scale, rp) = match w {
5930            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
5931                (bytes, *qtype, *row_bytes, *scale, *rp),
5932            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
5933        };
5934        let (mbytes, mrp) = match w {
5935            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5936            _ => (bytes, rp),
5937        };
5938        e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
5939                            qtype, row_bytes, scale, mrp, y)
5940    }
5941
5942    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
5943    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
5944    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
5945    #[allow(clippy::too_many_arguments)]
5946    pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
5947                                         pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5948                                         embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5949                                         n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5950                                         sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
5951                                         ring: Option<(&mut CudaSlice<u32>, usize)>)
5952                                         -> Result<(), Box<dyn std::error::Error>> {
5953        let n_embd = self.cfg.n_embd as usize;
5954        let eps = self.cfg.rms_eps;
5955        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
5956        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
5957        let n_layers = self.layers.len();
5958        let mut has_carry = false;
5959        for il in 0..n_layers {
5960            if !has_carry {
5961                e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
5962                                     eps, &mut sl.hq, &mut sl.hd_)?;
5963            }
5964            has_carry = true;
5965            let layer = &self.layers[il];
5966            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5967            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
5968            e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
5969            let next_norm = if il + 1 < n_layers {
5970                Some(self.layers[il + 1].attn_norm.float_data())
5971            } else { None };
5972            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
5973            std::mem::swap(&mut sl.x, &mut sl.xn);
5974        }
5975        e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
5976        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
5977        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
5978        {
5979            let (zq, zd) = (&sl.zq, &sl.zd);
5980            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
5981            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
5982            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
5983        }
5984        self.gemma4_suppress(e, &mut sl.logits, 1)?;
5985        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
5986        if let Some((ring, base)) = ring {
5987            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
5988            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
5989            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
5990            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
5991        }
5992        e.inc_seqlen(pos_d)?;
5993        if cap_bucket_max.is_none() { cache.pos += 1; }
5994        Ok(())
5995    }
5996
5997    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
5998    #[allow(clippy::too_many_arguments)]
5999    fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
6000                                     il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
6001                                     cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
6002                                     -> Result<(), Box<dyn std::error::Error>> {
6003        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6004        let eps = self.cfg.rms_eps;
6005        let aux = self.gemma4_aux.as_ref().unwrap();
6006        {
6007            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
6008            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
6009            if swa {
6010                if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
6011                                            &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
6012                    return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
6013                }
6014            } else {
6015                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
6016                    return Err("slotted step: fused2 unavailable".into());
6017                }
6018                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
6019                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
6020            }
6021        }
6022        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
6023        // kernel-for-kernel (graph stream-identity gate).
6024        let ff = if swa { None } else {
6025            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6026        };
6027        let kvl = cache.kv[il].as_mut().unwrap();
6028        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6029        if crate::Engine::qkv_append_on() {
6030            // append fold (2026-07-23): mirrors dc_into.
6031            e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
6032                fa.k_norm.float_data(), &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
6033                pos_d, nh, nkv, base, 1.0, ff, eps,
6034                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6035        } else {
6036            e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6037                                &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
6038                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
6039            e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
6040                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
6041                                     kv_fp8)?;
6042        }
6043        e.inc_seqlen(&mut kvl.len_d)?;
6044        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
6045        let k_view = e.view_u8(&kvl.k, kvl.k.len());
6046        let v_view = e.view_u8(&kvl.v, kvl.v.len());
6047        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
6048        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6049        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
6050        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
6051        // the dc_into arm branch-for-branch (stream gate).
6052        let mut fa_q8 = false;
6053        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
6054            e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
6055                             1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6056                             Some((&kvl.len_d, -1)), false, false,
6057                             Some((&mut sl.zq, &mut sl.zd)))?;
6058            fa_q8 = true;
6059        } else if swa && b_swa > win && hd == 256 && rows_on {
6060            e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
6061                               &kvl.len_d, -1, 1, scale, win,
6062                               kvl.k_tok_bytes, kvl.v_tok_bytes,
6063                               Some((&mut sl.zq, &mut sl.zd)))?;
6064            fa_q8 = true;
6065        } else {
6066            let b = if swa { b_swa } else { b_glob };
6067            e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
6068                           scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6069                           swa && crate::Engine::wkv_on())?;
6070        }
6071        if !fa_q8 {
6072            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
6073            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
6074        }
6075        {
6076            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
6077            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
6078            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
6079        }
6080        Ok(())
6081    }
6082
6083    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
6084    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
6085    fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
6086                                 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
6087                                 -> Result<(), Box<dyn std::error::Error>> {
6088        let n_embd = self.cfg.n_embd as usize;
6089        let eps = self.cfg.rms_eps;
6090        let bits = layer.gemma4.as_ref().unwrap();
6091        let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
6092        else { return Err("slotted tail: dense ffn only".into()) };
6093        e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
6094                       &mut sl.zsh, n_embd, 1, eps)?;
6095        let n_ff = ffn_gate.out_features();
6096        {
6097            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
6098            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
6099        }
6100        {
6101            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
6102            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
6103            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
6104                return Err("slotted tail: ffn fused2 unavailable".into());
6105            }
6106        }
6107        debug_assert!(e.uses_q8_1_fast(ffn_down));
6108        {
6109            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
6110            let upv = e.view(upr, n_ff);
6111            let up_all = upv.slice(0..n_ff);
6112            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
6113            e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
6114                                      &mut sl.actq, &mut sl.actd)?;
6115        }
6116        {
6117            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
6118            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
6119            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
6120        }
6121        e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
6122        match next_norm {
6123            Some(w) => {
6124                e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
6125                                               &mut sl.xn, n_embd, 1, eps,
6126                                               &mut sl.hq, &mut sl.hd_)?;
6127            }
6128            None => {
6129                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
6130            }
6131        }
6132        Ok(())
6133    }
6134
6135    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
6136    #[allow(clippy::too_many_arguments)]
6137    fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6138                             hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6139                             pos_d: &CudaSlice<i32>, cache: &mut Cache,
6140                             cap_bucket_max: Option<(usize, usize)>)
6141                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6142        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6143        let eps = self.cfg.rms_eps;
6144        let aux = self.gemma4_aux.as_ref().unwrap();
6145        let (q0, k0, v0) = if swa {
6146            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
6147                Some(t3) => t3,
6148                None => {
6149                    let h0 = e.zeros(0)?;
6150                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
6151                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
6152                     e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
6153                }
6154            }
6155        } else {
6156            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
6157                Some(p) => p,
6158                None => {
6159                    let h0 = e.zeros(0)?;
6160                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
6161                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
6162                }
6163            };
6164            let v0 = e.clone_dtod(&k0)?;
6165            (q0, k0, v0)
6166        };
6167        let mut q = e.uninit(nh * hd)?;
6168        let mut k = e.uninit(nkv * hd)?;
6169        let mut v = e.uninit(nkv * hd)?;
6170        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
6171        let ff = if swa { None } else {
6172            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6173        };
6174        let kvl = cache.kv[il].as_mut().unwrap();
6175        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6176        if crate::Engine::qkv_append_on() {
6177            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
6178            e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
6179                fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
6180                pos_d, nh, nkv, base, 1.0, ff, eps,
6181                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6182        } else {
6183            e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6184                                &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
6185                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
6186            e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
6187                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
6188        }
6189        e.inc_seqlen(&mut kvl.len_d)?;
6190        let mut attn = e.uninit(nh * hd)?;
6191        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
6192        // rides g4_matvec_m1_into instead of matmul's internal quantize.
6193        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6194        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
6195        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
6196        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
6197        // (gemma4_e4b_attn, +0.65% valid window).
6198        match cap_bucket_max {
6199            None => {
6200                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
6201                // decode (SWA layers attend the last `sliding_window` keys); the device
6202                // counters carry only the append slot + the graph seam.
6203                kvl.len += 1;
6204                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6205                if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
6206                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6207                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
6208                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
6209                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
6210                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
6211                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6212                    e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
6213                                     scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6214                                     Some((&kvl.len_d, -1)), false, false,
6215                                     Some((&mut aq8, &mut ad8)))?;
6216                    fa_q8 = Some((aq8, ad8));
6217                } else if swa && kvl.len > win && hd == 256
6218                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6219                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
6220                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
6221                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
6222                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6223                    e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
6224                                       1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
6225                                       Some((&mut aq8, &mut ad8)))?;
6226                    fa_q8 = Some((aq8, ad8));
6227                } else {
6228                    let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
6229                                          else { (0, kvl.len) };
6230                    let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6231                                                 (off_tok + t_kv) * kvl.k_tok_bytes);
6232                    let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6233                                                 (off_tok + t_kv) * kvl.v_tok_bytes);
6234                    e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
6235                                kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
6236                }
6237            }
6238            Some((b_swa, b_glob)) => {
6239                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
6240                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
6241                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
6242                // the RUNG max for the rows family (kernels derive per-replay splits from
6243                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
6244                let k_view = e.view_u8(&kvl.k, kvl.k.len());
6245                let v_view = e.view_u8(&kvl.v, kvl.v.len());
6246                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
6247                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6248                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
6249                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6250                    e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
6251                                     1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6252                                     Some((&kvl.len_d, -1)), false, false,
6253                                     Some((&mut aq8, &mut ad8)))?;
6254                    fa_q8 = Some((aq8, ad8));
6255                } else if swa && b_swa > win && hd == 256 && rows_on {
6256                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
6257                    e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6258                                       &kvl.len_d, -1, 1, scale, win,
6259                                       kvl.k_tok_bytes, kvl.v_tok_bytes,
6260                                       Some((&mut aq8, &mut ad8)))?;
6261                    fa_q8 = Some((aq8, ad8));
6262                } else {
6263                    let b = if swa { b_swa } else { b_glob };
6264                    e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
6265                                   scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6266                                   swa && crate::Engine::wkv_on())?;
6267                }
6268            }
6269        }
6270        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
6271        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
6272        if let Some((aq8, ad8)) = fa_q8 {
6273            let mut y = e.uninit(fa.wo.out_features())?;
6274            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
6275            return Ok(y);
6276        }
6277        Ok(e.matmul(&fa.wo, &attn, 1)?)
6278    }
6279
6280    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
6281    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
6282    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
6283    /// views in-graph); caller gates and falls back to the dc-eager loop.
6284    pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
6285                                 cache: &mut Cache, max_new: usize, eos: &[u32],
6286                                 mut on_token: impl FnMut(u32) -> bool)
6287                                 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
6288        if self.is_gemma4_e4b() {
6289            return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
6290        }
6291        use crate::decode::StopReason;
6292        let n_vocab = self.output.out_features();
6293        let n_embd = self.cfg.n_embd as usize;
6294        let embd_gpu = self.embd_gpu.get_or_init(|| {
6295            e.upload_u8(&self.embd.raw).expect("embed table upload")
6296        });
6297        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6298        for kvl in cache.kv.iter_mut().flatten() {
6299            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6300        }
6301        let mut token_d = e.stream().clone_htod(&[first_token])?;
6302        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
6303        let g4 = self.cfg.gemma4.as_ref().unwrap();
6304        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
6305        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
6306        let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
6307            .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
6308        let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
6309            .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
6310        let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
6311                                                  (cudarc::driver::CudaGraph,
6312                                                   Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
6313        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
6314        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
6315        let mut slots = self.g4_dc_slots(e)?;
6316        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
6317        // baked at the door entry (the modulo keeps every capture valid indefinitely).
6318        const RING: usize = 64;
6319        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
6320        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
6321        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
6322        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
6323        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
6324        const DRAIN: usize = 1;
6325        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
6326        let ring_base = prompt_pos;
6327        let mut out = Vec::with_capacity(max_new);
6328        let mut reason = StopReason::MaxNew;
6329        let mut next = first_token;
6330        let mut captures = 0usize;
6331        for _ in 0..max_new {
6332            out.push(next);
6333            if eos.contains(&next) { reason = StopReason::Eos; break; }
6334            if !on_token(next) { reason = StopReason::Callback; break; }
6335            let t_kv = cache.pos + 1;
6336            // Bucket key per ARM (graph arc step 3):
6337            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
6338            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
6339            //    the component collapses to a single marker).
6340            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
6341            //    at/above it — the kernel derives splits from len_d per replay, so buckets
6342            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
6343            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6344            let f512 = crate::fa512_min_tkv();
6345            let key_s = if t_kv > win { (true, usize::MAX) }
6346                        else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
6347            let (key_g, rung_end) = if t_kv >= f512 {
6348                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
6349                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
6350                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
6351                ((true, end), end)
6352            } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
6353            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
6354            if !graphs.contains_key(&key) {
6355                let bucket_max = (t_kv, rung_end);
6356                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
6357                let snap = cache.snapshot(e)?;
6358                let pos_save = e.dtoh_i32_one(&pos_d)?;
6359                let len_save: Vec<Option<i32>> = cache.kv.iter()
6360                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
6361                let tok_save = e.dtoh_u32_one(&token_d)?;
6362                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
6363                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
6364                // regression class, and this door's measured -8.8%. The keeper pins warmup
6365                // transients so the captured graph holds kernel nodes only.
6366                let graph = {
6367                    let tok_ref = &mut token_d;
6368                    let pos_ref = &mut pos_d;
6369                    let cache_ref = &mut *cache;
6370                    let slots_ref = &mut slots;
6371                    let ring_ref = &mut ring;
6372                    e.capture_graph_retained_flags(
6373                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
6374                        |e| {
6375                        // self-feeding: the argmax writes token_d itself.
6376                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
6377                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
6378                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
6379                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
6380                                                           cache_ref, n_vocab, Some(bucket_max),
6381                                                           sl, tok_ref, Some((rg, ring_base)))
6382                    })?
6383                };
6384                cache.rollback(e, &snap, 0)?;
6385                e.set_i32_one(&mut pos_d, pos_save)?;
6386                for (il, ls) in len_save.iter().enumerate() {
6387                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
6388                        e.set_i32_one(&mut kvl.len_d, *v)?;
6389                    }
6390                }
6391                e.set_u32_one(&mut token_d, tok_save)?;
6392                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
6393                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
6394                        eprintln!("[graph-census] {c:?}");
6395                    }
6396                }
6397                graphs.insert(key, graph);
6398                captures += 1;
6399            }
6400            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
6401            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
6402            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
6403            // the budget; capture warmups already emitted their tokens through the ring.
6404            let mut chunk = 1usize;
6405            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
6406                .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
6407            while chunk < drain_cap && out.len() + chunk < max_new {
6408                let t_next = cache.pos + 1 + chunk;
6409                let key_s2 = if t_next > win { (true, usize::MAX) }
6410                             else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
6411                let key_g2 = if t_next >= f512 {
6412                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
6413                } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
6414                if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
6415                chunk += 1;
6416            }
6417            let g = &graphs.get(&key).unwrap().0;
6418            for _ in 0..chunk { g.launch()?; }
6419            e.stream().synchronize()?;
6420            let ringh = e.dtoh_u32(&ring)?;
6421            for j in 0..chunk {
6422                let pos_j = cache.pos + j;
6423                let tok_j = ringh[(pos_j - ring_base) % RING];
6424                cache.pos += 0; // advanced below in one shot
6425                if j + 1 == chunk { next = tok_j; }
6426                else {
6427                    out.push(tok_j);
6428                    if eos.contains(&tok_j) || !on_token(tok_j) {
6429                        reason = if eos.contains(&tok_j) { StopReason::Eos }
6430                                 else { StopReason::Callback };
6431                        // roll device/host state back to the stop point.
6432                        let keep = cache.pos + j + 1;
6433                        e.set_i32_one(&mut pos_d, keep as i32)?;
6434                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
6435                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
6436                            kvl.len = keep;
6437                        }
6438                        cache.pos = keep;
6439                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
6440                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
6441                        }
6442                        return Ok((out, reason));
6443                    }
6444                }
6445            }
6446            cache.pos += chunk;
6447            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
6448        }
6449        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
6450            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
6451        }
6452        Ok((out, reason))
6453    }
6454
6455    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
6456    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
6457    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
6458    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
6459    /// logits (host) + advances cache.pos by t.
6460    pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
6461                                       cache: &mut Cache)
6462                                       -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6463        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
6464    }
6465
6466    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
6467    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
6468    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
6469    pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
6470                                          cache: &mut Cache)
6471                                          -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6472        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6473        let t = tokens.len();
6474        let n_vocab = self.output.out_features();
6475        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
6476        for i in 0..t {
6477            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
6478        }
6479        Ok((e.dtoh_u32(&toks)?, hn))
6480    }
6481
6482    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
6483    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
6484    pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6485                                              pos0: usize, cache: &mut Cache)
6486                                              -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6487        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
6488        let n_vocab = self.output.out_features();
6489        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6490        for i in 0..t {
6491            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6492        }
6493        Ok((vam, hn))
6494    }
6495
6496    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
6497    /// llama's h_nextn convention).
6498    pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
6499                                         cache: &mut Cache)
6500                                         -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6501        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6502        let t = tokens.len();
6503        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6504        e.softcap(&mut ld, cap, t * self.output.out_features())?;
6505        Ok((e.dtoh(&ld)?, hn))
6506    }
6507
6508    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
6509    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
6510    pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
6511                                        -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
6512        Ok(VerifyStreamScratch {
6513            pos_d: e.htod_i32(&vec![0i32; cap])?,
6514            row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
6515        })
6516    }
6517
6518    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
6519    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
6520    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
6521    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
6522    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
6523    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
6524    /// sync, exactly the turnaround the burst exists to remove.
6525    pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6526                                            ctr: &CudaSlice<i32>, hint: usize,
6527                                            cache: &mut Cache,
6528                                            scr: &mut VerifyStreamScratch)
6529                                            -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6530        let n_embd = self.cfg.n_embd as usize;
6531        let eps = self.cfg.rms_eps;
6532        assert!(t <= scr.row_ctrs.len() && t <= 64);
6533        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
6534        for i in 0..t {
6535            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
6536        }
6537        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
6538        let embd_gpu = self.embd_gpu.get_or_init(|| {
6539            e.upload_u8(&self.embd.raw).expect("embed table upload")
6540        });
6541        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6542        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
6543        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6544        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6545        let n_layers = self.layers.len();
6546        for (il, layer) in self.layers.iter().enumerate() {
6547            let (hq, hdq) = match h_carry.take() {
6548                Some(p) => p,
6549                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6550            };
6551            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6552            let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
6553                                                    hint, row_ctrs)?;
6554            let mut cur = e.uninit(t * n_embd)?;
6555            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6556            let next_norm = if il + 1 < n_layers {
6557                Some(self.layers[il + 1].attn_norm.float_data())
6558            } else { None };
6559            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6560            x = xn;
6561            h_carry = hn;
6562            self.dflash_tap(e, cache, il, &x, t)?;
6563        }
6564        let mut hn = e.uninit(t * n_embd)?;
6565        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6566        let ld = e.matmul(&self.output, &hn, t)?;
6567        let n_vocab = self.output.out_features();
6568        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6569        for i in 0..t {
6570            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6571        }
6572        Ok((vam, hn))
6573    }
6574
6575    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
6576    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
6577    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
6578    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
6579    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
6580    /// kernel later if it shows in the profile).
6581    fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
6582                  -> Result<(), Box<dyn std::error::Error>> {
6583        let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
6584        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
6585        let h = taps.hidden;
6586        let n_taps = taps.layer_ids.len();
6587        debug_assert_eq!(taps.t, t);
6588        let xv = e.view(x, t * h);
6589        for r in 0..t {
6590            let row = xv.slice(r * h..(r + 1) * h);
6591            e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
6592        }
6593        Ok(())
6594    }
6595
6596    fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
6597                           tok_dev: Option<&CudaSlice<u32>>)
6598                           -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6599        let n_embd = self.cfg.n_embd as usize;
6600        let eps = self.cfg.rms_eps;
6601        let t = tokens.len();
6602        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6603        let pos_d = e.htod_i32(&pos)?;
6604        let mut x = match tok_dev {
6605            Some(td) => {
6606                let embd_gpu = self.embd_gpu.get_or_init(|| {
6607                    e.upload_u8(&self.embd.raw).expect("embed table upload")
6608                });
6609                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6610                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
6611            }
6612            None => e.htod(&self.embd.gather(n_embd, tokens))?,
6613        };
6614        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6615        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6616        let n_layers = self.layers.len();
6617        for (il, layer) in self.layers.iter().enumerate() {
6618            let (hq, hdq) = match h_carry.take() {
6619                Some(p) => p,
6620                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6621            };
6622            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6623            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
6624            let mut cur = e.uninit(t * n_embd)?;
6625            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6626            let next_norm = if il + 1 < n_layers {
6627                Some(self.layers[il + 1].attn_norm.float_data())
6628            } else { None };
6629            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6630            x = xn;
6631            h_carry = hn;
6632            self.dflash_tap(e, cache, il, &x, t)?;
6633        }
6634        let mut hn = e.uninit(t * n_embd)?;
6635        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6636        let mut ld = e.matmul(&self.output, &hn, t)?;
6637        self.gemma4_suppress(e, &mut ld, t)?;   // before the per-row argmax consumers
6638        cache.pos += t;
6639        Ok((ld, hn))
6640    }
6641
6642    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
6643    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
6644    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
6645    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
6646    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
6647    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
6648    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
6649    #[allow(clippy::too_many_arguments)]
6650    fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6651                                 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6652                                 pos_d: &CudaSlice<i32>, t: usize,
6653                                 cache: &mut Cache, hint: usize,
6654                                 row_ctrs: &[CudaSlice<i32>])
6655                                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6656        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6657        let eps = self.cfg.rms_eps;
6658        let aux = self.gemma4_aux.as_ref().unwrap();
6659        let h0 = e.zeros(0)?;
6660        let h = &h0;
6661        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
6662        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
6663        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6664        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6665        let fused_qkv = if f2b {
6666            if swa {
6667                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6668                    .map(|(a, b, c)| (a, b, Some(c)))
6669            } else {
6670                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6671                    .map(|(a, b)| (a, b, None))
6672            }
6673        } else { None };
6674        let (q0, k0, v0) = match fused_qkv {
6675            Some((a, b, cv)) => {
6676                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6677                (a, b, v)
6678            }
6679            None => {
6680                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6681                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6682                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6683                         else { e.clone_dtod(&k0)? };
6684                (q0, k0, v0)
6685            }
6686        };
6687        let mut q = e.uninit(t * nh * hd)?;
6688        let mut k = e.uninit(t * nkv * hd)?;
6689        let mut v = e.uninit(t * nkv * hd)?;
6690        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
6691        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
6692        let ff = if swa { None } else {
6693            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6694        };
6695        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6696                            &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6697                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
6698        let kvl = cache.kv[il].as_mut().unwrap();
6699        // append at the DEVICE slot; the counter advances by t on-device.
6700        e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
6701                                      kvl.kv_dim_k, kvl.kv_dim_v,
6702                                      kvl.k_tok_bytes, kvl.v_tok_bytes,
6703                                      (!swa && crate::Engine::gkv_on())
6704                                          || (swa && crate::Engine::wkv_on()))?;
6705        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
6706        // the sole len writer after this round's attention (base stays = old len, plus = 0).
6707        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6708        let mut attn = e.uninit(t * nh * hd)?;
6709        let k_view = e.view_u8(&kvl.k, kvl.k.len());
6710        let v_view = e.view_u8(&kvl.v, kvl.v.len());
6711        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
6712        // and a stable window regime — the same rung/regime keys as the draft graph).
6713        if swa && hint + 1 >= win {
6714            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
6715            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
6716            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6717                               &kvl.len_d, 0, t, scale, win,
6718                               kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6719        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
6720            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
6721            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
6722            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
6723            // Burst entry gates the horizon onto one side of the crossover, so hint decides
6724            // for every row.
6725            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
6726            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
6727            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
6728            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
6729            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
6730            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
6731            // any bucket >= the live length is exact.
6732            let bucket = (hint + t + 2).next_power_of_two()
6733                .min(crate::fa512_min_tkv().saturating_sub(1));
6734            let qv = e.view(&q, t * nh * hd);
6735            for i in 0..t {
6736                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
6737                let mut q_one = e.uninit(nh * hd)?;
6738                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6739                let mut a_one = e.uninit(nh * hd)?;
6740                e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
6741                               &row_ctrs[i], bucket, scale,
6742                               kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
6743                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6744            }
6745        } else if hd == 512 {
6746            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
6747            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
6748            e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
6749                             kvl.k_tok_bytes, kvl.v_tok_bytes,
6750                             Some((&kvl.len_d, 0)), false, false, None)?;
6751        } else {
6752            // hd256 under-window: v4 device-len rows twin.
6753            e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6754                                &kvl.len_d, hint + t, t, scale,
6755                                kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6756                                swa && crate::Engine::wkv_on())?;
6757        }
6758        Ok(e.matmul(&fa.wo, &attn, t)?)
6759    }
6760
6761    fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6762                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6763                          pos_d: &CudaSlice<i32>, t: usize,
6764                          cache: &mut Cache)
6765                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6766        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6767        let eps = self.cfg.rms_eps;
6768        let aux = self.gemma4_aux.as_ref().unwrap();
6769        let n_embd = self.cfg.n_embd as usize;
6770        let _ = n_embd;
6771
6772        let h0 = e.zeros(0)?;
6773        let h = &h0;
6774        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
6775        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
6776        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6777        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6778        let fused_qkv = if f2b {
6779            if swa {
6780                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6781                    .map(|(a, b, c)| (a, b, Some(c)))
6782            } else {
6783                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6784                    .map(|(a, b)| (a, b, None))
6785            }
6786        } else { None };
6787        let (q0, k0, v0) = match fused_qkv {
6788            Some((a, b, cv)) => {
6789                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6790                (a, b, v)
6791            }
6792            None => {
6793                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6794                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6795                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6796                         else { e.clone_dtod(&k0)? };
6797                (q0, k0, v0)
6798            }
6799        };
6800        let mut q = e.uninit(t * nh * hd)?;
6801        let mut k = e.uninit(t * nkv * hd)?;
6802        let mut v = e.uninit(t * nkv * hd)?;
6803        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
6804        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
6805        let ff = if swa { None } else {
6806            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6807        };
6808        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6809                            &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6810                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
6811        let kvl = cache.kv[il].as_mut().unwrap();
6812        let base_len = kvl.len;
6813        e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
6814                                   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()))?;
6815        kvl.len += t;
6816        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6817        let mut attn = e.uninit(t * nh * hd)?;
6818        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
6819        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
6820        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
6821            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
6822            // decode rides the SAME symbol at t=1 (parity law).
6823            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
6824        if rows_ok && (!swa || base_len + t <= win) {
6825            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6826            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6827            if hd == 512 {
6828                // device-len twin: sync the counter to the verify base (async arg-store).
6829                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6830                e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
6831                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6832                                 Some((&kvl.len_d, 0)), false,
6833                                 swa && crate::Engine::wkv_on(), None)?;
6834            } else {
6835                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
6836                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
6837                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
6838                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6839                e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6840                                    &kvl.len_d, base_len + t, t, scale,
6841                                    kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6842                                    swa && crate::Engine::wkv_on())?;
6843            }
6844            return Ok(e.matmul(&fa.wo, &attn, t)?);
6845        }
6846        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
6847        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
6848        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
6849        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
6850        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
6851        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
6852        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
6853        if hd == 256 && swa && base_len + 1 >= win
6854            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6855            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6856            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6857            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6858            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
6859                               t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6860            return Ok(e.matmul(&fa.wo, &attn, t)?);
6861        }
6862        for i in 0..t {
6863            let avail = base_len + i + 1;
6864            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
6865            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6866                                         (off_tok + t_kv) * kvl.k_tok_bytes);
6867            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6868                                         (off_tok + t_kv) * kvl.v_tok_bytes);
6869            let qi = e.view(&q, t * nh * hd);
6870            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
6871            let mut q_one = e.uninit(nh * hd)?;
6872            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6873            let mut a_one = e.uninit(nh * hd)?;
6874            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
6875            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
6876            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
6877            if swa && avail > win && hd == 256
6878                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6879                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6880                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6881                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6882                e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
6883                                   1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6884            } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
6885                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6886                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6887                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6888                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6889                e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
6890                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6891                                 Some((&kvl.len_d, 0)), false, false, None)?;
6892            } else {
6893                e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
6894                            kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
6895            }
6896            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6897        }
6898        Ok(e.matmul(&fa.wo, &attn, t)?)
6899    }
6900
6901    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
6902    /// h_seed = pre-output_norm hidden). Advances cache.pos.
6903    pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
6904                                       -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6905        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
6906        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
6907        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
6908        // unsplit rather than guessing a fence.
6909        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
6910            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
6911        }
6912        if crate::pp::pp_cuts(self.layers.len()).is_some() {
6913            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
6914        }
6915        let n_embd = self.cfg.n_embd as usize;
6916        let eps = self.cfg.rms_eps;
6917        let pos_d = e.htod_i32(&[cache.pos as i32])?;
6918        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
6919        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
6920        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
6921        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
6922        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6923        let n_layers = self.layers.len();
6924        for (il, layer) in self.layers.iter().enumerate() {
6925            let (hq, hdq) = match h_carry.take() {
6926                Some(p) => p,
6927                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
6928            };
6929            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6930            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
6931            let mut cur = e.uninit(n_embd)?;
6932            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6933            let next_norm = if il + 1 < n_layers {
6934                Some(self.layers[il + 1].attn_norm.float_data())
6935            } else { None };
6936            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6937            x = xn;
6938            h_carry = hn;
6939        }
6940        let mut hn = e.uninit(n_embd)?;
6941        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
6942        let h_seed = e.clone_dtod(&x)?;
6943        let mut ld = e.matmul(&self.output, &hn, 1)?;
6944        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6945        e.softcap(&mut ld, cap, self.output.out_features())?;   // R4 on device (262k host tanh ~ms/step)
6946        self.gemma4_suppress(e, &mut ld, 1)?;
6947        let logits = e.dtoh(&ld)?;
6948        cache.pos += 1;
6949        Ok((logits, h_seed))
6950    }
6951
6952    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
6953    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
6954    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
6955    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
6956    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
6957    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
6958    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
6959    fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
6960                            pos_d: &CudaSlice<i32>, cache: &mut Cache)
6961                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6962        let n_embd = self.cfg.n_embd as usize;
6963        let eps = self.cfg.rms_eps;
6964        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6965        for il in lo..hi {
6966            let layer = &self.layers[il];
6967            let (hq, hdq) = match h_carry.take() {
6968                Some(p) => p,
6969                // range head: il == lo — norm against THIS layer's attn_norm.
6970                None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
6971            };
6972            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6973            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
6974            let mut cur = e.uninit(n_embd)?;
6975            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6976            let next_norm = if il + 1 < hi {
6977                Some(self.layers[il + 1].attn_norm.float_data())
6978            } else { None };
6979            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6980            x = xn;
6981            h_carry = hn;
6982        }
6983        Ok(x)
6984    }
6985
6986    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
6987    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
6988    /// boundary handoff — same choreography as the generic arm (decode.rs), same
6989    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
6990    /// stage 1 = layers [split, n) + output_norm + softcapped head.
6991    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
6992    fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
6993                                -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6994        if crate::pp::pp2_streams_off() {
6995            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
6996        }
6997        let rt = crate::pp::Pp2Rt::get(e)?;
6998        let e0 = rt.engine(0, e);
6999        let e1 = rt.engine(1, e);
7000        let n_embd = self.cfg.n_embd as usize;
7001        let eps = self.cfg.rms_eps;
7002
7003        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
7004        let (pos_d, slot) = {
7005            let _st0 = rt.enter(0);
7006            let pos_d = e0.htod_i32(&[cache.pos as i32])?;
7007            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
7008            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7009            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
7010            let slot = rt.tx(0, &x, n_embd)?;
7011            (pos_d, slot)
7012        };
7013
7014        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
7015        let _st1 = rt.enter(1);
7016        let x = rt.rx(0, slot, n_embd)?;
7017        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
7018
7019        let mut hn = e1.uninit(n_embd)?;
7020        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7021        let h_seed = e1.clone_dtod(&x)?;
7022        let mut ld = e1.matmul(&self.output, &hn, 1)?;
7023        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7024        e1.softcap(&mut ld, cap, self.output.out_features())?;
7025        self.gemma4_suppress(e1, &mut ld, 1)?;
7026        let logits = e1.dtoh(&ld)?;
7027        cache.pos += 1;
7028        Ok((logits, h_seed))
7029    }
7030
7031    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
7032    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
7033    fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
7034                                           split: usize)
7035                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7036        let n_embd = self.cfg.n_embd as usize;
7037        let eps = self.cfg.rms_eps;
7038        let pos_d = e.htod_i32(&[cache.pos as i32])?;
7039
7040        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
7041        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
7042        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7043        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
7044
7045        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
7046        let boundary_tx = e.clone_dtod(&x)?;
7047        let boundary_rx = e.clone_dtod(&boundary_tx)?;
7048
7049        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
7050        let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
7051
7052        let mut hn = e.uninit(n_embd)?;
7053        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
7054        let h_seed = e.clone_dtod(&x)?;
7055        let mut ld = e.matmul(&self.output, &hn, 1)?;
7056        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7057        e.softcap(&mut ld, cap, self.output.out_features())?;
7058        self.gemma4_suppress(e, &mut ld, 1)?;
7059        let logits = e.dtoh(&ld)?;
7060        cache.pos += 1;
7061        Ok((logits, h_seed))
7062    }
7063}
7064
7065// ============================ step35 (Step-3.7-Flash) ==================================
7066// Node-for-node vs llama.cpp src/models/step35.cpp:216-300. WHY THIS IS A DEDICATED MIXER
7067// FAMILY and not a few branches inside the generic `full_attn*` chain:
7068//
7069//   1. `n_head` IS PER LAYER (`step35.attention.head_count` is an ARRAY: 64 on full-attn
7070//      layers, 96 on SWA). Every generic site reads the `cfg.n_head` scalar, so wq/wo/attn_gate
7071//      shapes and the FA head counts would be wrong on 33 of 45 layers.
7072//   2. `cfg.rope_dim_count` is 128 for this arch, but FULL-attn layers rotate only 64 dims
7073//      (upstream halves `n_rot_full` after the generic loader defaults it). A generic mixer
7074//      would rotate 128 dims on the full layers — silently wrong, plausible-looking logits.
7075//   3. Dual rope base (5e6 full / 1e4 SWA) + `rope_freqs.weight` on FULL layers ONLY.
7076//   4. SWA 3:1 (`sliding_window` 512, pattern [false,true,true,true]).
7077//   5. The head-wise gate is a SEPARATE `attn_gate.weight [n_embd, n_head_l]` (one pre-sigmoid
7078//      scalar per head, broadcast over head_dim) applied to the attention output BEFORE wo —
7079//      not the qwen35 fused-in-wq per-DIM gate that `attn_out_gate()`/`q_gate_split` handle.
7080//      Its input is the POST-attn_norm hidden (`cur`), not the residual.
7081//
7082// Attention scale is the DEFAULT 1/sqrt(n_embd_head_k) (step35.cpp:255) — NOT gemma4's 1.0.
7083impl HybridModel {
7084    /// Per-layer attention geometry: (head_dim, n_kv, n_head, rope_base, scale, is_swa).
7085    pub(crate) fn step35_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
7086        let s = self.cfg.step35.as_ref().unwrap();
7087        let hd = self.cfg.head_dim_k as usize;
7088        (hd,
7089         s.n_head_kv(il as u32) as usize,        // 8 (uniform on 3.7-Flash)
7090         s.n_head(il as u32) as usize,           // 64 full / 96 SWA
7091         s.rope_base(il as u32),                 // 5e6 full / 1e4 SWA
7092         1.0 / (hd as f32).sqrt(),               // step35.cpp:255 kq_scale
7093         s.is_swa(il as u32))
7094    }
7095
7096    /// step35 attention core: everything from the q/k/v projection outputs through the head-wise
7097    /// gate, returning the GATED attention output PRE-`wo` (the caller runs wo, so this composes
7098    /// with both the plain `matmul(&fa.wo, ..)` sites and the f16/into-slab GEMM sites).
7099    ///
7100    /// `hg` = the POST-attn_norm hidden state (upstream's `cur`) — the gate projection's input.
7101    /// `cache`:
7102    ///   * `Some` => PRIME mode: append this chunk's post-rope K / raw V rows into the resident
7103    ///     quantized cache and attend THROUGH the cache view, exactly like the generic
7104    ///     `full_attn_prime_fa_dispatch` (quantize-then-attend on chunk 0 too, so chunk size
7105    ///     cannot decide where a precision edge falls — the grain-free chunk-invariance
7106    ///     contract, lane/chunkinv-flip).
7107    ///   * `None` => pure prefill (forward/forward_last/t2probe): attend over this batch's f32
7108    ///     q/k/v, no cache side effect.
7109    ///
7110    /// SWA WINDOW, and why prefill needs `sdpa_naive_w_quantized_view`: the view is trimmed to
7111    /// the oldest key ANY query in this chunk can reach (`off = base_len - (win-1)`), which
7112    /// bounds the view at `win-1+t` rows, but the mask is still required — inside the chunk,
7113    /// query `qt` may only see view keys `[qt, qt+win-1]`, so the earlier keys the trimmed view
7114    /// still contains must be masked per query. memra's window convention
7115    /// (`sdpa_naive_w_f32`: mask `t < q_pos-(window-1)`, `q_pos = (T_kv-T)+qt`) is bit-for-bit
7116    /// upstream's `LLAMA_SWA_TYPE_STANDARD` (`llama-hparams.h:359`: mask `p1-p0 >= n_swa`), and
7117    /// step35.cpp:6 sets exactly that swa_type — verified verbatim on both sides.
7118    ///
7119    /// Note the f32 `sdpa_naive_w` floor's shared memory is `t_kv*4` bytes, so the SWA arm needs
7120    /// `t_kv = win-1+t <= 12287` — i.e. it relies on chunked prefill (MEMRA_PRIME_CHUNK, default
7121    /// 4096). A monolithic 32k prime would exceed the 48 KiB dynamic-smem default.
7122    ///
7123    /// SWA ARM SELECTION IS KEYED ON `seq_end`, NOT ON THE CHUNK'S OWN `t_kv` (lane/step35-chunkfix,
7124    /// 2026-08-07). `seq_end` = the ABSOLUTE end position of the whole prime call
7125    /// (`cache.pos + prompt_len` at entry; `t` on the cacheless prefill path) — a property of the
7126    /// REQUEST, identical at every MEMRA_PRIME_CHUNK value. The old predicate `swa && t_kv > win`
7127    /// read the chunk's own extent, which made kernel selection — and therefore the logits, the
7128    /// hidden rows, and the generated text — a function of the chunk size:
7129    ///   a chunk [b,e) has off = max(0, b-(win-1)), t_kv = e-off, so b <= win-1 => off=0 => FA iff
7130    ///   e <= win, while every later chunk has t_kv = t+(win-1) > win. The FA rows were therefore a
7131    ///   contiguous PREFIX [0,P) with P = c*floor(win/c) for c <= win (else 0), and the output
7132    ///   depended only on P. Measured (research/step37-p2-20260806, closed-form receipt
7133    ///   `raw/chunkinv-step35-GAP2-CONFIRMED-20260807.txt`): at T=4883, c=512 and c=64 (P=512) are
7134    ///   byte-identical to each other but DIFFER from the c=4096 default (P=0) by maxdiff 1.813e0
7135    ///   with greedy text diverging at step 6; c=513 (P=0) is bit-identical to the default. A
7136    ///   one-token change in a documented machine-config knob changed the answer.
7137    /// So the pre-fix comment "same cache bytes, same numeric class" was false as written: same
7138    /// bytes, DIFFERENT numeric class — swapping `fa_prefill_view_ws` for the f32 windowed floor on
7139    /// the same rows moves the logits by ~1.8.
7140    ///
7141    /// Keying on `seq_end` makes P identically 0 for every chunk size, so MEMRA_PRIME_CHUNK is a
7142    /// pure memory/transient knob again (research/step35-chunkfix-20260807). It is
7143    /// correct-by-construction rather than a tolerance argument — the windowed kernel computes the
7144    /// same masked attention the FA arm computed (for e <= win the window mask is a no-op under
7145    /// causal, which is exactly why FA was legal there), it is now simply the ONLY arm used once
7146    /// the request passes the window. It also cannot move the shipped default: at chunk=4096 every
7147    /// chunk of a `seq_end > win` prime ALREADY had t_kv > win (chunk 0 has t_kv = min(chunk,
7148    /// seq_end) > win; every later chunk has t_kv >= t+win-1 > win since `PRIME_MIN_T` keeps
7149    /// t >= 16), and a `seq_end <= win` prime keeps every chunk on FA exactly as before — so this
7150    /// is a no-op on both default regimes and only removes the FA prefix at small chunk values.
7151    /// The `t_kv <= 12287` smem ceiling is respected: the chunks whose arm changes are precisely
7152    /// those with t_kv <= win = 512.
7153    #[allow(clippy::too_many_arguments)]
7154    fn step35_attn_pre_wo(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
7155                          hg: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
7156                          cache: Option<&mut Cache>, il: usize, seq_end: usize)
7157                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7158        let (hd, nkv, nh, rbase, scale, swa) = self.step35_geom(il);
7159        let eps = self.cfg.rms_eps;
7160        let s = self.cfg.step35.as_ref().unwrap();
7161        let win = s.sliding_window as usize;
7162        let n_rot = s.n_rot(il as u32) as usize;
7163
7164        let v = g3.pop().unwrap();
7165        let k0 = g3.pop().unwrap();
7166        let q0 = g3.pop().unwrap();
7167
7168        // q/k RMSNorm over head_dim rows (attn_q_norm/attn_k_norm [128] F32), then the
7169        // per-layer PARTIAL NEOX rope: n_rot = 64 on full layers (dims 64..127 pass through
7170        // unrotated), 128 on SWA. rope_freqs (llama3-style per-dim factors) on FULL only.
7171        let mut q = e.uninit(t * nh * hd)?;
7172        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh * t, eps)?;
7173        let mut k = e.uninit(t * nkv * hd)?;
7174        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv * t, eps)?;
7175        let ff = if swa { None } else {
7176            self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
7177        };
7178        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, t, rbase, 1.0, ff)?;
7179
7180        let mut attn = e.uninit(t * nh * hd)?;
7181        match cache {
7182            Some(cache) => {
7183                let base_len = {
7184                    let kvl = cache.kv[il].as_mut().unwrap();
7185                    assert!(kvl.len + t <= cache.max_ctx, "step35 prime: KV overflow");
7186                    let base_len = kvl.len;
7187                    e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
7188                                               kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
7189                                               kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
7190                    kvl.len += t;
7191                    let new_len = kvl.len as i32;
7192                    e.set_i32_one(&mut kvl.len_d, new_len)?;
7193                    base_len
7194                };
7195                let kvl = cache.kv[il].as_ref().unwrap();
7196                // Both segmentation seams restore the FULL pre-fix arithmetic: their legacy
7197                // arm predicate (`t_kv` below, per-call `seq_end` in prime_cache) AND the
7198                // unaligned view offset here. Both halves are load-bearing for the canaries:
7199                // on the FA default the predicate arms agree bitwise wherever they can differ
7200                // (a t_kv<=win view has no maskable key, so windowed==unwindowed FA bit-for-bit).
7201                // That made predicate-only seams INERT under Lever A — first chunkinv35c, then
7202                // tickinv35c. The raw offset restores the live segmentation-variant mechanism
7203                // on the current FA path: its tile grid starts at the chunk/call boundary.
7204                // Read per layer call, never in a measured default.
7205                let legacy_tkv = std::env::var("MEMRA_STEP35_SWA_TKV").as_deref() == Ok("1");
7206                let legacy_calllocal =
7207                    std::env::var("MEMRA_PRIME_CALLLOCAL").as_deref() == Ok("1");
7208                // SWA: trim the view to the oldest key any query in this chunk can reach —
7209                // ALIGNED DOWN to the FA tile size (BK=32). The raw trim is chunk-DEPENDENT
7210                // (off = base_len-(win-1), and base_len is a chunk boundary), and the FA
7211                // kernel's online-softmax recurrence groups keys into BK tiles relative to
7212                // the VIEW START — so an unaligned off regroups the same absolute keys into
7213                // different tiles at different chunk sizes = different (m,l) rounding =
7214                // chunk-dependent bits. chunkinv35 caught exactly this on the first FA-arm
7215                // battery (first_div == chunk size; raw/leverA-gates-20260807T135541Z.log).
7216                // Aligning off to 32 pins tiles to ABSOLUTE key positions for every chunk
7217                // size; the <=31 extra leading keys are older than EVERY query's window
7218                // (all queries sit at >= base_len, so keys < base_len-(win-1) are masked
7219                // for all of them) and a fully-masked key is an exact-0.0 no-op in both
7220                // kernels (NEG_INF -> p=0.0; l+=0.0 and O+=0.0 are bitwise identity), so
7221                // the floor arm's bits do not move either (gated: G2f, battery 2).
7222                let off = if swa {
7223                    let raw = base_len.saturating_sub(win - 1);
7224                    if legacy_tkv || legacy_calllocal { raw } else { raw & !31usize }
7225                } else {
7226                    0
7227                };
7228                let t_kv = base_len + t - off;
7229                let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
7230                                             (off + t_kv) * kvl.k_tok_bytes);
7231                let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
7232                                             (off + t_kv) * kvl.v_tok_bytes);
7233                // CHUNK-INVARIANT PREDICATE: `seq_end` (the request's absolute end position), not
7234                // this chunk's `t_kv`. See the doc note — keying on t_kv made P = c*floor(win/c)
7235                // rows take FA and the output a function of MEMRA_PRIME_CHUNK.
7236                // MEMRA_STEP35_SWA_TKV=1 is the ROLLBACK SEAM to the pre-fix arithmetic — BOTH
7237                // halves: this predicate AND the unaligned view offset above (`legacy_tkv`).
7238                // It is what gives the step35 chunkinv gate its canary teeth: chunk-VARIANT by
7239                // construction, so the invariance assertion MUST break under it (the seam whose
7240                // absence made the original canary inert — GAP 1 in research/step37-p2-20260806;
7241                // and a predicate-only seam went inert AGAIN under the FA default, battery 2's
7242                // CANARY UNEXPECTEDLY MATCHED — see the offset comment). Read per call, not
7243                // cached (probes flip it in-process). Never on in a measured default run.
7244                let swa_naive = if legacy_tkv { t_kv > win } else { seq_end > win };
7245                if swa && swa_naive {
7246                    // Windowed mask needed (see the doc note). DEFAULT since lane/pp-prefill
7247                    // 2026-08-07: the windowed hd128 FA stamp (`fa_prefill_view_ws_w_hd128`) —
7248                    // the anatomy profile measured the f32 floor at 565 ms/layer on a pp4096
7249                    // (41% of the whole prime) while the unwindowed hd128 FA family did the
7250                    // strictly harder causal-4096 in 3.3 ms. NOTE t_kv can be <= win here
7251                    // (a chunk-0 view on a small chunk); the windowed kernel handles that
7252                    // identically to the unwindowed one modulo the mask, which is the point.
7253                    // NEW NUMERIC CLASS vs the floor (bf16-MMA online softmax vs f32 serial),
7254                    // selected on `seq_end` like every arm here, so the class is uniform for
7255                    // the whole request at every MEMRA_PRIME_CHUNK — chunkinv holds by the
7256                    // same construction as the chunkfix. MEMRA_STEP35_SWA_FA=0 = rollback to
7257                    // the f32 floor (the previous numeric config, kept as the A/B seam).
7258                    if std::env::var("MEMRA_STEP35_SWA_FA").as_deref() == Ok("0") {
7259                        e.sdpa_naive_w_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh,
7260                                                      nkv, t, t_kv, scale, true, win,
7261                                                      kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7262                    } else {
7263                        e.fa_prefill_view_ws_w_hd128(&q, &k_view, &v_view, &mut attn, hd, nh,
7264                                                     nkv, t, t_kv, scale, true, win,
7265                                                     kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7266                    }
7267                } else if std::env::var("MEMRA_NOFA").is_ok() {
7268                    e.sdpa_naive_quantized_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7269                                                t, t_kv, scale, true,
7270                                                kvl.k_tok_bytes, kvl.v_tok_bytes)?;
7271                } else {
7272                    // seq_end <= win (or a full-attn layer): no query in the WHOLE request can
7273                    // reach past the window, so the window mask is a no-op under causal and every
7274                    // chunk rides the hd128 dequant-once FA prefill — one arm for the whole
7275                    // request either way, which is what makes the chunk size arithmetic-free.
7276                    e.fa_prefill_view_ws(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7277                                         t, t_kv, scale, true,
7278                                         kvl.k_tok_bytes, kvl.v_tok_bytes,
7279                                         crate::Engine::kv_fp8_on())?;
7280                }
7281            }
7282            None => {
7283                // Cacheless prefill (forward/forward_last/t2probe) is MONOLITHIC — there is no
7284                // chunk loop on this path, so seq_end == t and the predicate is unchanged. The
7285                // assert pins that: if a chunked cacheless prefill is ever added, it must thread
7286                // seq_end here too or it re-opens the same door.
7287                debug_assert_eq!(seq_end, t, "step35 cacheless prefill is monolithic (seq_end == t)");
7288                if swa && seq_end > win {
7289                    e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
7290                } else if std::env::var("MEMRA_NOFA").is_ok() {
7291                    e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7292                } else {
7293                    e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7294                }
7295            }
7296        }
7297
7298        // HEAD-WISE GATE (step35.cpp:267-285): gate = attn_gate(cur) -> [t, n_head_l];
7299        // attn *= sigmoid(gate) broadcast over head_dim. BEFORE wo.
7300        let gw = fa.attn_gate.as_ref()
7301            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
7302        let gt = e.matmul(gw, hg, t)?;
7303        let mut ag = e.uninit(t * nh * hd)?;
7304        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, t)?;
7305        Ok(ag)
7306    }
7307
7308    /// step35 PREFILL mixer, no cache side effect (the `full_attn` contract: `forward`,
7309    /// `forward_last`, t2probe). Post-`wo`.
7310    pub(crate) fn step35_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
7311                              pos_d: &CudaSlice<i32>, t: usize, il: usize)
7312                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7313        let g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
7314        // Monolithic path: the request ends at t (no chunk loop). See step35_attn_pre_wo's note.
7315        let ag = self.step35_attn_pre_wo(e, fa, g3, h, pos_d, t, None, il, t)?;
7316        Ok(e.matmul(&fa.wo, &ag, t)?)
7317    }
7318
7319    /// step35 PRIME mixer (the `full_attn_prime` contract: append this chunk's K/V into the
7320    /// resident quantized cache, attend through the cache view). Post-`wo`.
7321    ///
7322    /// `seq_end` = the ABSOLUTE end position of the whole prime request (chunk-invariant); it
7323    /// selects the SWA arm. See `step35_attn_pre_wo`'s doc note for why it cannot be this chunk's
7324    /// own extent.
7325    #[allow(clippy::too_many_arguments)]
7326    pub(crate) fn step35_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
7327                                    hx: Option<&CudaSlice<u8>>, pos_d: &CudaSlice<i32>, t: usize,
7328                                    cache: &mut Cache, il: usize, seq_end: usize)
7329                                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7330        let g3 = match hx {
7331            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
7332            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
7333        };
7334        let ag = self.step35_attn_pre_wo(e, fa, g3, h, pos_d, t, Some(cache), il, seq_end)?;
7335        Ok(e.matmul(&fa.wo, &ag, t)?)
7336    }
7337
7338    /// step35 T=1 decode attention (post-`wo`, matching `full_attn_decode_pre`'s contract).
7339    /// `pre_q` = the attn-input norm's q8_1 pair when the caller took the norm-fusion lever
7340    /// (then `h` is a zero-length placeholder and EVERY projection here — including the gate —
7341    /// must be on the q8_1 fast path; `mixer_in_q8_1_fast` enforces that for step35 by also
7342    /// requiring `attn_gate`).
7343    ///
7344    /// SWA decode is a token-aligned VIEW OFFSET into the quantized cache (the gemma4 R6
7345    /// pattern): keys carry absolute rope and the mask is purely positional, so the single
7346    /// query at `len-1` attending the last `win` rows IS the windowed result — no mask kernel.
7347    #[allow(clippy::too_many_arguments)]
7348    pub(crate) fn step35_decode_attn(&self, e: &Engine, fa: &FullAttnLayer, il: usize,
7349                          h: &CudaSlice<f32>,
7350                          pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7351                          pos_d: &CudaSlice<i32>, cache: &mut Cache)
7352                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7353        let (hd, nkv, nh, rbase, scale, swa) = self.step35_geom(il);
7354        let eps = self.cfg.rms_eps;
7355        let s = self.cfg.step35.as_ref().unwrap();
7356        let win = s.sliding_window as usize;
7357        let n_rot = s.n_rot(il as u32) as usize;
7358        let n_embd = self.cfg.n_embd as usize;
7359        let gw = fa.attn_gate.as_ref()
7360            .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
7361
7362        let (q0, k0, v0, gt) = match pre_q {
7363            Some((hq, hdq)) => {
7364                debug_assert!(e.uses_q8_1_fast(gw),
7365                    "step35 pre-quantized decode requires attn_gate on the q8_1 fast path \
7366                     (h is a zero-length placeholder here) — see mixer_in_q8_1_fast");
7367                let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
7368                    Some(t3) => t3,
7369                    None => (e.matmul_pre(&fa.wq, hq, hdq, h, 1)?,
7370                             e.matmul_pre(&fa.wk, hq, hdq, h, 1)?,
7371                             e.matmul_pre(&fa.wv, hq, hdq, h, 1)?),
7372                };
7373                let gt = e.matmul_pre(gw, hq, hdq, h, 1)?;
7374                (a, b, c, gt)
7375            }
7376            None => {
7377                if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk)
7378                    && e.uses_q8_1_fast(&fa.wv) && e.uses_q8_1_fast(gw) {
7379                    let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
7380                    let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
7381                        Some(t3) => t3,
7382                        None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
7383                                 e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
7384                                 e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
7385                    };
7386                    let gt = e.matmul_pre(gw, &hq, &hdq, h, 1)?;
7387                    (a, b, c, gt)
7388                } else {
7389                    (e.matmul(&fa.wq, h, 1)?, e.matmul(&fa.wk, h, 1)?,
7390                     e.matmul(&fa.wv, h, 1)?, e.matmul(gw, h, 1)?)
7391                }
7392            }
7393        };
7394
7395        let mut q = e.uninit(nh * hd)?;
7396        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
7397        let mut k = e.uninit(nkv * hd)?;
7398        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
7399        let ff = if swa { None } else {
7400            self.step35_aux.as_ref().and_then(|a| a.rope_freqs.as_ref())
7401        };
7402        e.rope_neox2(&mut q, &mut k, pos_d, hd, n_rot, nh, nkv, 1, rbase, 1.0, ff)?;
7403
7404        if std::env::var("MEMRA_NOFA").is_ok() {
7405            return Err("MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV \
7406                        cache; unset MEMRA_NOFA to use fa_decode".into());
7407        }
7408        let kvl = cache.kv[il].as_mut().unwrap();
7409        e.append_kv_quantized(&k, &v0, &mut kvl.k, &mut kvl.v, kvl.len,
7410                              kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
7411                              crate::Engine::kv_fp8_on())?;
7412        kvl.len += 1;
7413        let (off, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
7414        let k_view = e.view_u8_range(&kvl.k, off * kvl.k_tok_bytes,
7415                                     (off + t_kv) * kvl.k_tok_bytes);
7416        let v_view = e.view_u8_range(&kvl.v, off * kvl.v_tok_bytes,
7417                                     (off + t_kv) * kvl.v_tok_bytes);
7418        let mut attn = e.uninit(nh * hd)?;
7419        e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
7420                          kvl.k_tok_bytes, kvl.v_tok_bytes, crate::Engine::kv_fp8_on())?;
7421
7422        let mut ag = e.uninit(nh * hd)?;
7423        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
7424        Ok(e.matmul(&fa.wo, &ag, 1)?)
7425    }
7426}
7427
7428// ===================================================================================== //
7429//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
7430//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
7431//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
7432//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
7433//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
7434//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
7435// ===================================================================================== //
7436impl HybridModel {
7437    pub fn is_gemma4_e4b(&self) -> bool {
7438        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
7439    }
7440
7441    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
7442    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
7443    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
7444    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
7445        let g = self.cfg.gemma4.as_ref().unwrap();
7446        let swa = g.swa_pattern[il];
7447        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
7448        let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
7449        let nh = fa.wq.out_features() / hd;
7450        let nkv = fa.wk.out_features() / hd;
7451        (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
7452    }
7453
7454    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
7455    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
7456        self.layers[il].gemma4.as_ref()
7457            .and_then(|b| b.e4b.as_ref())
7458            .and_then(|e4| e4.kv_share.map(|t| t as usize))
7459    }
7460
7461    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
7462    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
7463    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
7464    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
7465    fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
7466                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7467        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
7468        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
7469    }
7470
7471    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
7472    fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
7473                             x_scaled: &CudaSlice<f32>, t: usize)
7474                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7475        let aux = self.gemma4_aux.as_ref().unwrap();
7476        let m = aux.e4b.as_ref().unwrap();
7477        let n_embd = self.cfg.n_embd as usize;
7478        let n_layer = self.layers.len();
7479        let width = m.n_epl * n_layer;
7480        let tbl = m.tok_tbl_gpu.get_or_init(|| {
7481            e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
7482        });
7483        let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
7484                                             m.tok_embd_row_bytes)?;
7485        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
7486        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
7487        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
7488        let mut pn = e.uninit(t * width)?;
7489        e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
7490                   self.cfg.rms_eps)?;
7491        let mut out = e.uninit(t * width)?;
7492        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
7493        Ok(out)
7494    }
7495
7496    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
7497    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
7498    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
7499    /// already holds this forward's rows — the target runs earlier in the stack).
7500    #[allow(clippy::too_many_arguments)]
7501    fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
7502                       hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
7503                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
7504                       dc_bucket: Option<usize>)
7505                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7506        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
7507        let eps = self.cfg.rms_eps;
7508        let aux = self.gemma4_aux.as_ref().unwrap();
7509        let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
7510        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
7511        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
7512        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
7513        let h0 = e.zeros(0)?;
7514        let h = &h0;
7515
7516        let ff = if swa { None } else {
7517            Some(aux.rope_freqs.as_ref().expect("e4b global rope needs rope_freqs.weight"))
7518        };
7519        let share = self.gemma4_e4b_kv_target(il);
7520        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
7521        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7522        let mut q;
7523        if let Some(_tgt) = share {
7524            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
7525            q = e.uninit(t * nh * hd)?;
7526            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
7527            // empty; q0 stands in for the unused k/v pointers).
7528            let mut kdummy = e.uninit(1)?;
7529            let mut vdummy = e.uninit(1)?;
7530            e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
7531                                fa.q_norm.float_data(), &aux.ones,
7532                                &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
7533                                pos_d, nh, 1, base, 1.0, ff, eps)?;
7534        } else {
7535            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
7536            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
7537            // q|k|v rows — the cat norm+rope twin consumes it directly.
7538            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
7539            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
7540            q = e.uninit(t * nh * hd)?;
7541            let mut k = e.uninit(t * nkv * hd)?;
7542            let mut v = e.uninit(t * nkv * hd)?;
7543            if t == 1 && cat.is_some() {
7544                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
7545                e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
7546                                        &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
7547                                        pos_d, nh, nkv, base, 1.0, ff, eps)?;
7548            } else {
7549                let (q0, k0, v0) = match if t == 1 {
7550                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
7551                } else {
7552                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
7553                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
7554                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7555                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
7556                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
7557                    } else { None }
7558                } {
7559                    Some(triple) => triple,
7560                    None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
7561                             e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
7562                             e.matmul_pre(&fa.wv, hq, hdq, h, t)?),   // E4B: real v (K != V)
7563                };
7564                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
7565                // the normed rows; V ones-rms, never roped).
7566                e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
7567                                    fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v,
7568                                    hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
7569            }
7570            let kvl = cache.kv[il].as_mut().unwrap();
7571            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
7572            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
7573            // degenerate tok-0 stream, 2026-07-12).
7574            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7575            if dc_bucket.is_some() {
7576                // DC arm (graph serving): append at the len_d slot, advance the counter
7577                // in-stream — replay-correct, no host len in the launch args. Host mirrors
7578                // are NOT touched here (the replay loop owns them; a bump at capture-record
7579                // time would double-count the capture iteration).
7580                debug_assert!(t == 1);
7581                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
7582                e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
7583                                                 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
7584                                                 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
7585            } else {
7586                e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
7587                                           kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
7588                                           kvl.v_tok_bytes, cls)?;
7589                kvl.len += t;
7590            }
7591            kv_f32 = Some((k, v));
7592        }
7593        // attention: per-row causal fa over the (own or target) quantized cache. The cache
7594        // already contains this forward's rows in both arms; row i attends [.., base+i].
7595        let kvl_idx = share.unwrap_or(il);
7596        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
7597        let base_len = kvl.len - t;   // pre-append length (target appended this forward too)
7598        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
7599        let mut attn = e.uninit(t * nh * hd)?;
7600        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
7601        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
7602        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
7603        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
7604        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
7605        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
7606        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
7607        //     rows (the T=K verify kernel; the target appended this forward's rows already).
7608        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
7609        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
7610        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
7611        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
7612            if let Some((kf, vf)) = &kv_f32 {
7613                if hd == 256 && t <= win {
7614                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
7615                    return Ok(e.matmul(&fa.wo, &attn, t)?);
7616                }
7617                if hd == 256 && swa && t > win {
7618                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
7619                                   win)?;
7620                    return Ok(e.matmul(&fa.wo, &attn, t)?);
7621                }
7622                if hd == 512 && !swa {
7623                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
7624                                       true)?;
7625                    return Ok(e.matmul(&fa.wo, &attn, t)?);
7626                }
7627            } else if share.is_some() {
7628                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7629                let k_view = e.view_u8(&kvl.k, kvl.k.len());
7630                let v_view = e.view_u8(&kvl.v, kvl.v.len());
7631                if hd == 256 && (!swa || t <= win) {
7632                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
7633                    e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
7634                                      scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7635                    return Ok(e.matmul(&fa.wo, &attn, t)?);
7636                }
7637                // remaining shared classes (swa above the window; hd512 globals): dequant
7638                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
7639                let kv_dim = nkv * hd;
7640                let mut kf = e.uninit(t * kv_dim)?;
7641                let mut vf = e.uninit(t * kv_dim)?;
7642                e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
7643                                         t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7644                if hd == 512 {
7645                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
7646                                       true)?;
7647                } else {
7648                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
7649                                   win)?;
7650                }
7651                return Ok(e.matmul(&fa.wo, &attn, t)?);
7652            }
7653        }
7654        if let Some(bucket) = dc_bucket {
7655            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
7656            // fa_decode_dc over the live counter. len_d already advanced past this token
7657            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
7658            // counter (advanced when the target ran earlier in the stack).
7659            assert!(t == 1);
7660            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
7661            // and under the window every live t_kv sits below it — cap the capture bucket
7662            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
7663            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
7664            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
7665            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
7666                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
7667            } else { bucket };
7668            let k_view = e.view_u8(&kvl.k, kvl.k.len());
7669            let v_view = e.view_u8(&kvl.v, kvl.v.len());
7670            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
7671            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
7672            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
7673            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
7674            // captured into the dc graph like any other launch. Extending the cascade to
7675            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
7676            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
7677            // MEMRA_WPF=0 rollback seam.
7678            if crate::Engine::wpf_level() >= 1 {
7679                e.prefetch_weight_l2(&fa.wo)?;
7680            }
7681            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
7682            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
7683            if e.uses_q8_1_fast(&fa.wo) {
7684                let mut oq = e.alloc_i8_uninit(nh * hd)?;
7685                let mut od = e.zeros(nh * hd / 32)?;
7686                e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7687                                  &kvl.len_d, bucket, scale,
7688                                  kvl.k_tok_bytes, kvl.v_tok_bytes, g,
7689                                  Some((&mut oq, &mut od)))?;
7690                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
7691            }
7692            e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
7693                           &kvl.len_d, bucket, scale,
7694                           kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
7695            return Ok(e.matmul(&fa.wo, &attn, t)?);
7696        }
7697        for i in 0..t {
7698            let avail = base_len + i + 1;
7699            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
7700            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
7701                                         (off_tok + t_kv) * kvl.k_tok_bytes);
7702            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
7703                                         (off_tok + t_kv) * kvl.v_tok_bytes);
7704            let qv = e.view(&q, t * nh * hd);
7705            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
7706            let mut q_one = e.uninit(nh * hd)?;
7707            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
7708            let mut a_one = e.uninit(nh * hd)?;
7709            // read class MUST match the append class (globals are e4m3 under gkv): the
7710            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
7711            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
7712            e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
7713                        kvl.k_tok_bytes, kvl.v_tok_bytes,
7714                        (!swa && crate::Engine::gkv_on())
7715                            || (swa && crate::Engine::wkv_on()))?;
7716            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
7717        }
7718        Ok(e.matmul(&fa.wo, &attn, t)?)
7719    }
7720
7721    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
7722    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
7723    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
7724    /// layer; does NOT advance cache.pos (caller owns pos).
7725    fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
7726                        head_last: bool)
7727                        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7728        let n_embd = self.cfg.n_embd as usize;
7729        let t = tokens.len();
7730        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7731        let pos_d = e.htod_i32(&pos)?;
7732        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7733        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7734        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
7735        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
7736    }
7737
7738    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
7739    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
7740    /// eager chain by construction: SAME functions, not twins).
7741    fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
7742                             pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
7743                             dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
7744                             -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7745        let n_embd = self.cfg.n_embd as usize;
7746        let eps = self.cfg.rms_eps;
7747        let n_layer = self.layers.len();
7748        let mut x = x_in;
7749        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
7750        let n_epl = aux_e4b.n_epl;
7751
7752        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
7753        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
7754        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
7755        // head rides matmul_pre too. First layer's pair comes from a standalone fused
7756        // norm+quant.
7757        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
7758        for il in 0..n_layer {
7759            let layer = &self.layers[il];
7760            let (hq, hdq) = match h_carry.take() {
7761                Some(p) => p,
7762                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7763            };
7764            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
7765            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
7766            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
7767            let bits = layer.gemma4.as_ref().unwrap();
7768            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
7769            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
7770            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
7771            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
7772            // the fused single-phase reduction is NOT FP-order-identical to the unfused
7773            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
7774            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
7775            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
7776            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
7777            // gate dropped, decode AND verify ride the same fused chain — parity by
7778            // construction, VERIFY-GATE 0.000e0.
7779            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
7780            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
7781                e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
7782            let mut resid = e.uninit(t * n_embd)?;
7783            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
7784            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
7785            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
7786            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
7787            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
7788            let g = if fuse_exit {
7789                // sn here = RAW f0 (post_ffw deferred).
7790                let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
7791                                                  &attn_out, &mut resid, n_embd, t,
7792                                                  self.cfg.rms_eps)?;
7793                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
7794            } else {
7795                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
7796                e.matmul(&e4b.inp_gate, &resid, t)?
7797            };
7798            let mut act = e.uninit(t * n_epl)?;
7799            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
7800                let ipv = e.view(&inp_pl, n_epl * n_layer);
7801                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
7802                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
7803                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
7804            } else {
7805                let mut inp_this = e.uninit(t * n_epl)?;
7806                e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
7807                                    il * n_epl)?;
7808                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
7809                e.matmul(&e4b.proj, &act, t)?
7810            };
7811            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
7812            // ONE launch (glue-fusion lane; last layer emits through output_norm).
7813            let next_norm = if il + 1 < n_layer {
7814                self.layers[il + 1].attn_norm.float_data()
7815            } else {
7816                self.output_norm.float_data()
7817            };
7818            let mut xn = e.uninit(t * n_embd)?;
7819            let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
7820                                                         &resid, bits.layer_scale, next_norm,
7821                                                         &mut xn, n_embd, t, eps)?;
7822            h_carry = Some(pair);
7823            x = xn;
7824        }
7825        // the head consumes the last layer's fused (output_norm) emit. head_last callers
7826        // (prime, last_only forward) need only the final row's logits — the all-T head is
7827        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
7828        let (oq, odq) = h_carry.take().unwrap();
7829        let h0 = e.zeros(0)?;
7830        let hm = if head_last { 1 } else { t };
7831        let (hq, hd) = if head_last && t > 1 {
7832            let mut q1 = e.uninit_i8(n_embd)?;
7833            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
7834            let nb = n_embd / 32;
7835            let mut d1 = e.uninit(nb)?;
7836            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
7837            (q1, d1)
7838        } else {
7839            (oq, odq)
7840        };
7841        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
7842        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
7843        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
7844        // Logit-returning callers (host logits / spec prime) keep the capped emit.
7845        if cap_logits {
7846            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7847            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
7848        }
7849        self.gemma4_suppress(e, &mut ld, hm)?;  // mask both capped and argmax-only consumers
7850        Ok((ld, x))
7851    }
7852
7853    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
7854    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
7855    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
7856    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
7857    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
7858    /// covers exactly the layers that appended).
7859    pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
7860                                                  t: usize, pos0: usize, cache: &mut Cache)
7861                                                  -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7862        let n_embd = self.cfg.n_embd as usize;
7863        let eps = self.cfg.rms_eps;
7864        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7865        let pos_d = e.htod_i32(&pos)?;
7866        let embd_gpu = self.embd_gpu.get_or_init(|| {
7867            e.upload_u8(&self.embd.raw).expect("embed table upload")
7868        });
7869        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
7870        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
7871        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7872        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
7873        let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
7874                                                  false)?;
7875        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
7876        // emit is already capped, matching the eager chain bit-for-bit).
7877        let n_vocab = self.output.out_features();
7878        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
7879        for i in 0..t {
7880            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
7881        }
7882        let mut hn = e.uninit(t * n_embd)?;
7883        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7884        cache.pos += t;
7885        Ok((vam, hn))
7886    }
7887
7888    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
7889    /// prime path — mirror of `gemma4_decode_step_t_h`).
7890    pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
7891                                             cache: &mut Cache)
7892                                             -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7893        let n_embd = self.cfg.n_embd as usize;
7894        let eps = self.cfg.rms_eps;
7895        let t = tokens.len();
7896        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
7897        let mut hn = e.uninit(t * n_embd)?;
7898        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7899        cache.pos += t;
7900        Ok((e.dtoh(&ld)?, hn))
7901    }
7902
7903    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
7904    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
7905    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
7906    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
7907    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
7908    pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
7909                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7910                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7911                                      n_vocab: usize, bucket: usize)
7912                                      -> Result<(), Box<dyn std::error::Error>> {
7913        let n_embd = self.cfg.n_embd as usize;
7914        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7915        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7916        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7917        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
7918                                                  false, false)?;
7919        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
7920        e.inc_seqlen(pos_d)?;
7921        Ok(())
7922    }
7923
7924    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
7925    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
7926    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
7927    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
7928    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
7929    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
7930    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
7931    #[allow(clippy::too_many_arguments)]
7932    pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7933                                     pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7934                                     embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7935                                     n_vocab: usize)
7936                                     -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7937        let n_embd = self.cfg.n_embd as usize;
7938        let eps = self.cfg.rms_eps;
7939        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7940        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7941        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7942        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
7943                                                  false)?;
7944        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7945        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
7946        e.inc_seqlen(pos_d)?;
7947        cache.pos += 1;
7948        let _ = eps;
7949        Ok(tok_out)
7950    }
7951
7952    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
7953    /// pre-output_norm hidden). Advances cache.pos.
7954    pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
7955                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7956        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
7957        let logits = e.dtoh(&ld)?;
7958        cache.pos += 1;
7959        Ok((logits, x))
7960    }
7961
7962    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
7963    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
7964    /// fast; the prefill fa arms come later.
7965    pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7966                                   -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7967        // Err, not assert (2026-08-07, lane/gemma4-serve-gaps): same served-chunked-prompt
7968        // process-kill as gemma4_prime — refuse per-request.
7969        if cache.pos != 0 {
7970            return Err("e4b prime is fresh-prompt only (v0) — prime the full prompt in one \
7971                        call or decode tokenwise".into());
7972        }
7973        let n_embd = self.cfg.n_embd as usize;
7974        let t = tokens.len();
7975        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
7976        cache.pos += t;
7977        let last = e.dtoh(&ld)?;   // head_last: ld is already the final row only
7978        let xv = e.view(&x, t * n_embd);
7979        let row = xv.slice((t - 1) * n_embd..t * n_embd);
7980        let mut h_seed = e.uninit(n_embd)?;
7981        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
7982        Ok((last, h_seed, x))
7983    }
7984
7985    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
7986    pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7987                                     -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7988        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
7989        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
7990        Ok(e.dtoh(&ld)?)   // head_last already reduced to the final row when last_only
7991    }
7992}
7993
7994#[cfg(test)]
7995mod page_prefetch_tests {
7996    use super::{
7997        grouped_worker_prefetch_position, page_prefetch_positions,
7998        page_prefetch_window_from_values, worker_prefetch_positions,
7999    };
8000
8001    #[test]
8002    fn page_prefetch_window_keeps_existing_opt_in_default() {
8003        assert_eq!(page_prefetch_window_from_values(false, None), 0);
8004        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
8005        assert_eq!(page_prefetch_window_from_values(true, None), 1);
8006        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
8007        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
8008        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
8009    }
8010
8011    #[test]
8012    fn rolling_page_prefetch_advises_each_future_expert_once() {
8013        let advised: Vec<_> = (0..7)
8014            .flat_map(|position| page_prefetch_positions(position, 7, 3))
8015            .collect();
8016        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
8017
8018        let one_ahead: Vec<_> = (0..4)
8019            .flat_map(|position| page_prefetch_positions(position, 4, 1))
8020            .collect();
8021        assert_eq!(one_ahead, vec![1, 2, 3]);
8022        assert!(page_prefetch_positions(0, 4, 0).is_empty());
8023    }
8024
8025    #[test]
8026    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
8027        assert_eq!(grouped_worker_prefetch_position(0, None), None);
8028        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
8029            .chain((0..4).filter_map(|position| {
8030                grouped_worker_prefetch_position(4, Some(position))
8031            }))
8032            .collect();
8033        assert_eq!(positions, vec![0, 1, 2, 3]);
8034        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
8035    }
8036
8037    #[test]
8038    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
8039        let queued: Vec<_> = (0..8)
8040            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
8041            .collect();
8042        assert_eq!(queued, (0..8).collect::<Vec<_>>());
8043
8044        let one_at_a_time: Vec<_> = (0..4)
8045            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
8046            .collect();
8047        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
8048        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
8049    }
8050}
8051
8052pub struct G4DcSlots {
8053    x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
8054    hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
8055    q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
8056    q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
8057    attn: CudaSlice<f32>, o: CudaSlice<f32>,
8058    attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
8059    zq: CudaSlice<i8>, zd: CudaSlice<f32>,
8060    gate: CudaSlice<f32>, up: CudaSlice<f32>,
8061    act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
8062    f0: CudaSlice<f32>, sn: CudaSlice<f32>,
8063    hn: CudaSlice<f32>, logits: CudaSlice<f32>,
8064}