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/// Deterministic in-token expert prefetch. `MEMRA_MOE_PREFETCH=1` overlaps memory-source H2D on the
86/// copy stream; selecting the opt-in worker spill backend enables the same known-next hook for disk.
87fn moe_prefetch_enabled() -> bool {
88    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
89    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREFETCH").as_deref() == Ok("1")
90        || crate::spill_pread::worker_enabled())
91}
92
93/// Best-effort OS page-cache prefetch distance for mmap-backed expert ranges. Independent of the
94/// H2D copy-stream experiment so storage->RAM and RAM->HBM overlap can be measured separately.
95/// The opt-in default stays one expert to preserve the original experiment; spill rigs can widen
96/// it with `MEMRA_MOE_PAGE_PREFETCH_WINDOW` to cover NVMe latency.
97fn moe_page_prefetch_window() -> usize {
98    static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
99    *W.get_or_init(|| page_prefetch_window_from_values(
100        std::env::var("MEMRA_MOE_PAGE_PREFETCH").as_deref() == Ok("1"),
101        std::env::var("MEMRA_MOE_PAGE_PREFETCH_WINDOW").ok().as_deref(),
102    ))
103}
104
105fn page_prefetch_window_from_values(enabled: bool, raw_window: Option<&str>) -> usize {
106    if !enabled {
107        return 0;
108    }
109    raw_window
110        .and_then(|value| value.parse().ok())
111        .unwrap_or(1)
112}
113
114/// Return only the newly exposed positions in a rolling lookahead window. Position zero seeds the
115/// full window; each later position adds one expert at the far edge. Thus widening the window does
116/// not repeatedly issue `MADV_WILLNEED` for the same range.
117fn page_prefetch_positions(
118    position: usize,
119    len: usize,
120    window: usize,
121) -> std::ops::Range<usize> {
122    if window == 0 || position >= len {
123        return len..len;
124    }
125    let (start, count) = if position == 0 {
126        (1, window)
127    } else {
128        (position.saturating_add(window), 1)
129    };
130    let start = start.min(len);
131    start..start.saturating_add(count).min(len)
132}
133
134/// Grouped worker-I/O schedule: prime the first active expert before the loop, then queue exactly
135/// one known-next expert at each iteration. Returning positions keeps expert ordering authoritative.
136fn grouped_worker_prefetch_position(order_len: usize, current: Option<usize>) -> Option<usize> {
137    let position = current.map_or(0, |position| position.saturating_add(1));
138    (position < order_len).then_some(position)
139}
140
141/// Fill the worker ring with complete experts, retaining one pinned buffer for an unexpected
142/// demand miss. Each expert has gate/up/down extents, so depth 16 admits a rolling five-expert
143/// window. Position zero primes the current expert too: its three independent reads can run in
144/// parallel instead of demand-serializing gate, up, and down before any useful GPU work exists.
145fn worker_prefetch_window() -> usize {
146    static WINDOW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
147    *WINDOW.get_or_init(|| {
148        let automatic = crate::spill_pread::configured_depth().saturating_sub(1) / 3;
149        std::env::var("MEMRA_SPILL_WORKER_EXPERT_WINDOW")
150            .ok()
151            .and_then(|value| value.parse::<usize>().ok())
152            .unwrap_or(automatic.max(1))
153    })
154}
155
156/// Return only positions newly exposed by a rolling worker-I/O window. Unlike mmap page advice,
157/// this includes the current expert when the window is seeded so all three current projections
158/// enter the CPU pool together.
159fn worker_prefetch_positions(position: usize, len: usize, window: usize) -> std::ops::Range<usize> {
160    if window == 0 || position >= len {
161        return len..len;
162    }
163    let (start, count) = if position == 0 {
164        (0, window)
165    } else {
166        (position.saturating_add(window).saturating_sub(1), 1)
167    };
168    let start = start.min(len);
169    start..start.saturating_add(count).min(len)
170}
171
172/// LAUNCH-STRUCTURE STAGE 3 gate (MEMRA_MOE_DEV, default ON; `=0` restores host routing). The
173/// zero-DtoH device-dispatch path for fully-resident layers: router top-k output stays on device,
174/// expert weight pointers come from the per-layer device table. Requires the fused router (the
175/// dev path consumes the device sel/w directly), so MEMRA_FUSED_ROUTER=0 also disables it.
176fn moe_dev_enabled() -> bool {
177    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
178    *E.get_or_init(|| std::env::var("MEMRA_MOE_DEV").map(|v| v != "0").unwrap_or(true)
179        && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")))
180}
181
182/// MoE EXPERT dp4a gate (MEMRA_MOE_Q8, default ON; `=0` restores the Stage-A f32-dequant expert
183/// kernels). Applies when gate/up/down expert qtypes are all in the dp4a body set (IQ3_S/IQ4_XS).
184/// FP-order differs from Stage-A (int dp4a + warp tree) — argmax/run-gen/stream-identity gates
185/// arbitrate; the sequential and fused q8 paths ship as a matched pair (MEMRA_MOE_GATE contract).
186fn moe_q8_enabled() -> bool {
187    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
188    *E.get_or_init(|| std::env::var("MEMRA_MOE_Q8").map(|v| v != "0").unwrap_or(true))
189}
190
191/// gemma4 fast-arm gate: qtypes with an `expert_dot_g` dp4a body (superset used by the gelu
192/// dev arm; the qwen q8 arms keep their own battery-gated q8_expert_supported policy).
193fn expert_dp4a_supported(qt: i32) -> bool {
194    qt == crate::QT_Q4_0 || qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS
195        || qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K
196}
197
198fn q8_expert_supported(qt: i32) -> bool {
199    // k-quant arms added 2026-07-06 (Q3_K/Q4_K/Q6_K bodies for the UD tail layers). Briefly
200    // default-excluded the same day when they appeared to break 35B real-prompt spec — the
201    // ACTUAL culprit was the MoE router's cuBLASLt n-dependence (d994271); with the router
202    // decode-exact at verify t, the k-quant arms pass the full spec battery (p1/p2/p3 + raw
203    // K=1..8) and are DEFAULT ON again (+9 tok/s: 148.9 -> 157.9). MEMRA_MOE_Q8_KQ=0 excludes.
204    static KQ: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205    let kq = *KQ.get_or_init(|| {
206        std::env::var("MEMRA_MOE_Q8_KQ").map(|v| v != "0").unwrap_or(true)
207    });
208    // NVFP4 experts: DEFAULT ON (2026-07-17). The M3-era "decode-vs-verify MISMATCH 3.4e1"
209    // that had this excluded was the missing per-expert macro-scale fold, fixed in the
210    // dev-kernel epilogues + moe_w_scale_by_expert; the 35B ct-NVFP4 artifact now runs the
211    // q8 arm at parity with the IQ4_XS daily (174-178 tok/s, spec K=1..8 exact). M3/Hy3
212    // never reach the q8 arms regardless (sigmoid-router cfg gates on pairs/dev/gdec).
213    // MEMRA_MOE_Q8_NVFP4=0 restores the f32 arm.
214    let nvfp4_q8 = std::env::var("MEMRA_MOE_Q8_NVFP4").map(|v| v != "0").unwrap_or(true);
215    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || (nvfp4_q8 && qt == crate::QT_NVFP4)
216        || (kq && (qt == crate::QT_Q3_K || qt == crate::QT_Q4_K || qt == crate::QT_Q6_K))
217}
218
219/// The decode-once (_dec) and IQ-MMA expert kernels dequant via IQ-specific extractors —
220/// k-quant tensors must fall to the _em dot path instead.
221fn q8_expert_dec_supported(qt: i32) -> bool {
222    qt == crate::QT_IQ3_S || qt == crate::QT_IQ4_XS || qt == crate::QT_Q4_0
223}
224
225/// Grouped-f16 door (MEMRA_MOE_F16G) per-projection admission: the qtype has a dequant-to-f16
226/// kernel in cu/moe_f16_grouped.cu AND the projection's k dimension tiles its block size.
227/// Round 49 widened coverage to q35's UD mix (gate/up IQ3_S x39 + Q3_K x1 + IQ4_XS x1; down
228/// IQ4_XS x37 + Q6_K x3 + Q4_K x1) — the round-47 IQ4_XS/Q4_0-only table admitted ~1 of 41
229/// q35 layers, which is why that cell measured FLAT.
230fn f16g_proj_ok(qt: i32, in_f: usize) -> bool {
231    match qt {
232        crate::QT_Q4_0 => in_f % 32 == 0,
233        crate::QT_IQ4_XS | crate::QT_IQ3_S | crate::QT_Q3_K | crate::QT_Q4_K
234        | crate::QT_Q6_K => in_f % 256 == 0,
235        _ => false,
236    }
237}
238
239/// STAGE 3 prewarm gate (MEMRA_MOE_PREWARM, default ON; `=0` leaves residency organic). One-shot
240/// per layer: force-admit every block while FREE slots cover the whole layer (never evicts).
241fn moe_prewarm_enabled() -> bool {
242    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
243    *E.get_or_init(|| std::env::var("MEMRA_MOE_PREWARM").map(|v| v != "0").unwrap_or(true))
244}
245
246/// During a discarded fixed-residency profile, admit CPU-routed misses after their current-token
247/// CPU result is complete. The current result and numeric path are unchanged; later warmup tokens
248/// can then vote for and exercise those experts on GPU before the cache is frozen.
249fn cpu_expert_profile_admit_enabled() -> bool {
250    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
251    *E.get_or_init(|| std::env::var("MEMRA_CPU_EXPERT_FREEZE_PROFILE_ADMIT").as_deref() == Ok("1"))
252}
253
254/// Minimum prompt length for the BATCHED cache prime (`prime_cache`). Below this the tokenwise
255/// decode loop wins anyway (the batched path's GEMM dispatch needs m>=16, and the stateful conv
256/// kernel needs T >= d_conv-1). Callers: generate / generate_spec.
257pub const PRIME_MIN_T: usize = 16;
258
259impl HybridModel {
260    /// CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance): `MEMRA_PRIME_TRACE=<path>`
261    /// appends one JSONL row per (chunk, layer) with a hash of that layer's last-row
262    /// post-residual hidden. Diagnostic only — never on in a measured or gated run
263    /// (it forces a dtoh + host hash per layer).
264    fn prime_trace_path() -> Option<&'static str> {
265        static P: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
266        P.get_or_init(|| std::env::var("MEMRA_PRIME_TRACE").ok())
267            .as_deref()
268    }
269
270    /// Prefill forward over `tokens`; returns logits [T, n_vocab] (host f32).
271    pub fn forward(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
272        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, false); }
273        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, false); }
274        let cfg = &self.cfg;
275        let n_embd = cfg.n_embd as usize;
276        let t = tokens.len();
277        let eps = cfg.rms_eps;
278        let pos: Vec<i32> = (0..t as i32).collect();
279        let pos_d = e.htod_i32(&pos)?;
280
281        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
282
283        for (il, layer) in self.layers.iter().enumerate() {
284            // attn_norm
285            let mut h = e.uninit(t * n_embd)?;
286            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
287
288            let mixed = match &layer.mixer {
289                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t)?,
290                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
291                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
292            };
293
294            // residual 1
295            let mut x1 = e.uninit(t * n_embd)?;
296            e.add(&x, &mixed, &mut x1, t * n_embd)?;
297
298            // pre-FFN norm (post_attention_norm), FFN (Dense or MoE), residual 2
299            let mut z = e.uninit(t * n_embd)?;
300            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
301            let ffn_out = match &layer.ffn {
302                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
303                    let n_ff = ffn_gate.out_features();
304                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
305                    let up = g2.pop().unwrap();
306                    let gate = g2.pop().unwrap();
307                    let mut act = e.uninit(t * n_ff)?;
308                    Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, t * n_ff)?;
309                    e.matmul(ffn_down, &act, t)?
310                }
311                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
312            };
313            let mut x2 = e.uninit(t * n_embd)?;
314            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
315            x = x2;
316        }
317
318        let mut hn = e.uninit(t * n_embd)?;
319        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
320        let logits = e.matmul(&self.output, &hn, t)?;
321        Ok(e.dtoh(&logits)?)
322    }
323
324    /// Prefill that returns ONLY the last token's logits — the common case (greedy/sample needs
325    /// just the final position to start decode). Runs the trunk over all T, then the lm_head
326    /// (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T.
327    /// On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill
328    /// cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].
329    pub fn forward_last(&self, e: &Engine, tokens: &[u32]) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
330        if self.cfg.gemma4.is_some() { return self.gemma4_forward(e, tokens, true); }
331        let cfg = &self.cfg;
332        let n_embd = cfg.n_embd as usize;
333        let t = tokens.len();
334        let eps = cfg.rms_eps;
335        let pos: Vec<i32> = (0..t as i32).collect();
336        let pos_d = e.htod_i32(&pos)?;
337
338        let mut x = self.embed(e, tokens)?;   // [T, n_embd]
339        // MEMRA_LAYER_PROBE=1: synchronize + print after every stage — bisects an in-graph
340        // ILLEGAL_ADDRESS to (layer, stage) at ~1 line of output per layer (M3 bring-up tool).
341        let probe = std::env::var("MEMRA_LAYER_PROBE").is_ok();
342        for (il, layer) in self.layers.iter().enumerate() {
343            let mut h = e.uninit(t * n_embd)?;
344            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
345            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} norm ok"); }
346            let mixed = match &layer.mixer {
347                Mixer::Full(fa) => self.full_attn(e, fa, &h, &pos_d, t)?,
348                Mixer::Linear(la) => self.linear_attn(e, la, &h, t)?,
349                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
350            };
351            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} mixer ok"); }
352            let mut x1 = e.uninit(t * n_embd)?;
353            e.add(&x, &mixed, &mut x1, t * n_embd)?;
354            let mut z = e.uninit(t * n_embd)?;
355            e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
356            let ffn_out = match &layer.ffn {
357                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
358                    let n_ff = ffn_gate.out_features();
359                    let mut g2 = e.matmul_group(&[ffn_gate, ffn_up], &z, t)?;
360                    let up = g2.pop().unwrap();
361                    let gate = g2.pop().unwrap();
362                    let mut act = e.uninit(t * n_ff)?;
363                    Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, t * n_ff)?;
364                    e.matmul(ffn_down, &act, t)?
365                }
366                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
367            };
368            if probe { e.stream().synchronize()?; eprintln!("[probe] L{il} ffn ok"); }
369            let mut x2 = e.uninit(t * n_embd)?;
370            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
371            x = x2;
372        }
373        // norm over all T, then slice the LAST row and run lm_head on that single row.
374        let mut hn = e.uninit(t * n_embd)?;
375        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
376        let last = e.view(&hn, t * n_embd);            // [T, n_embd]
377        let last_row = last.slice((t - 1) * n_embd..t * n_embd);  // [1, n_embd]
378        let mut hlast = e.uninit(n_embd)?;
379        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
380        let logits = e.matmul(&self.output, &hlast, 1)?;   // [1, n_vocab] — lm_head on ONE row
381        Ok(e.dtoh(&logits)?)
382    }
383
384    /// BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): `forward_last`'s batched
385    /// prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime's
386    /// ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput.
387    ///   (a) full-attn layers append their T post-RoPE K/V rows into `cache.kv[il]` via the SAME
388    ///       per-row quantize kernel as the decode append (bit-identical cache bytes per row);
389    ///   (b) linear layers run STATEFULLY from the cache's current recurrent state (zero at a
390    ///       fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in,
391    ///       state_out) whose internal sequential t-loop equals T chained T=1 steps — but with
392    ///       the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the
393    ///       spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode
394    ///       argmax gate is the accuracy authority, exactly as for forward_last);
395    ///   (c) `cache.pos`/KV len/len_d advance by T.
396    /// Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd],
397    /// hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec's prompt_h).
398    /// FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within `tokens` alone.
399    /// forward_last itself stays untouched (kernel-check / run-gen gate on it).
400    pub fn prime_cache(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
401                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
402        let n_embd = self.cfg.n_embd as usize;
403        let t = tokens.len();
404        // SESSION CONTINUATION (2026-07-05): cache.pos > 0 = priming a NEW SUFFIX onto a live
405        // session cache — every chunk (including the first) takes the continuation arm
406        // (fa_prefill_view over the quantized past + this chunk). Fresh prime (pos==0) unchanged.
407        assert!(t >= PRIME_MIN_T, "prime_cache needs T >= {PRIME_MIN_T} (caller gates)");
408        assert!(cache.pos + t <= cache.max_ctx, "prime_cache: prompt exceeds cache max_ctx");
409
410        // CHUNKED PRIME (2026-07-05, the long-ctx OOM fix): the monolithic prime allocates
411        // per-layer transients proportional to T (gate/up/act = T*n_ff*4B EACH — 1.5GB apiece at
412        // 16k on the 27B), which OOMs a 24GB card around 16k prompt tokens. Chunk the prompt:
413        // each chunk runs the full layer stack with transients sized to the chunk, appending its
414        // K/V to the resident quantized cache and carrying the GDN conv-ring + recurrent state
415        // through `cache.recur` (linear_attn_prime is already stateful — a chunk boundary is
416        // exactly the state carry it was built for). Full-attn chunks after the first attend to
417        // the QUANTIZED past KV via fa_prefill_view (the spec-verify pattern) — same numeric
418        // class as decode reading the cache. Prompts <= one chunk take the ORIGINAL monolithic
419        // body byte-for-byte (chunk 0 short-circuits to the f32 fa_prefill path).
420        // MEMRA_PRIME_CHUNK sets the chunk size (tokens); 0 disables chunking (monolithic).
421        if self.is_gemma4_e4b() {
422            return self.gemma4_e4b_prime(e, tokens, cache);
423        }
424        if self.cfg.gemma4.is_some() {
425            // gemma4 v0: monolithic fresh-prompt prime (chunked/continuation arms later).
426            return self.gemma4_prime(e, tokens, cache);
427        }
428        let chunk: usize = std::env::var("MEMRA_PRIME_CHUNK").ok()
429            .and_then(|v| v.parse().ok()).unwrap_or(4096);
430        // CHUNK-ORDER INVARIANCE (lane/chunk-invariance, 2026-08-05; vLLM #38561 shape).
431        // MEMRA_PRIME_CHUNK is documented as a memory-transient knob, but it also decides
432        // the prefill's ARITHMETIC, so two rigs with different values produced different
433        // greedy text for the same prompt (research/session-affinity-20260805: 97- and
434        // 149-token prompts). ROOT CAUSE, measured in research/chunk-invariance-20260805
435        // (VERDICT.md) — and it is NOT what docs originally said:
436        //   * NOT trunk GEMM m / reduction order. REFUTED: the prefill GEMM is m-INVARIANT
437        //     (rows [0,32) bit-identical at m=32 vs m=33..80, both quantized wq and the
438        //     output head), so growing a chunk cannot move an existing row's value.
439        //   * NOT the GDN scan segmentation. REFUTED: MEMRA_GDN_CHUNKED=0 (sequential scan,
440        //     no WY segmentation at all) still diverges, so vLLM's mamba-boundary fix does
441        //     not describe our leak.
442        //   * IT IS the attention numeric CLASS edge in full_attn_prime_fa_dispatch, which
443        //     selects on `base_len == 0`: chunk 0 attends over this batch's f32 K/V
444        //     (fa_prefill) while every later chunk attends over the q8_0/q5_1 quantized KV
445        //     cache (fa_prefill_view_ws). The chunk size therefore decides WHERE in the
446        //     prompt that precision edge falls. Signature: per-row maxdiff is exactly 0.0
447        //     before the first boundary and O(1) right after, first_div_pos == chunk size.
448        // The grain-free fix (full_attn_prime_fa_dispatch below) removed that class edge at
449        // the source — every row is in one numeric class, so the chunk size no longer steers
450        // arithmetic and MEMRA_PRIME_CHUNK is a pure memory/transient knob again. The interim
451        // MEMRA_PRIME_INVARIANT/MEMRA_PRIME_GRAIN pin-the-boundary door was superseded by that
452        // fix and KILLED at v0.71 per the flags doctrine (the jsonl + VERDICT.md are the
453        // record); the chunkinv gate asserts byte-identity across chunk sizes naked.
454        if chunk == 0 || t <= chunk {
455            return self.prime_chunk(e, tokens, cache);
456        }
457        let mut hiddens = e.uninit(t * n_embd)?;
458        let mut last: Option<(Vec<f32>, CudaSlice<f32>)> = None;
459        let mut start = 0usize;
460        while start < t {
461            // keep the tail chunk >= PRIME_MIN_T (the stateful conv needs T >= d_conv-1).
462            let mut end = (start + chunk).min(t);
463            if t - end > 0 && t - end < PRIME_MIN_T { end = t; }
464            let (l, hs, x) = self.prime_chunk(e, &tokens[start..end], cache)?;
465            e.copy_into(&mut hiddens, start * n_embd, &x, (end - start) * n_embd)?;
466            last = Some((l, hs));
467            start = end;
468        }
469        let (logits, h_seed) = last.unwrap();
470        Ok((logits, h_seed, hiddens))
471    }
472
473    /// One prime chunk: the full layer stack over `tokens`, continuing from the cache's current
474    /// state (`cache.pos` = tokens already primed; 0 = fresh). Positions/RoPE are absolute
475    /// task #21 de-broadcast: the q/k head count PREP must emit so the scan's consumers
476    /// agree. Compact (num_k) ONLY when the scan will take the chunked+mma route AND
477    /// num_k == num_v/2 (the engine-side hint mirrors this exact formula); everything
478    /// else (s128 verify tier, chunked-off, mma-off) keeps the broadcast layout.
479    fn gdn_hk(e: &Engine, t: usize, num_v: usize, num_k: usize) -> usize {
480        if Engine::gdn_db_on()
481            && Engine::gdn_chunked_enabled() && t >= 16
482            && e.gdn_mma_enabled(Engine::gdn_chunk_size())
483            && num_k * 2 == num_v
484        {
485            num_k
486        } else {
487            num_v
488        }
489    }
490
491    /// task #17: gate for the fused fp16-operand epilogues (silu_mul/gated_rmsnorm/sig_mul
492    /// `_f16out` twins). Bit-identical class (the twins emit the cvt kernel's exact halves),
493    /// but seam-gated (MEMRA_F16OUT=0) so a bit-check can arbitrate, and OFF under verify-exact
494    /// (matmul_group skips the f16 lane there; the twins must not resurrect it).
495    fn f16out_on(e: &Engine, t: usize) -> bool {
496        crate::f16_ffi::pp_f16_enabled() && t >= 16 && !e.verify_exact_on()
497            && std::env::var("MEMRA_F16OUT").as_deref() != Ok("0")
498    }
499
500    /// (cache.pos + i). Returns (last-row logits, h_seed, this chunk's hidden stack [T, n_embd]).
501    /// See HybridModel::prime_slabs — the eager prime's resident trunk transients.
502    pub fn prime_slabs_get(&self, e: &Engine, t: usize, n_embd: usize, n_ff_max: usize)
503                           -> Result<std::sync::MutexGuard<'_, Option<PrimeSlabs>>, Box<dyn std::error::Error>> {
504        let mut g = self.prime_slabs.lock().unwrap();
505        let need_new = match g.as_ref() { None => true, Some(sl) => sl.t_cap < t };
506        if need_new {
507            *g = Some(PrimeSlabs {
508                t_cap: t,
509                h: e.uninit(t * n_embd)?,
510                x1: e.uninit(t * n_embd)?,
511                z: e.uninit(t * n_embd)?,
512                act: e.uninit(t * n_ff_max)?,
513                xa: e.uninit(t * n_embd)?,
514                xb: e.uninit(t * n_embd)?,
515                h16: e.alloc_u8_uninit(t * n_embd * 2)?,
516                z16: e.alloc_u8_uninit(t * n_embd * 2)?,
517                gate: e.uninit(t * n_ff_max)?,
518                up: e.uninit(t * n_ff_max)?,
519                ffn_out: e.uninit(t * n_embd)?,
520                seg_glue: Vec::new(),
521                mixed: e.uninit(t * n_embd)?,
522                seg_mid: Vec::new(),
523                seg_t: 0,
524            });
525        }
526        Ok(g)
527    }
528
529    fn prime_chunk(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
530                       -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
531        let cfg = &self.cfg;
532        let n_embd = cfg.n_embd as usize;
533        let t = tokens.len();
534        let eps = cfg.rms_eps;
535        let base = cache.pos;
536        let pos: Vec<i32> = (base as i32..(base + t) as i32).collect();
537        let pos_d = e.htod_i32(&pos)?;
538
539        let x_embed = self.embed(e, tokens)?;   // [T, n_embd]
540        // task #14: fuse the fp16 GEMM-operand emission into the trunk norms (kills the
541        // standalone convert launches). Only when the f16 lane serves and T reaches the
542        // GEMM tier; bit-identical either way.
543        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
544        // PRIME SLABS (piecewise foundation): trunk transients in resident buffers —
545        // ~224 fewer alloc/free calls per prime and FROZEN Lt operand addresses
546        // (nvjet's alignment-variant selection becomes run-to-run stable). Every slab is
547        // fully overwritten before use; x ping-pongs xa<->xb; the hidden-stack return
548        // clones the final x (the slab cannot leave). Non-slab fallback: MEMRA_PRIME_SLABS=0.
549        let n_ff_max = self.layers.iter().map(|l| match &l.ffn {
550            crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
551            _ => n_embd,
552        }).max().unwrap_or(n_embd).max(n_embd);
553        let use_slabs = std::env::var("MEMRA_PRIME_SLABS").as_deref() != Ok("0");
554        let mut slab_guard = if use_slabs {
555            Some(self.prime_slabs_get(e, t, n_embd, n_ff_max)?)
556        } else {
557            None
558        };
559        let mut x_own;   // fallback storage when slabs are off
560        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>);
561        let (mut x_cur, mut x_nxt, sl): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, Option<SlabRefs>);
562        let mut seg: Option<(&mut Vec<Option<cudarc::driver::CudaGraph>>, &mut Vec<Option<cudarc::driver::CudaGraph>>, &mut CudaSlice<f32>, &mut usize)> = None;
563        let mut x_own2;
564        match slab_guard.as_mut() {
565            Some(g) => {
566                let slabs = g.as_mut().unwrap();
567                e.copy_into(&mut slabs.xa, 0, &x_embed, t * n_embd)?;
568                let PrimeSlabs { xa, xb, h, x1, z, act, h16, z16, gate, up, ffn_out, seg_glue, mixed, seg_mid, seg_t, .. } = slabs;
569                x_cur = xa;
570                x_nxt = xb;
571                seg = Some((seg_glue, seg_mid, mixed, seg_t));
572                sl = Some((h, x1, z, act, h16, z16, gate, up, ffn_out));
573            }
574            None => {
575                x_own = x_embed;
576                x_own2 = e.uninit(t * n_embd)?;
577                x_cur = &mut x_own;
578                x_nxt = &mut x_own2;
579                sl = None;
580            }
581        }
582        let mut alloc_h; let mut alloc_x1; let mut alloc_z; let mut alloc_act;
583        let mut alloc_h16; let mut alloc_z16;
584        let mut alloc_gate; let mut alloc_up; let mut alloc_fo;
585        let (h, x1, z, act): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
586        let (h16, z16): (&mut CudaSlice<u8>, &mut CudaSlice<u8>);
587        let (sl_gate, sl_up, sl_fo): (&mut CudaSlice<f32>, &mut CudaSlice<f32>, &mut CudaSlice<f32>);
588        match sl {
589            Some((a, b, c, d, e16, f16b, g, u, fo)) => {
590                h = a; x1 = b; z = c; act = d; h16 = e16; z16 = f16b;
591                sl_gate = g; sl_up = u; sl_fo = fo;
592            }
593            None => {
594                alloc_h = e.uninit(t * n_embd)?;
595                alloc_x1 = e.uninit(t * n_embd)?;
596                alloc_z = e.uninit(t * n_embd)?;
597                alloc_act = e.uninit(t * n_ff_max)?;
598                alloc_h16 = e.alloc_u8_uninit(t * n_embd * 2)?;
599                alloc_z16 = e.alloc_u8_uninit(t * n_embd * 2)?;
600                alloc_gate = e.uninit(t * n_ff_max)?;
601                alloc_up = e.uninit(t * n_ff_max)?;
602                alloc_fo = e.uninit(t * n_embd)?;
603                h = &mut alloc_h; x1 = &mut alloc_x1; z = &mut alloc_z; act = &mut alloc_act;
604                h16 = &mut alloc_h16; z16 = &mut alloc_z16;
605                sl_gate = &mut alloc_gate; sl_up = &mut alloc_up; sl_fo = &mut alloc_fo;
606            }
607        }
608        // piecewise increment 3: layer-0 norm hoisted; the per-layer tail fuses
609        // [down-add + NEXT layer's attn-norm] into one captured S-glue segment
610        // (all-slab IO, zero in-graph allocations). Capture happens lazily on the
611        // first prime at this t (capture does not execute -> launch right after).
612        let n_layers = self.layers.len();
613        // OPT-IN (2026-07-26 interleaved verdict): 2-kernel segments measured NET
614        // -0.7%-to-neutral under the interleaved A/B protocol — one cuGraphLaunch costs
615        // about what two kernel submissions do. The earlier "+0.9%" was cross-run clock
616        // drift (the repo's interleaved-A/B law exists for exactly this). Larger segments
617        // (S-prep/S-attn, 7-9 kernels) remain the open hypothesis; the core-split
618        // machinery stays (byte-identical) as their foundation.
619        let use_seg = f16fuse && seg.is_some()
620            && std::env::var("MEMRA_PRIME_SEG").as_deref() == Ok("1");
621        if let Some((sg, sm, _, st)) = seg.as_mut() {
622            if **st != t {
623                sg.clear();
624                sg.extend((0..n_layers).map(|_| None));
625                sm.clear();
626                sm.extend((0..n_layers).map(|_| None));
627                **st = t;
628            }
629        }
630        {
631            let layer0 = &self.layers[0];
632            if f16fuse {
633                e.rms_norm_f16out(x_cur, layer0.attn_norm.float_data(), h, h16, n_embd, t, eps)?;
634            } else {
635                e.rms_norm(x_cur, layer0.attn_norm.float_data(), h, n_embd, t, eps)?;
636            }
637        }
638        for (il, layer) in self.layers.iter().enumerate() {
639            let hx16 = if f16fuse { Some(&*h16) } else { None };
640            if use_seg {
641                // core-split path: projections -> _inner core -> out-GEMM INTO the mixed
642                // slab (no copies) -> S-mid segment [add + post-norm] as one graph launch.
643                let (pre, pre16, w_out) = match &layer.mixer {
644                    Mixer::Full(fa) => {
645                        let g3 = match hx16 {
646                            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
647                            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
648                        };
649                        let (pre, pre16) = self.full_attn_prime_core_inner(e, fa, g3, &pos_d, t, cache, il)?;
650                        (pre, pre16, &fa.wo)
651                    }
652                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
653                    Mixer::Linear(la) => {
654                        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
655                        let g4 = match hx16 {
656                            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
657                            None => e.matmul_group(&ws, h, t)?,
658                        };
659                        let (pre, pre16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, None)?;
660                        (pre, pre16, &la.ssm_out)
661                    }
662                };
663                {
664                    let (_, sm, mslab, _) = seg.as_mut().unwrap();
665                    let pre_n = pre.len() / t;
666                    let xh_pre = match pre16 {
667                        Some(x) => x,
668                        None => e.f16_act(&pre, t * pre_n, pre_n)?,
669                    };
670                    if !e.try_f16_gemm_pre_into(w_out, &xh_pre, t, mslab)? {
671                        let y = e.matmul(w_out, &pre, t)?;
672                        e.copy_into(mslab, 0, &y, t * n_embd)?;
673                    }
674                    if sm[il].is_none() {
675                        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
676                        let w_post = layer.post_attn_norm.float_data();
677                        e.stream().synchronize()?;
678                        e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
679                        let r = (|| -> Result<(), Box<dyn std::error::Error>> {
680                            e.add(x_cur, mslab, x1, t * n_embd)?;
681                            e.rms_norm_f16out(x1, w_post, z, z16, n_embd, t, eps)?;
682                            Ok(())
683                        })();
684                        let g = e.stream().end_capture(
685                            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
686                        r?;
687                        sm[il] = Some(g?.ok_or("S-mid capture produced no graph")?);
688                    }
689                    sm[il].as_ref().unwrap().launch()?;
690                }
691            } else {
692                let mixed = match &layer.mixer {
693                    Mixer::Full(fa) => self.full_attn_prime(e, fa, h, hx16, &pos_d, t, cache, il)?,
694                    Mixer::Linear(la) => self.linear_attn_prime(e, la, h, hx16, t, cache, il)?,
695                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
696                };
697                if f16fuse {
698                    // round 28: residual+norm in ONE kernel (add_rms_norm precedent,
699                    // bit-identical) — the standalone add pass disappears.
700                    e.add_rms_norm_f16out(x_cur, &mixed, layer.post_attn_norm.float_data(),
701                                          x1, z, z16, n_embd, t, eps)?;
702                } else {
703                    e.add(x_cur, &mixed, x1, t * n_embd)?;
704                    e.rms_norm(x1, layer.post_attn_norm.float_data(), z, n_embd, t, eps)?;
705                }
706            }
707            let zx16 = if f16fuse { Some(&*z16) } else { None };
708            match &layer.ffn {
709                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
710                    let n_ff = ffn_gate.out_features();
711                    // gate/up INTO boundary slabs (piecewise increment 2); fall back to
712                    // the allocating group + copy when a mirror is missing.
713                    let mut into_ok = false;
714                    if let Some(xh) = zx16 {
715                        into_ok = e.try_f16_gemm_pre_into(ffn_gate, xh, t, sl_gate)?
716                            && e.try_f16_gemm_pre_into(ffn_up, xh, t, sl_up)?;
717                    }
718                    if !into_ok {
719                        let mut g2 = match zx16 {
720                            Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], z, xh, t)?,
721                            None => e.matmul_group(&[ffn_gate, ffn_up], z, t)?,
722                        };
723                        let up_y = g2.pop().unwrap();
724                        let gate_y = g2.pop().unwrap();
725                        e.copy_into(sl_gate, 0, &gate_y, t * n_ff)?;
726                        e.copy_into(sl_up, 0, &up_y, t * n_ff)?;
727                    }
728                    // task #17: the silu arm's f16out twin emits the down GEMM's fp16
729                    // operand in-epilogue; non-silu activations keep the standalone convert.
730                    let act16 = if Self::f16out_on(e, t) && self.cfg.m3.is_none() {
731                        let mut a16 = e.alloc_u8_uninit(t * n_ff * 2)?;
732                        e.silu_mul_f16out(sl_gate, sl_up, act, &mut a16, t * n_ff)?;
733                        Some(a16)
734                    } else {
735                        Self::ffn_act(e, &self.cfg, sl_gate, sl_up, act, t * n_ff)?;
736                        None
737                    };
738                    // down GEMM into the ffn_out slab (f16 arm; fallback copies)
739                    let xh_act = match act16 {
740                        Some(x) => x,
741                        None => e.f16_act(act, t * n_ff, n_ff)?,
742                    };
743                    if !e.try_f16_gemm_pre_into(ffn_down, &xh_act, t, sl_fo)? {
744                        let y = e.matmul(ffn_down, &*act, t)?;
745                        e.copy_into(sl_fo, 0, &y, t * n_embd)?;
746                    }
747                }
748                crate::hybrid::Ffn::Moe(m) => {
749                    let y = self.moe_ffn_il(e, m, z, t, il as u16)?;
750                    e.copy_into(sl_fo, 0, &y, t * n_embd)?;
751                }
752            }
753            if use_seg && il + 1 < n_layers {
754                // S-glue segment: [add + next attn-norm(+f16out)] — one cuGraphLaunch
755                let w_next = self.layers[il + 1].attn_norm.float_data();
756                let (sg, _, _, _) = seg.as_mut().unwrap();
757                if sg[il].is_none() {
758                    use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
759                    e.stream().synchronize()?;
760                    e.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
761                    let r = (|| -> Result<(), Box<dyn std::error::Error>> {
762                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
763                        e.rms_norm_f16out(x_nxt, w_next, h, h16, n_embd, t, eps)?;
764                        Ok(())
765                    })();
766                    let g = e.stream().end_capture(
767                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
768                    r?;
769                    sg[il] = Some(g?.ok_or("S-glue capture produced no graph")?);
770                }
771                sg[il].as_ref().unwrap().launch()?;
772            } else {
773                if il + 1 < n_layers {
774                    let w_next = self.layers[il + 1].attn_norm.float_data();
775                    if f16fuse {
776                        e.add_rms_norm_f16out(x1, sl_fo, w_next, x_nxt, h, h16, n_embd, t, eps)?;
777                    } else {
778                        e.add(x1, sl_fo, x_nxt, t * n_embd)?;
779                        e.rms_norm(x_nxt, w_next, h, n_embd, t, eps)?;
780                    }
781                } else {
782                    e.add(x1, sl_fo, x_nxt, t * n_embd)?;
783                }
784            }
785            // CHUNK-INVARIANCE BISECT SEAM (lane/chunk-invariance, 2026-08-05): dump the
786            // per-layer post-residual hidden for the LAST row of this chunk, keyed by
787            // absolute position, so two runs at different MEMRA_PRIME_CHUNK can be diffed
788            // layer-by-layer to find the FIRST diverging layer. Diagnostic only —
789            // unset (the default) costs one OnceLock read per layer.
790            if let Some(path) = Self::prime_trace_path() {
791                let row = (base + t - 1) as usize;
792                let host = e.dtoh(x_nxt)?;
793                let last = &host[(t - 1) * n_embd..t * n_embd];
794                use std::io::Write as _;
795                let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
796                let mut h64: u64 = 0xcbf29ce484222325;
797                for v in last {
798                    h64 ^= v.to_bits() as u64;
799                    h64 = h64.wrapping_mul(0x100000001b3);
800                }
801                writeln!(f, "{{\"pos\":{row},\"layer\":{il},\"t\":{t},\"base\":{base},\
802                             \"hash\":\"{h64:016x}\",\"v0\":{:.9e},\"v1\":{:.9e},\"v2\":{:.9e}}}",
803                         last[0], last[1], last[2])?;
804            }
805            std::mem::swap(&mut x_cur, &mut x_nxt);
806        }
807        // hidden-stack return: clone the final x out of the slab
808        let mut x = e.uninit(t * n_embd)?;
809        e.copy_into(&mut x, 0, x_cur, t * n_embd)?;
810        drop(slab_guard);
811
812        // h_seed = LAST row of x BEFORE output_norm (MTP-PLAN §A default) or AFTER it
813        // (MEMRA_SPEC_HPOST — reference convention; hn is computed just below either way, so
814        // the post-norm copy happens after hn exists).
815        let mut h_seed = e.uninit(n_embd)?;
816        if !crate::spec::spec_hpost() {
817            e.copy_view_into(&mut h_seed, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
818        }
819        // last-row logits, exactly like forward_last (norm all T — per-row op — then lm_head on 1 row).
820        let mut hn = e.uninit(t * n_embd)?;
821        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
822        if crate::spec::spec_hpost() {
823            e.copy_view_into(&mut h_seed, 0, &hn.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
824        }
825        let last = e.view(&hn, t * n_embd);
826        let last_row = last.slice((t - 1) * n_embd..t * n_embd);
827        let mut hlast = e.uninit(n_embd)?;
828        e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
829        let logits = e.matmul(&self.output, &hlast, 1)?;
830        cache.pos += t;
831        // Hidden stack handed to generate_spec as prompt_h: pre-norm x (default) or the full
832        // post-norm stack hn (MEMRA_SPEC_HPOST).
833        Ok((e.dtoh(&logits)?, h_seed, if crate::spec::spec_hpost() { hn } else { x }))
834    }
835
836    /// CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2):
837    /// prime_chunk's layer stack with every capture hazard hoisted — `x` is the
838    /// PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay
839    /// eager, one launch), `pos_d` is a baked device param (fresh prime = 0..T, constant
840    /// per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is
841    /// NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk
842    /// (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len
843    /// bookkeeping still runs on the host per call — the real replay path moves the write
844    /// slot to the len_d device counter (increment 3).
845    /// GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers
846    /// (`logits_out` [n_vocab], `h_seed_out` [n_embd]) — every internal allocation drops
847    /// INSIDE the capture region (alloc+free node pairs). Retaining an in-capture
848    /// allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding
849    /// 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn't survive a launch
850    /// anyway — the decode GraphSession's pre-allocated-output pattern is the law here.
851    pub fn prime_chunk_captured(&self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>,
852                                t: usize, cache: &mut Cache,
853                                len_d: &CudaSlice<i32>,
854                                logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>)
855                                -> Result<(), Box<dyn std::error::Error>> {
856        let cfg = &self.cfg;
857        let n_embd = cfg.n_embd as usize;
858        let eps = cfg.rms_eps;
859        let f16fuse = crate::f16_ffi::pp_f16_enabled() && t >= 16;
860        let mut x = e.uninit(t * n_embd)?;
861        e.copy_into(&mut x, 0, x_in, t * n_embd)?;
862        for (il, layer) in self.layers.iter().enumerate() {
863            let mut h = e.uninit(t * n_embd)?;
864            let mut hx16: Option<CudaSlice<u8>> = None;
865            if f16fuse {
866                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
867                e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut b16, n_embd, t, eps)?;
868                hx16 = Some(b16);
869            } else {
870                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
871            }
872            let mixed = match &layer.mixer {
873                Mixer::Full(fa) => self.full_attn_prime(e, fa, &h, hx16.as_ref(), pos_d, t, cache, il)?,
874                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
875                Mixer::Linear(la) => {
876                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
877                    let g4 = match hx16.as_ref() {
878                        Some(xh) => e.matmul_group_xh(&ws, &h, xh, t)?,
879                        None => e.matmul_group(&ws, &h, t)?,
880                    };
881                    self.linear_attn_prime_core_pad(e, la, g4, t, cache, il, Some(len_d))?
882                }
883            };
884            let mut x1 = e.uninit(t * n_embd)?;
885            e.add(&x, &mixed, &mut x1, t * n_embd)?;
886            let mut z = e.uninit(t * n_embd)?;
887            let mut zx16: Option<CudaSlice<u8>> = None;
888            if f16fuse {
889                let mut b16 = e.alloc_u8_uninit(t * n_embd * 2)?;
890                e.rms_norm_f16out(&x1, layer.post_attn_norm.float_data(), &mut z, &mut b16, n_embd, t, eps)?;
891                zx16 = Some(b16);
892            } else {
893                e.rms_norm(&x1, layer.post_attn_norm.float_data(), &mut z, n_embd, t, eps)?;
894            }
895            let ffn_out = match &layer.ffn {
896                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
897                    let n_ff = ffn_gate.out_features();
898                    let mut g2 = match &zx16 {
899                        Some(xh) => e.matmul_group_xh(&[ffn_gate, ffn_up], &z, xh, t)?,
900                        None => e.matmul_group(&[ffn_gate, ffn_up], &z, t)?,
901                    };
902                    let up = g2.pop().unwrap();
903                    let gate = g2.pop().unwrap();
904                    let mut act = e.uninit(t * n_ff)?;
905                    Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, t * n_ff)?;
906                    e.matmul(ffn_down, &act, t)?
907                }
908                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
909            };
910            let mut x2 = e.uninit(t * n_embd)?;
911            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
912            x = x2;
913        }
914        // device-indexed TRUE last row (pads sit past it in a bucketed graph)
915        if !crate::spec::spec_hpost() {
916            e.row_gather_dev(&x, h_seed_out, len_d, n_embd)?;
917        }
918        let mut hn = e.uninit(t * n_embd)?;
919        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
920        if crate::spec::spec_hpost() {
921            e.row_gather_dev(&hn, h_seed_out, len_d, n_embd)?;
922        }
923        let mut hlast = e.uninit(n_embd)?;
924        e.row_gather_dev(&hn, &mut hlast, len_d, n_embd)?;
925        let logits = e.matmul(&self.output, &hlast, 1)?;
926        let nv = logits.len();
927        e.copy_into(logits_out, 0, &logits, nv)?;
928        Ok(())
929    }
930
931    /// Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the
932    /// trunk's token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on
933    /// the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win
934    /// the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer
935    /// CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection
936    /// buffers (D2D row copies; the mixers' own out-projections stay per-seq this
937    /// increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs
938    /// batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the
939    /// per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view
940    /// over the quantized past; Linear: the stateful pad_view twin — the same state
941    /// carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths
942    /// stay byte-identical (gated on !carried). gemma4 models have no continuation
943    /// prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls
944    /// back to single-chunk serving).
945    /// NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class
946    /// as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).
947    pub fn prime_cache_batch(&self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache])
948                             -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
949        let cfg = &self.cfg;
950        let n_embd = cfg.n_embd as usize;
951        let eps = cfg.rms_eps;
952        let b = prompts.len();
953        assert!(b >= 1 && b == caches.len());
954        let pos0s: Vec<usize> = caches.iter().map(|c| c.pos).collect();
955        let carried = pos0s.iter().any(|&p| p > 0);
956        if carried && cfg.gemma4.is_some() {
957            return Err("prime_cache_batch: gemma4 has no continuation prime (v0 fresh-only)".into());
958        }
959        let ts: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
960        for &t in &ts { assert!(t >= PRIME_MIN_T, "prime_cache_batch needs T >= {PRIME_MIN_T}"); }
961        for (s, c) in caches.iter().enumerate() {
962            assert!(c.pos + ts[s] <= c.max_ctx, "prime_cache_batch: prompt exceeds cache max_ctx");
963        }
964        let total: usize = ts.iter().sum();
965        let offs: Vec<usize> = ts.iter().scan(0usize, |a, &t| { let o = *a; *a += t; Some(o) }).collect();
966        // per-seq positions (fresh: 0..T_s; continuation: pos0..pos0+T_s)
967        let pos_ds: Vec<CudaSlice<i32>> = ts.iter().zip(&pos0s)
968            .map(|(&t, &p0)| e.htod_i32(&(p0 as i32..(p0 + t) as i32).collect::<Vec<_>>()))
969            .collect::<Result<_, _>>()?;
970        // split a concat [total, dim] buffer into per-seq copies
971        let split = |e: &Engine, y: &CudaSlice<f32>, dim: usize|
972                     -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
973            let mut out = Vec::with_capacity(b);
974            for s in 0..b {
975                let mut ys = e.uninit(ts[s] * dim)?;
976                e.copy_view_into(&mut ys, 0, &y.slice(offs[s] * dim..(offs[s] + ts[s]) * dim), ts[s] * dim)?;
977                out.push(ys);
978            }
979            Ok(out)
980        };
981
982        let cat_tokens: Vec<u32> = prompts.iter().flat_map(|p| p.iter().copied()).collect();
983        let mut x = self.embed(e, &cat_tokens)?;   // [total, n_embd]
984        for (il, layer) in self.layers.iter().enumerate() {
985            let mut h = e.uninit(total * n_embd)?;
986            let mut hx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
987            e.rms_norm_f16out(&x, layer.attn_norm.float_data(), &mut h, &mut hx16, n_embd, total, eps)?;
988            // mixer: projection GROUP on the concat (m = total), stateful core per seq
989            let mut mixed = e.uninit(total * n_embd)?;
990            match &layer.mixer {
991                Mixer::Full(fa) => {
992                    let g3 = e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], &h, &hx16, total)?;
993                    // task #18 (attn side): the WHOLE attn core is varlen for fresh gated
994                    // batches — split/QK-norm/RoPE/append (attn_pre_vl8, view inputs: the
995                    // q/k/v split copies vanish) + ONE varlen FA. Per-block math identical
996                    // everywhere (bit-gateable); MEMRA_FA_VL=0 or a non-bf16kv config falls
997                    // back to the per-seq dispatch.
998                    let (n_head, n_head_kv, head_dim) =
999                        (self.cfg.n_head as usize, self.cfg.n_head_kv as usize, self.cfg.head_dim_k as usize);
1000                    let fa_scale = 1.0 / (head_dim as f32).sqrt();
1001                    let use_favl = !carried
1002                        && (2..=8).contains(&b)
1003                        && (head_dim == 256 || head_dim == 128)
1004                        && self.cfg.attn_out_gate()
1005                        && std::env::var("MEMRA_NOFA").is_err()
1006                        && std::env::var("MEMRA_FA_FLOOR").is_err()
1007                        && std::env::var("MEMRA_FA_PP_W2").as_deref() != Ok("1")
1008                        && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0")
1009                        && std::env::var("MEMRA_FA_VL").as_deref() != Ok("0");
1010                    if use_favl {
1011                        let (qf_w, kf_w, vf_w) =
1012                            (fa.wq.out_features(), fa.wk.out_features(), fa.wv.out_features());
1013                        struct APre {
1014                            q: CudaSlice<f32>, gate: Option<CudaSlice<f32>>,
1015                            qn: CudaSlice<f32>, kn: CudaSlice<f32>,
1016                        }
1017                        let mut aps = Vec::with_capacity(b);
1018                        for &t in ts.iter().take(b) {
1019                            aps.push(APre {
1020                                q: e.uninit(t * n_head * head_dim)?,
1021                                gate: Some(e.uninit(t * n_head * head_dim)?),
1022                                qn: e.uninit(t * n_head * head_dim)?,
1023                                kn: e.uninit(t * n_head_kv * head_dim)?,
1024                            });
1025                        }
1026                        let (kv_dim_k, kv_dim_v, ktb, vtb) = {
1027                            let kvl = caches[0].kv[il].as_ref().unwrap();
1028                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
1029                        };
1030                        let pargs: Vec<crate::AttnPreVl> = (0..b).map(|s| {
1031                            let (o, t) = (offs[s], ts[s]);
1032                            let kvl = caches[s].kv[il].as_ref().unwrap();
1033                            assert!(kvl.len == 0 && kvl.len + t <= caches[s].max_ctx,
1034                                    "prime_cache_batch attn vl: fresh + capacity");
1035                            crate::AttnPreVl {
1036                                qf: e.addr_f32v(&g3[0].slice(o * qf_w..(o + t) * qf_w)),
1037                                kf: e.addr_f32v(&g3[1].slice(o * kf_w..(o + t) * kf_w)),
1038                                vf: e.addr_f32v(&g3[2].slice(o * vf_w..(o + t) * vf_w)),
1039                                q: e.addr_f32(&aps[s].q),
1040                                gate: e.addr_f32(aps[s].gate.as_ref().unwrap()),
1041                                qn: e.addr_f32(&aps[s].qn), kn: e.addr_f32(&aps[s].kn),
1042                                kc: e.addr_u8(&kvl.k), vc: e.addr_u8(&kvl.v),
1043                                t: t as i32, pad: 0,
1044                            }
1045                        }).collect();
1046                        e.attn_pre_vl8(&pargs, fa.q_norm.float_data(), fa.k_norm.float_data(),
1047                                       head_dim, self.cfg.rope_dim_count as usize, n_head, n_head_kv,
1048                                       self.cfg.rms_eps, self.cfg.rope_freq_base, 1.0,
1049                                       kv_dim_k, kv_dim_v, ktb, vtb)?;
1050                        for s in 0..b {
1051                            let kvl = caches[s].kv[il].as_mut().unwrap();
1052                            kvl.len += ts[s];
1053                            let new_len = kvl.len as i32;
1054                            e.set_i32_one(&mut kvl.len_d, new_len)?;
1055                        }
1056                        let mut attns = Vec::with_capacity(b);
1057                        let mut mirrors = Vec::with_capacity(b);
1058                        for &t in ts.iter().take(b) {
1059                            attns.push(e.uninit(t * n_head * head_dim)?);
1060                            let n = t * n_head_kv * head_dim;
1061                            mirrors.push((e.alloc_u8_uninit(n * 2)?, e.alloc_u8_uninit(n * 2)?));
1062                        }
1063                        // FA3 batched twin (round 31): TMA-swizzled wgmma vl when the
1064                        // promoted single-seq config is on; else the mma favl.
1065                        let fa3_on = match std::env::var("MEMRA_FA3").as_deref() {
1066                            Ok("0") => false,
1067                            Ok("1") => true,
1068                            _ => cfg!(memra_hopper_mma),
1069                        };
1070                        if fa3_on {
1071                            let mut q16s = Vec::with_capacity(b);
1072                            let mut v16s = Vec::with_capacity(b);
1073                            for s in 0..b {
1074                                let t = ts[s];
1075                                let mut q16 = e.alloc_u8_uninit(t * n_head * head_dim * 2)?;
1076                                e.f32_to_bf16_into(&aps[s].qn, &mut q16, t * n_head * head_dim)?;
1077                                let mut k16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1078                                e.f32_to_bf16_into(&aps[s].kn, &mut k16, t * n_head_kv * head_dim)?;
1079                                let mut v16 = e.alloc_u8_uninit(t * n_head_kv * head_dim * 2)?;
1080                                e.f32_to_bf16_v(&g3[2].slice(offs[s] * vf_w..(offs[s] + t) * vf_w),
1081                                                &mut v16, t * n_head_kv * head_dim)?;
1082                                q16s.push(q16);
1083                                v16s.push((k16, v16));
1084                            }
1085                            let mut qp = [core::ptr::null::<core::ffi::c_void>(); 8];
1086                            let mut kp = qp;
1087                            let mut vp = qp;
1088                            let mut op = [core::ptr::null_mut::<f32>(); 8];
1089                            let mut tsv = [0i32; 8];
1090                            for s in 0..b {
1091                                qp[s] = e.addr_u8(&q16s[s]) as *const core::ffi::c_void;
1092                                kp[s] = e.addr_u8(&v16s[s].0) as *const core::ffi::c_void;
1093                                vp[s] = e.addr_u8(&v16s[s].1) as *const core::ffi::c_void;
1094                                op[s] = e.addr_f32(&attns[s]) as *mut f32;
1095                                tsv[s] = ts[s] as i32;
1096                            }
1097                            let rc = unsafe {
1098                                crate::fa3_vl_raw(qp.as_ptr(), kp.as_ptr(), vp.as_ptr(), op.as_ptr(),
1099                                                  tsv.as_ptr(), b as i32, n_head as i32,
1100                                                  n_head_kv as i32, head_dim as i32, fa_scale,
1101                                                  e.stream().cu_stream() as *mut core::ffi::c_void)
1102                            };
1103                            if rc != 0 {
1104                                return Err(format!("memra_fa3_vl rc={rc}").into());
1105                            }
1106                        } else {
1107                            let fargs: Vec<crate::FaSeqVl> = (0..b).map(|s| crate::FaSeqVl {
1108                                q: e.addr_f32(&aps[s].qn), k16: e.addr_u8(&mirrors[s].0),
1109                                v16: e.addr_u8(&mirrors[s].1), o: e.addr_f32(&attns[s]),
1110                                kf: e.addr_f32(&aps[s].kn),
1111                                vf: e.addr_f32v(&g3[2].slice(offs[s] * vf_w..(offs[s] + ts[s]) * vf_w)),
1112                                t: ts[s] as i32, pad: 0,
1113                            }).collect();
1114                            e.fa_prefill_vl8(&fargs, head_dim, n_head, n_head_kv, fa_scale)?;
1115                        }
1116                        for (s, attn) in attns.into_iter().enumerate() {
1117                            let (attn_g, ag16) = self.full_attn_prime_post_fa(
1118                                e, attn, &aps[s].gate, ts[s], n_head, head_dim)?;
1119                            let mut done = false;
1120                            if let Some(xh) = &ag16 {
1121                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1122                            }
1123                            if !done {
1124                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1125                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1126                            }
1127                        }
1128                    } else {
1129                        let mut parts: Vec<Vec<CudaSlice<f32>>> = (0..b).map(|_| Vec::new()).collect();
1130                        for (w, y) in [&fa.wq, &fa.wk, &fa.wv].iter().zip(g3) {
1131                            for (s, ys) in split(e, &y, w.out_features())?.into_iter().enumerate() {
1132                                parts[s].push(ys);
1133                            }
1134                        }
1135                        for (s, g3s) in parts.into_iter().enumerate() {
1136                            // task #16 gather removal: wo writes into `mixed` at offs[s] directly.
1137                            let (attn_g, ag16) = self.full_attn_prime_core_inner(
1138                                e, fa, g3s, &pos_ds[s], ts[s], caches[s], il)?;
1139                            let mut done = false;
1140                            if let Some(xh) = &ag16 {
1141                                done = e.try_f16_gemm_pre_into_off(&fa.wo, xh, ts[s], &mut mixed, offs[s] * n_embd)?;
1142                            }
1143                            if !done {
1144                                let m = e.matmul(&fa.wo, &attn_g, ts[s])?;
1145                                e.copy_into(&mut mixed, offs[s] * n_embd, &m, ts[s] * n_embd)?;
1146                            }
1147                        }
1148                    }
1149                }
1150                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1151                Mixer::Linear(la) => {
1152                    // task #16: NO split copies (cores read row-offset views of the concat
1153                    // outputs; out-GEMMs write into `mixed` at offs[s]). task #18: the core
1154                    // itself is BATCHED — per-seq prep/K1-K3, then ONE varlen K4 + ONE
1155                    // varlen K5 launch for all sequences.
1156                    let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1157                    let g4 = e.matmul_group_xh(&ws, &h, &hx16, total)?;
1158                    let outs = self.linear_attn_prime_core_batch(e, la, &g4, &offs, &ts, caches, il)?;
1159                    for (s, (gn, gn16)) in outs.into_iter().enumerate() {
1160                        let (o, t) = (offs[s], ts[s]);
1161                        let mut done = false;
1162                        if let Some(xh) = &gn16 {
1163                            done = e.try_f16_gemm_pre_into_off(&la.ssm_out, xh, t, &mut mixed, o * n_embd)?;
1164                        }
1165                        if !done {
1166                            let m = e.matmul(&la.ssm_out, &gn, t)?;
1167                            e.copy_into(&mut mixed, o * n_embd, &m, t * n_embd)?;
1168                        }
1169                    }
1170                }
1171            }
1172            let mut x1 = e.uninit(total * n_embd)?;
1173            let mut z = e.uninit(total * n_embd)?;
1174            let mut zx16 = e.alloc_u8_uninit(total * n_embd * 2)?;
1175            e.add_rms_norm_f16out(&x, &mixed, layer.post_attn_norm.float_data(),
1176                                  &mut x1, &mut z, &mut zx16, n_embd, total, eps)?;
1177            let ffn_out = match &layer.ffn {
1178                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
1179                    let n_ff = ffn_gate.out_features();
1180                    let mut g2 = e.matmul_group_xh(&[ffn_gate, ffn_up], &z, &zx16, total)?;
1181                    let up = g2.pop().unwrap();
1182                    let gate = g2.pop().unwrap();
1183                    let mut act = e.uninit(total * n_ff)?;
1184                    // task #17 (batch trunk): silu twin emits the down GEMM's fp16 operand
1185                    // in-epilogue (nsys round-26: this trunk still paid 32 cvt passes).
1186                    if Self::f16out_on(e, total) && self.cfg.m3.is_none() {
1187                        let mut a16 = e.alloc_u8_uninit(total * n_ff * 2)?;
1188                        e.silu_mul_f16out(&gate, &up, &mut act, &mut a16, total * n_ff)?;
1189                        match e.try_f16_gemm_pre(ffn_down, &a16, total)? {
1190                            Some(y) => y,
1191                            None => e.matmul(ffn_down, &act, total)?,
1192                        }
1193                    } else {
1194                        Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, total * n_ff)?;
1195                        e.matmul(ffn_down, &act, total)?
1196                    }
1197                }
1198                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, total, il as u16)?,
1199            };
1200            let mut x2 = e.uninit(total * n_embd)?;
1201            e.add(&x1, &ffn_out, &mut x2, total * n_embd)?;
1202            x = x2;
1203        }
1204        // epilogue per seq (identical math to prime_chunk: norm all rows, lm_head on last row)
1205        let mut hn = e.uninit(total * n_embd)?;
1206        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, total, eps)?;
1207        // batched lm_head (nsys round-26: B sequential m=1 matvecs re-read the 600MB
1208        // lm_head weight B times — 2.2ms at B=6): gather the B last rows and run ONE
1209        // m=B GEMM (f16 lane; falls back to the per-seq matvec when no mirror). The
1210        // f16-vs-mmvq logits delta is the prefill GEMM numeric class — prime_batch_gate's
1211        // argmax battery arbitrates, same as every other prefill GEMM change.
1212        let mut hcat = e.uninit(b * n_embd)?;
1213        for s in 0..b {
1214            let last0 = (offs[s] + ts[s] - 1) * n_embd;
1215            e.copy_view_into(&mut hcat, s * n_embd, &hn.slice(last0..last0 + n_embd), n_embd)?;
1216        }
1217        let logits_cat = if b >= 2 { e.try_f16_gemm(&self.output, &hcat, b)? } else { None };
1218        let logits_host: Option<Vec<f32>> = match &logits_cat {
1219            Some(lc) => Some(e.dtoh(lc)?),
1220            None => None,
1221        };
1222        let n_vocab = self.output.out_features();
1223        let mut hidden_all = if crate::spec::spec_hpost() {
1224            split(e, &hn, n_embd)?
1225        } else {
1226            split(e, &x, n_embd)?
1227        };
1228        let mut out = Vec::with_capacity(b);
1229        for s in 0..b {
1230            let last0 = (offs[s] + ts[s] - 1) * n_embd;
1231            let mut h_seed = e.uninit(n_embd)?;
1232            if !crate::spec::spec_hpost() {
1233                e.copy_view_into(&mut h_seed, 0, &x.slice(last0..last0 + n_embd), n_embd)?;
1234            } else {
1235                e.copy_view_into(&mut h_seed, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1236            }
1237            let logits = match &logits_host {
1238                Some(lh) => lh[s * n_vocab..(s + 1) * n_vocab].to_vec(),
1239                None => {
1240                    let mut hlast = e.uninit(n_embd)?;
1241                    e.copy_view_into(&mut hlast, 0, &hn.slice(last0..last0 + n_embd), n_embd)?;
1242                    e.dtoh(&e.matmul(&self.output, &hlast, 1)?)?
1243                }
1244            };
1245            caches[s].pos += ts[s];
1246            out.push((logits, h_seed, hidden_all.remove(0)));
1247        }
1248        Ok(out)
1249    }
1250
1251    /// `full_attn` (batched prefill mixer) + the cache side-effect: append the T post-RoPE K/V
1252    /// rows into the resident quantized KV cache (q8_0 K / q5_1 V) and advance len/len_d. Row
1253    /// bytes are BIT-IDENTICAL to the decode append (same per-warp quant kernel per row; the
1254    /// batched `append_kv_quantized_rows` runs that exact warp math on a (block, token) grid).
1255    /// The attention itself is unchanged prefill math (fa_prefill over the f32 K/V).
1256    fn full_attn_prime(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>,
1257                       hx: Option<&CudaSlice<u8>>,
1258                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1259                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1260        // PROJ/CORE SPLIT (task #13, 2026-07-26): the q/k/v group projection is hoisted so
1261        // the cross-request batch driver can run it at m = sum_T over concatenated tokens;
1262        // this single-seq path composes proj+core identically (byte-for-byte the old body).
1263        // task #14: `hx` = the norm-fused fp16 twin of `h` (skips the convert launch).
1264        let g3 = match hx {
1265            Some(xh) => e.matmul_group_xh(&[&fa.wq, &fa.wk, &fa.wv], h, xh, t)?,
1266            None => e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?,
1267        };
1268        self.full_attn_prime_core(e, fa, g3, pos_d, t, cache, il)
1269    }
1270
1271    /// Everything after the q/k/v projections (split/QK-norm/RoPE/FA/gate/append + wo).
1272    /// `g3` = matmul_group([wq, wk, wv]) output rows for THIS sequence.
1273    /// Composes inner + wo (== the pre-split core; every existing caller unchanged).
1274    fn full_attn_prime_core(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1275                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1276                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1277        let (attn_g, ag16) = self.full_attn_prime_core_inner(e, fa, g3, pos_d, t, cache, il)?;
1278        if let Some(xh) = &ag16 {
1279            if let Some(y) = e.try_f16_gemm_pre(&fa.wo, xh, t)? {
1280                return Ok(y);
1281            }
1282        }
1283        Ok(e.matmul(&fa.wo, &attn_g, t)?)
1284    }
1285
1286    fn full_attn_prime_core_inner(&self, e: &Engine, fa: &FullAttnLayer, g3: Vec<CudaSlice<f32>>,
1287                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1288                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1289        let cfg = &self.cfg;
1290        let n_head = cfg.n_head as usize;
1291        let n_head_kv = cfg.n_head_kv as usize;
1292        let head_dim = cfg.head_dim_k as usize;
1293        let scale = 1.0 / (head_dim as f32).sqrt();
1294        let (pre, base_len) = self.full_attn_prime_pre_fa(e, fa, g3, pos_d, t, cache, il)?;
1295        let AttnPre { q, k, v, gate } = pre;
1296        let mut attn = e.uninit(t * n_head * head_dim)?;
1297        self.full_attn_prime_fa_dispatch(e, &q, &k, &v, &mut attn, base_len, t, cache, il,
1298                                         head_dim, n_head, n_head_kv, scale)?;
1299        self.full_attn_prime_post_fa(e, attn, &gate, t, n_head, head_dim)
1300    }
1301
1302    /// task #18 (attn side): projections tail through KV append — everything before the
1303    /// attention kernel. Returns the post-rope q/k, v, optional out-gate, and the KV rows
1304    /// present BEFORE this chunk's append (base_len; 0 == fresh).
1305    #[allow(clippy::type_complexity)]
1306    fn full_attn_prime_pre_fa(&self, e: &Engine, fa: &FullAttnLayer, mut g3: Vec<CudaSlice<f32>>,
1307                            pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, il: usize)
1308                            -> Result<(AttnPre, usize), Box<dyn std::error::Error>> {
1309        let cfg = &self.cfg;
1310        let n_head = cfg.n_head as usize;
1311        let n_head_kv = cfg.n_head_kv as usize;
1312        let head_dim = cfg.head_dim_k as usize;
1313        let eps = cfg.rms_eps;
1314
1315        // qwen35 fuses [q|gate] per head in wq (2*head_dim stride); M3/Hy3 have NO output gate
1316        // (attention_output_gate=false) — wq out = n_head*head_dim exactly, and q_gate_split
1317        // would read 2x out of bounds. `gated` keys both the split and the sigmoid epilogue.
1318        let gated = cfg.attn_out_gate();
1319        let v = g3.pop().unwrap();
1320        let mut k = g3.pop().unwrap();
1321        let qf = g3.pop().unwrap();
1322        let (mut q, gate) = if gated {
1323            let mut q = e.uninit(t * n_head * head_dim)?;
1324            let mut gate = e.uninit(t * n_head * head_dim)?;
1325            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
1326            (q, Some(gate))
1327        } else {
1328            (qf, None)
1329        };
1330
1331        let mut qn = e.uninit(t * n_head * head_dim)?;
1332        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
1333        q = qn;
1334        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
1335        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
1336        k = kn;
1337        let rope_dims = cfg.rope_dim_count as usize;
1338        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
1339        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
1340
1341        // CACHE SIDE-EFFECT: append the T post-rope K/V token rows (token-major [T, kv_dim] ==
1342        // the cache row layout) quantized into cache.kv[il], then advance len + device len_d.
1343        {
1344            let kvl = cache.kv[il].as_mut().unwrap();
1345            assert!(kvl.len + t <= cache.max_ctx, "prime_cache: KV overflow");
1346            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
1347                                       kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
1348                                       crate::Engine::kv_fp8_on())?;
1349            kvl.len += t;
1350            let new_len = kvl.len as i32;
1351            e.set_i32_one(&mut kvl.len_d, new_len)?;
1352        }
1353
1354        let base_len = {
1355            let kvl = cache.kv[il].as_ref().unwrap();
1356            kvl.len - t   // KV rows present BEFORE this chunk's append above
1357        };
1358        Ok((AttnPre { q, k, v, gate }, base_len))
1359    }
1360
1361    /// batched prefill attention. FRESH prime (no past KV): unchanged forward_last math over
1362    /// the f32 K/V of this batch. CONTINUATION chunk (past KV present): the chunk's queries
1363    /// must attend to [0 .. base+t) — run fa_prefill_view over the resident QUANTIZED cache
1364    /// (the spec-verify pattern; kernel's causal mask offsets by T_kv-T). Numerically this
1365    /// reads q8_0/q5_1-dequantized K/V for the past AND the current chunk — the same class as
1366    /// decode reading the cache; the run-gen/first-16 battery is the accuracy authority.
1367    #[allow(clippy::too_many_arguments)]
1368    fn full_attn_prime_fa_dispatch(&self, e: &Engine, q: &CudaSlice<f32>, k: &CudaSlice<f32>,
1369                            v: &CudaSlice<f32>, attn: &mut CudaSlice<f32>, base_len: usize,
1370                            t: usize, cache: &mut Cache, il: usize,
1371                            head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32)
1372                            -> Result<(), Box<dyn std::error::Error>> {
1373        // GRAIN-FREE CHUNK-INVARIANCE FIX (lane/chunkinv-flip, 2026-08-05): the base_len == 0
1374        // f32 special case (fa_prefill over this batch's f32 K/V) is DROPPED. pre_fa appends
1375        // the chunk's quantized rows into cache.kv[il] BEFORE this dispatch, so chunk 0 can
1376        // attend through the quantized cache exactly like every later chunk (quantize-then-
1377        // attend). One numeric class for every row => the chunk size cannot decide where a
1378        // precision edge falls, so chunked prefill is reduction-order-stable with NO grain
1379        // knob (the stronger fix VERDICT.md filed; supersedes the MEMRA_PRIME_INVARIANT door's
1380        // pin-the-boundary approach).
1381        // MEMRA_PRIME_F32CHUNK0=1 is the ROLLBACK SEAM to the old arithmetic (chunk 0 attends
1382        // f32 K/V) — flags-doctrine rollback door AND the chunkinv gate's canary injection:
1383        // with the fix unconditional, only re-introducing the class edge can prove the gate
1384        // still detects the mechanism. Never on in a measured default run.
1385        if base_len == 0 && std::env::var("MEMRA_PRIME_F32CHUNK0").as_deref() == Ok("1") {
1386            if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1387                e.sdpa_naive(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1388            } else {
1389                e.fa_prefill(q, k, v, attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1390            }
1391            return Ok(());
1392        }
1393        let kvl = cache.kv[il].as_ref().unwrap();
1394        let t_kv = base_len + t;
1395        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
1396        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
1397        // fa_prefill_q/_qw twins are stamped for head_dim 256 (qwen35) and 128 (M3) only;
1398        // other dims (and MEMRA_NOFA) take the naive quantized-view SDPA — SAME cache bytes,
1399        // same numeric class, so the uniform contract holds on the fallback too.
1400        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1401            e.sdpa_naive_quantized_view(q, &k_view, &v_view, attn, head_dim, n_head,
1402                                        n_head_kv, t, t_kv, scale, true,
1403                                        kvl.k_tok_bytes, kvl.v_tok_bytes)?;
1404            return Ok(());
1405        }
1406        // ARC B (2026-07-05): dequant-once workspace, DEFAULT ON. fa_prefill_q's inline
1407        // dequant re-reads+re-dequants the whole quantized KV stream from every one of the
1408        // T/64 x n_head CTAs (64x+ redundant at chunk=4096; 30.5% of the 32k prime wall).
1409        // fa_prefill_view_ws dequants K/V ONCE into a resident bf16 workspace then runs the
1410        // bit-identical bf16 twin (fa_prefill_qw) — same staged values, same FP order, token-
1411        // identical output (gate: MEMRA_PRIME_CHUNK=4096 ws-on vs ws-off vs monolithic).
1412        // MEMRA_PRIME_DEQW=0 reverts to the inline-dequant kernel.
1413        let deqw = std::env::var("MEMRA_PRIME_DEQW").map(|v| v != "0").unwrap_or(true);
1414        if deqw {
1415            e.fa_prefill_view_ws(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1416                                 t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1417                                 crate::Engine::kv_fp8_on())?;
1418        } else {
1419            e.fa_prefill_view(q, &k_view, &v_view, attn, head_dim, n_head, n_head_kv,
1420                              t, t_kv, scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes,
1421                              crate::Engine::kv_fp8_on())?;
1422        }
1423        Ok(())
1424    }
1425
1426    /// task #17: sig_mul_f16out fuses [sigmoid + mul + f16 convert] into one launch
1427    /// (bit-identical composition) and hands wo its fp16 operand directly.
1428    fn full_attn_prime_post_fa(&self, e: &Engine, attn: CudaSlice<f32>,
1429                            gate: &Option<CudaSlice<f32>>, t: usize,
1430                            n_head: usize, head_dim: usize)
1431                            -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1432        let (attn_g, ag16) = match gate {
1433            Some(gate) => {
1434                let n = t * n_head * head_dim;
1435                let mut ag = e.uninit(n)?;
1436                if Self::f16out_on(e, t) {
1437                    let mut a16 = e.alloc_u8_uninit(n * 2)?;
1438                    e.sig_mul_f16out(&attn, gate, &mut ag, &mut a16, n)?;
1439                    (ag, Some(a16))
1440                } else {
1441                    let mut gsig = e.uninit(n)?;
1442                    e.sigmoid(gate, &mut gsig, n)?;
1443                    e.mul(&attn, &gsig, &mut ag, n)?;
1444                    (ag, None)
1445                }
1446            }
1447            None => (attn, None),
1448        };
1449        Ok((attn_g, ag16))
1450    }
1451
1452    /// STATEFUL batched linear-attention prime: `linear_attn`'s prefill-dispatch pass (normal
1453    /// `e.matmul` — GEMM at m>=16 — plus the prefill repack/L2/glog kernels) but with the state
1454    /// carried THROUGH the cache like the spec verify does: carried-ring conv
1455    /// (ssm_conv1d_tm_state writes the final ring back) + ONE gdn_scan from cache.recur[il]'s
1456    /// current state (zero at a fresh prime) whose final state ping-pongs back into the cache.
1457    /// Wiring mirrors `linear_attn_verify_t` (spec.rs); dispatch mirrors `linear_attn` (prefill).
1458    fn linear_attn_prime(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>,
1459                         hx: Option<&CudaSlice<u8>>, t: usize,
1460                         cache: &mut Cache, il: usize)
1461                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1462        // PROJ/CORE SPLIT (task #13): see full_attn_prime — same hoist for the GDN 4-tuple.
1463        let ws = [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha];
1464        let g4 = match hx {
1465            Some(xh) => e.matmul_group_xh(&ws, h, xh, t)?,
1466            None => e.matmul_group(&ws, h, t)?,
1467        };
1468        self.linear_attn_prime_core(e, la, g4, t, cache, il)
1469    }
1470
1471    /// Everything after the GDN 4-tuple projections (conv/chunk stack/gated norm + ssm_out).
1472    fn linear_attn_prime_core(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1473                              t: usize, cache: &mut Cache, il: usize)
1474                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1475        self.linear_attn_prime_core_pad(e, la, g4.drain(..).collect(), t, cache, il, None)
1476    }
1477
1478    /// task #14: `pad_len` = device true length for PADDED prime graphs. Pads become
1479    /// identity GDN steps (gdn_pad_mask zeroes beta/g_log past the true length) and the
1480    /// conv ring writes back from the true tail. None = classic path, byte-identical.
1481    #[allow(clippy::too_many_arguments)]
1482    fn linear_attn_prime_core_pad_inner(&self, e: &Engine, la: &LinearAttnLayer, mut g4: Vec<CudaSlice<f32>>,
1483                              t: usize, cache: &mut Cache, il: usize,
1484                              pad_len: Option<&CudaSlice<i32>>)
1485                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1486        // shim over the view twin (task #16): full-range views of the owned buffers.
1487        let ssm = self.cfg.ssm.as_ref().unwrap();
1488        let d_state = ssm.state_size as usize;
1489        let num_k = ssm.group_count as usize;
1490        let num_v = ssm.time_step_rank as usize;
1491        let key_dim = d_state * num_k;
1492        let value_dim = d_state * num_v;
1493        let conv_dim = key_dim * 2 + value_dim;
1494        let alpha = g4.pop().unwrap();                   // [T, num_v]
1495        let beta_raw = g4.pop().unwrap();                // [T, num_v]
1496        let z = g4.pop().unwrap();                       // [T, value_dim]
1497        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
1498        self.linear_attn_prime_core_pad_view(
1499            e, la,
1500            &qkv_mixed.slice(0..t * conv_dim), &z.slice(0..t * value_dim),
1501            &beta_raw.slice(0..t * num_v), &alpha.slice(0..t * num_v),
1502            t, cache, il, pad_len)
1503    }
1504
1505    /// task #18: the GDN prep stage (conv + repack + l2 x2 + sigmoid + glog + pad-mask) —
1506    /// shared verbatim by the per-seq scan path and the varlen batched path.
1507    #[allow(clippy::too_many_arguments)]
1508    fn linear_attn_gdn_prep(&self, e: &Engine, la: &LinearAttnLayer,
1509                            qkv_mixed: &cudarc::driver::CudaView<f32>,
1510                            beta_raw: &cudarc::driver::CudaView<f32>,
1511                            alpha: &cudarc::driver::CudaView<f32>,
1512                            t: usize, cache: &mut Cache, il: usize,
1513                            pad_len: Option<&CudaSlice<i32>>)
1514                            -> Result<GdnPrep, Box<dyn std::error::Error>> {
1515        let cfg = &self.cfg;
1516        let ssm = cfg.ssm.as_ref().unwrap();
1517        let d_state = ssm.state_size as usize;       // 128
1518        let num_k = ssm.group_count as usize;        // 16
1519        let num_v = ssm.time_step_rank as usize;     // 32
1520        let d_conv = ssm.conv_kernel as usize;       // 4
1521        let key_dim = d_state * num_k;               // 2048
1522        let value_dim = d_state * num_v;             // 4096
1523        let conv_dim = key_dim * 2 + value_dim;      // 8192
1524        let eps = cfg.rms_eps;
1525        debug_assert!(t >= d_conv - 1, "stateful conv needs T >= pad (PRIME_MIN_T gates)");
1526
1527        // conv with CARRIED ring state + ring roll (state read + final-window write-back).
1528        // task #18 conv-fuse (default ON): conv + SiLU + repack in one pass — the 67MB
1529        // conv_out intermediate and its transposed re-read disappear (values bit-identical).
1530        // MEMRA_CONV_FUSE=0 reverts to the two-kernel chain.
1531        let rl = cache.recur[il].as_mut().unwrap();
1532        let hk = Self::gdn_hk(e, t, num_v, num_k);
1533        let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
1534        let hk = if conv_fuse { hk } else { num_v };   // de-broadcast rides the fused conv
1535        let mut q_g = e.uninit(d_state * hk * t)?;
1536        let mut k_g = e.uninit(d_state * hk * t)?;
1537        let mut v_g = e.uninit(d_state * num_v * t)?;
1538        if conv_fuse {
1539            e.ssm_conv1d_gdn_state_pad(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
1540                                  &mut q_g, &mut k_g, &mut v_g,
1541                                  conv_dim, t, d_conv, d_state, num_v, num_k, key_dim, hk, pad_len)?;
1542        } else {
1543            let mut conv_out = e.uninit(conv_dim * t)?;      // [conv_dim, T] channel-major, SiLU
1544            e.ssm_conv1d_tm_state_pad_v(qkv_mixed, &mut rl.conv_state, la.ssm_conv1d.float_data(),
1545                                  &mut conv_out, conv_dim, t, d_conv, pad_len)?;
1546            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)?;
1547        }
1548        let mut q_l2 = e.uninit(d_state * hk * t)?;
1549        // mirror-fold (round 35): q's bf16 twin (wgmma K45/K2 A-operand) in-epilogue too.
1550        // Emitted only where a consumer exists (the wgmma config) — on other arches the
1551        // alloc + epilogue stores would be pure waste.
1552        let qb16 = if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(32) {
1553            let mut qb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1554            e.l2_norm_pp(&q_g, &mut q_l2, Some(&mut qb), d_state, hk * t, eps)?;
1555            Some(qb)
1556        } else {
1557            e.l2_norm_pp(&q_g, &mut q_l2, None, d_state, hk * t, eps)?;
1558            None
1559        };
1560        let mut k_l2 = e.uninit(d_state * hk * t)?;
1561        // mirror-fold: emit k's bf16 twin (the chunked-scan kb16 mirror) in-epilogue
1562        let kb16 = if Engine::l2_v2_on(d_state) {
1563            let mut kb = e.alloc_u8_uninit(d_state * hk * t * 2)?;
1564            e.l2_norm_pp(&k_g, &mut k_l2, Some(&mut kb), d_state, hk * t, eps)?;
1565            Some(kb)
1566        } else {
1567            e.l2_norm_pp(&k_g, &mut k_l2, None, d_state, hk * t, eps)?;
1568            None
1569        };
1570        let mut beta = e.uninit(t * num_v)?;
1571        e.sigmoid_v(beta_raw, &mut beta, t * num_v)?;
1572        let mut g_log = e.uninit(t * num_v)?;
1573        e.gdn_glog_v(alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
1574        if let Some(len_d) = pad_len {
1575            e.gdn_pad_mask(&mut beta, &mut g_log, len_d, num_v, t)?;
1576        }
1577        Ok(GdnPrep { hk, q_l2, k_l2, v_g, beta, g_log, kb16, qb16 })
1578    }
1579
1580    /// task #18: batched GDN mixer core — per-seq prep + K1-K3, then ONE varlen K4 and
1581    /// ONE varlen K5 launch for all B sequences (full-machine occupancy vs B underfilled
1582    /// per-seq trains). Per-block math identical to the per-seq path (bit-gateable);
1583    /// MEMRA_GDN_VL=0 or a non-mma/non-C32 config falls back to the per-seq loop.
1584    #[allow(clippy::too_many_arguments)]
1585    fn linear_attn_prime_core_batch(&self, e: &Engine, la: &LinearAttnLayer,
1586                                    g4: &[CudaSlice<f32>], offs: &[usize], ts: &[usize],
1587                                    caches: &mut [&mut Cache], il: usize)
1588                                    -> Result<Vec<(CudaSlice<f32>, Option<CudaSlice<u8>>)>, Box<dyn std::error::Error>> {
1589        let ssm = self.cfg.ssm.as_ref().unwrap();
1590        let d_state = ssm.state_size as usize;
1591        let num_k = ssm.group_count as usize;
1592        let num_v = ssm.time_step_rank as usize;
1593        let key_dim = d_state * num_k;
1594        let value_dim = d_state * num_v;
1595        let conv_dim = key_dim * 2 + value_dim;
1596        let eps = self.cfg.rms_eps;
1597        let scale = 1.0 / (d_state as f32).sqrt();
1598        let b = ts.len();
1599        let c = Engine::gdn_chunk_size();
1600        // CONTINUATION batches (increment (b)) take the per-seq stateful pad_view twin
1601        // below — the varlen K4/K5 chain is fresh-only (zero initial state assumed).
1602        let carried = caches.iter().any(|c| c.pos > 0);
1603        let use_vl = !carried
1604            && (2..=8).contains(&b)
1605            && Engine::gdn_chunked_enabled() && ts.iter().all(|&t| t >= 16)
1606            && e.gdn_mma_enabled(c)
1607            && std::env::var("MEMRA_GDN_VL").as_deref() != Ok("0");
1608        if !use_vl {
1609            return (0..b).map(|s| {
1610                let (o, t) = (offs[s], ts[s]);
1611                self.linear_attn_prime_core_pad_view(
1612                    e, la,
1613                    &g4[0].slice(o * conv_dim..(o + t) * conv_dim),
1614                    &g4[1].slice(o * value_dim..(o + t) * value_dim),
1615                    &g4[2].slice(o * num_v..(o + t) * num_v),
1616                    &g4[3].slice(o * num_v..(o + t) * num_v),
1617                    t, caches[s], il, None)
1618            }).collect();
1619        }
1620        // increment 3: the ENTIRE core is varlen — allocs only per seq, then 13 launches
1621        // total for the whole batch: [conv, ring, repack, l2(q+k), gate-prep] +
1622        // [k-mirror, K1, K2, K3, w-mirror, K4, K5] + [gated-norm tail].
1623        struct SeqBufs {
1624            conv_out: CudaSlice<f32>, q_g: CudaSlice<f32>, k_g: CudaSlice<f32>, v_g: CudaSlice<f32>,
1625            q_l2: CudaSlice<f32>, k_l2: CudaSlice<f32>, beta: CudaSlice<f32>, g_log: CudaSlice<f32>,
1626            gn: CudaSlice<f32>, gn16: CudaSlice<u8>,
1627        }
1628        let d_conv = ssm.conv_kernel as usize;
1629        let f16o = Self::f16out_on(e, 16);
1630        let hk = Self::gdn_hk(e, 16, num_v, num_k);   // vl path is always chunked+mma
1631        let mut sb = Vec::with_capacity(b);
1632        let mut pres = Vec::with_capacity(b);
1633        for &t in ts.iter().take(b) {
1634            sb.push(SeqBufs {
1635                conv_out: e.uninit(conv_dim * t)?,
1636                q_g: e.uninit(d_state * hk * t)?,
1637                k_g: e.uninit(d_state * hk * t)?,
1638                v_g: e.uninit(d_state * num_v * t)?,
1639                q_l2: e.uninit(d_state * hk * t)?,
1640                k_l2: e.uninit(d_state * hk * t)?,
1641                beta: e.uninit(t * num_v)?,
1642                g_log: e.uninit(t * num_v)?,
1643                gn: e.uninit(d_state * num_v * t)?,
1644                gn16: e.alloc_u8_uninit(d_state * num_v * t * 2)?,
1645            });
1646            pres.push(e.gdn_chunk_alloc(num_v, t, c, hk)?);
1647        }
1648        let prep_args: Vec<crate::GdnPrepVl> = (0..b).map(|s| {
1649            let (o, t) = (offs[s], ts[s]);
1650            let rl = caches[s].recur[il].as_ref().unwrap();
1651            crate::GdnPrepVl {
1652                qkv: e.addr_f32v(&g4[0].slice(o * conv_dim..(o + t) * conv_dim)),
1653                conv_state: e.addr_f32(&rl.conv_state),
1654                conv_out: e.addr_f32(&sb[s].conv_out),
1655                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),
1656                q_l2: e.addr_f32(&sb[s].q_l2), k_l2: e.addr_f32(&sb[s].k_l2),
1657                beta_raw: e.addr_f32v(&g4[2].slice(o * num_v..(o + t) * num_v)),
1658                alpha: e.addr_f32v(&g4[3].slice(o * num_v..(o + t) * num_v)),
1659                beta: e.addr_f32(&sb[s].beta), g_log: e.addr_f32(&sb[s].g_log),
1660                o: e.addr_f32(&pres[s].o),
1661                z: e.addr_f32v(&g4[1].slice(o * value_dim..(o + t) * value_dim)),
1662                gn: e.addr_f32(&sb[s].gn), gn16: e.addr_u8(&sb[s].gn16),
1663                kb16: if Engine::l2_v2_on(d_state) { e.addr_u8(&pres[s].kb16) } else { 0 },
1664                qb16: if Engine::l2_v2_on(d_state) && e.gdn_wgmma_on(c) { e.addr_u8(&pres[s].qb16) } else { 0 },
1665                t: t as i32, pad: 0,
1666            }
1667        }).collect();
1668        let args: Vec<crate::GdnSeqVl> = (0..b).map(|s| {
1669            let rl = caches[s].recur[il].as_ref().unwrap();
1670            crate::GdnSeqVl {
1671                kb16: e.addr_u8(&pres[s].kb16), gcum: e.addr_f32(&pres[s].gcum),
1672                beta: e.addr_f32(&sb[s].beta), u: e.addr_f32(&pres[s].u),
1673                wb16: e.addr_u8(&pres[s].wb16), y: e.addr_u8(&pres[s].y16),
1674                ssnap: e.addr_u8(&pres[s].ssnap16),
1675                state_in: e.addr_f32(&rl.ssm_state), state_out: e.addr_f32(&rl.ssm_state_alt),
1676                q: e.addr_f32(&sb[s].q_l2), p: e.addr_f32(&pres[s].p),
1677                o: e.addr_f32(&pres[s].o),
1678                k: e.addr_f32(&sb[s].k_l2), v: e.addr_f32(&sb[s].v_g),
1679                g: e.addr_f32(&sb[s].g_log), a: e.addr_f32(&pres[s].a),
1680                w: e.addr_f32(&pres[s].w),
1681                t: ts[s] as i32, nc: pres[s].nc as i32,
1682            }
1683        }).collect();
1684        e.gdn_prep_vl8(&prep_args, la.ssm_conv1d.float_data(), la.ssm_dt.float_data(),
1685                       la.ssm_a.float_data(), conv_dim, d_conv, d_state, num_v, num_k, key_dim, hk, eps)?;
1686        // mirror-fold (round 27): l2 v2 emits kb16 in-epilogue; K3's store emits wb16 —
1687        // both standalone mirror launches vanish on the default config.
1688        if !Engine::l2_v2_on(d_state) {
1689            e.gdn_mirror_vl8(&args, num_v, 0, hk)?;
1690        }
1691        // task #22: wgmma-fused vl twins — qb16 mirrors + GdnWVl8 side-struct.
1692        let wq8: Option<crate::GdnWVl8> = if e.gdn_wgmma_on(c) {
1693            // qb16 emitted by the vl l2 mirror-fold when l2-v2 serves; bulk cvt otherwise
1694            if !Engine::l2_v2_on(d_state) {
1695                for s in 0..b {
1696                    e.f32_to_bf16_into(&sb[s].q_l2, &mut pres[s].qb16, d_state * hk * ts[s])?;
1697                }
1698            }
1699            let mut wa = [crate::GdnWVl::default(); 8];
1700            for s in 0..b {
1701                wa[s] = crate::GdnWVl { qb16: e.addr_u8(&pres[s].qb16), pb16: e.addr_u8(&pres[s].pb16) };
1702            }
1703            Some(crate::GdnWVl8(wa))
1704        } else { None };
1705        e.gdn_chunk_k123_vl8(&args, num_v, hk, wq8.as_ref())?;
1706        e.gdn_chunk_vl8(&args, num_v, scale, hk, wq8.as_ref())?;
1707        if f16o {
1708            e.gdn_tail_vl8(&prep_args, la.ssm_norm.float_data(), d_state, num_v, eps)?;
1709        }
1710        // per-seq state swap (+ non-f16out tail fallback)
1711        let mut out = Vec::with_capacity(b);
1712        for (s, bufs) in sb.into_iter().enumerate() {
1713            let rl = caches[s].recur[il].as_mut().unwrap();
1714            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1715            let (o, t) = (offs[s], ts[s]);
1716            let SeqBufs { mut gn, gn16, .. } = bufs;
1717            if f16o {
1718                out.push((gn, Some(gn16)));
1719            } else {
1720                let z_v = g4[1].slice(o * value_dim..(o + t) * value_dim);
1721                e.gated_rmsnorm_zv(&pres[s].o, la.ssm_norm.float_data(), &z_v, &mut gn,
1722                                   d_state, num_v * t, eps)?;
1723                out.push((gn, None));
1724            }
1725        }
1726        Ok(out)
1727    }
1728
1729    /// task #16: view-consuming GDN prime core — the batched prime hands row-offset
1730    /// views of the CONCAT projection outputs directly (no per-seq split copies).
1731    /// Same kernels, same values, byte-identical to the Vec shim above.
1732    #[allow(clippy::too_many_arguments)]
1733    fn linear_attn_prime_core_pad_view(&self, e: &Engine, la: &LinearAttnLayer,
1734                              qkv_mixed: &cudarc::driver::CudaView<f32>,
1735                              z: &cudarc::driver::CudaView<f32>,
1736                              beta_raw: &cudarc::driver::CudaView<f32>,
1737                              alpha: &cudarc::driver::CudaView<f32>,
1738                              t: usize, cache: &mut Cache, il: usize,
1739                              pad_len: Option<&CudaSlice<i32>>)
1740                              -> Result<(CudaSlice<f32>, Option<CudaSlice<u8>>), Box<dyn std::error::Error>> {
1741        let cfg = &self.cfg;
1742        let ssm = cfg.ssm.as_ref().unwrap();
1743        let d_state = ssm.state_size as usize;       // 128
1744        let num_v = ssm.time_step_rank as usize;     // 32
1745        let eps = cfg.rms_eps;
1746        let scale = 1.0 / (d_state as f32).sqrt();
1747
1748        let prep = self.linear_attn_gdn_prep(e, la, qkv_mixed, beta_raw, alpha, t, cache, il, pad_len)?;
1749
1750        // ONE gdn_scan over T from the cache's CURRENT state (zero at fresh prime); the final
1751        // state lands in the spare buffer and ping-pongs back (stable resident pointers, the
1752        // decode-determinism discipline from linear_attn_decode_inner). A4: `gdn_scan_prefill`
1753        // dispatches the chunked WY form under MEMRA_GDN_CHUNKED (prefill-only seam; decode +
1754        // verify keep the sequential kernel).
1755        let mut o = e.uninit(d_state * num_v * t)?;
1756        let rl = cache.recur[il].as_mut().unwrap();
1757        {
1758            let crate::cache::RecurLayer { ssm_state, ssm_state_alt, .. } = rl;
1759            e.gdn_scan_prefill(&prep.q_l2, &prep.k_l2, &prep.v_g, &prep.g_log, &prep.beta,
1760                               prep.kb16.as_ref(), prep.qb16.as_ref(), ssm_state, ssm_state_alt, &mut o, num_v, t, scale,
1761                               prep.hk)?;
1762        }
1763        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1764
1765        // gated RMSNorm + out projection (prefill dispatch). task #17: the f16out twin also
1766        // emits the ssm_out GEMM's fp16 operand in-epilogue (kills the standalone convert).
1767        let mut gn = e.uninit(d_state * num_v * t)?;
1768        let gn16 = if Self::f16out_on(e, t) {
1769            let mut g16 = e.alloc_u8_uninit(d_state * num_v * t * 2)?;
1770            e.gated_rmsnorm_f16out_zv(&o, la.ssm_norm.float_data(), z, &mut gn, &mut g16,
1771                                      d_state, num_v * t, eps)?;
1772            Some(g16)
1773        } else {
1774            e.gated_rmsnorm_zv(&o, la.ssm_norm.float_data(), z, &mut gn, d_state, num_v * t, eps)?;
1775            None
1776        };
1777        Ok((gn, gn16))
1778    }
1779
1780    /// Task #15 core-split wrapper: composes inner + ssm_out (== the old core).
1781    #[allow(clippy::too_many_arguments)]
1782    fn linear_attn_prime_core_pad(&self, e: &Engine, la: &LinearAttnLayer, g4: Vec<CudaSlice<f32>>,
1783                              t: usize, cache: &mut Cache, il: usize,
1784                              pad_len: Option<&CudaSlice<i32>>)
1785                              -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1786        let (gn, gn16) = self.linear_attn_prime_core_pad_inner(e, la, g4, t, cache, il, pad_len)?;
1787        if let Some(xh) = &gn16 {
1788            if let Some(y) = e.try_f16_gemm_pre(&la.ssm_out, xh, t)? {
1789                return Ok(y);
1790            }
1791        }
1792        Ok(e.matmul(&la.ssm_out, &gn, t)?)
1793    }
1794
1795    /// Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).
1796    pub fn full_attn(&self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
1797                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1798        let cfg = &self.cfg;
1799        let _n_embd = cfg.n_embd as usize;
1800        let n_head = cfg.n_head as usize;
1801        let n_head_kv = cfg.n_head_kv as usize;
1802        let head_dim = cfg.head_dim_k as usize;
1803        let eps = cfg.rms_eps;
1804        let scale = 1.0 / (head_dim as f32).sqrt();
1805
1806        // qwen35: wq output = head_dim*2*n_head (fused [q|gate] per head). M3/Hy3: NO output
1807        // gate — wq out = n_head*head_dim, no split (see prime-path note).
1808        let gated = cfg.attn_out_gate();
1809        // grouped: one f16 activation convert feeds q/k/v (matmul_group)
1810        let mut g3 = e.matmul_group(&[&fa.wq, &fa.wk, &fa.wv], h, t)?;
1811        let v = g3.pop().unwrap();
1812        let mut k = g3.pop().unwrap();
1813        let qf = g3.pop().unwrap();
1814        let (mut q, gate) = if gated {
1815            let mut q = e.uninit(t * n_head * head_dim)?;
1816            let mut gate = e.uninit(t * n_head * head_dim)?;
1817            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
1818            (q, Some(gate))
1819        } else {
1820            (qf, None)
1821        };
1822
1823        // QK-norm (per head_dim row), then partial RoPE.
1824        let mut qn = e.uninit(t * n_head * head_dim)?;
1825        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * t, eps)?;
1826        q = qn;
1827        let mut kn = e.uninit(t * n_head_kv * head_dim)?;
1828        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * t, eps)?;
1829        k = kn;
1830        let rope_dims = cfg.rope_dim_count as usize;
1831        e.rope_neox(&mut q, pos_d, head_dim, rope_dims, n_head, t, cfg.rope_freq_base, 1.0)?;
1832        e.rope_neox(&mut k, pos_d, head_dim, rope_dims, n_head_kv, t, cfg.rope_freq_base, 1.0)?;
1833
1834        // SDPA
1835        let mut attn = e.uninit(t * n_head * head_dim)?;
1836        // hand-written FlashAttention prefill (head_dim 256/128 stamped twins). MEMRA_NOFA
1837        // falls back to naive sdpa.
1838        if std::env::var("MEMRA_NOFA").is_ok() || !(head_dim == 256 || head_dim == 128) {
1839            // head_dim gate: see prime-path note (fa_prefill is stamped at 256 and 128 only).
1840            e.sdpa_naive(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1841        } else {
1842            e.fa_prefill(&q, &k, &v, &mut attn, head_dim, n_head, n_head_kv, t, t, scale, true)?;
1843        }
1844
1845        // output gate: attn * sigmoid(gate) — qwen35 only (M3 has no gate).
1846        let attn_g = match &gate {
1847            Some(gate) => {
1848                let mut gsig = e.uninit(t * n_head * head_dim)?;
1849                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
1850                let mut ag = e.uninit(t * n_head * head_dim)?;
1851                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
1852                ag
1853            }
1854            None => attn,
1855        };
1856
1857        // o projection
1858        let o = e.matmul(&fa.wo, &attn_g, t)?;
1859        Ok(o)
1860    }
1861
1862    /// Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).
1863    pub fn linear_attn(&self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize)
1864                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1865        let cfg = &self.cfg;
1866        let _n_embd = cfg.n_embd as usize;
1867        let ssm = cfg.ssm.as_ref().unwrap();
1868        let d_state = ssm.state_size as usize;       // 128
1869        let num_k = ssm.group_count as usize;        // 16
1870        let num_v = ssm.time_step_rank as usize;     // 32
1871        let d_conv = ssm.conv_kernel as usize;       // 4
1872        let head_k = d_state; let head_v = d_state;
1873        let key_dim = head_k * num_k;                // 2048
1874        let value_dim = head_v * num_v;              // 4096
1875        let conv_dim = key_dim * 2 + value_dim;      // 8192
1876        let eps = cfg.rms_eps;
1877        let scale = 1.0 / (d_state as f32).sqrt();
1878
1879        // projections
1880        // grouped: one f16 activation convert feeds all four projections (matmul_group)
1881        let mut g4 = e.matmul_group(&[&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha], h, t)?;
1882        let alpha = g4.pop().unwrap();                   // [T, num_v]
1883        let beta_raw = g4.pop().unwrap();                // [T, num_v]
1884        let z = g4.pop().unwrap();                       // [T, value_dim]
1885        let qkv_mixed = g4.pop().unwrap();               // [T, conv_dim] token-major
1886
1887        // conv + GDN repack, FUSED (2026-07-03): ssm_conv1d_gdn reads qkv_mixed [T, conv_dim]
1888        // token-major DIRECTLY (causal window rows t-pad..t, rows<0 = zero prefill state), applies
1889        // the 8-tap conv + SiLU, and scatters straight into the GDN [d_state, num_v, T] q/k/v
1890        // layout with the modulo head-repeat. Replaces transpose + zeros + conv_left_pad +
1891        // ssm_conv1d + qkv_to_gdn_repack (5 launches, conv_in/conv_out scratch + a 16MB@T=512
1892        // round-trip). BIT-IDENTICAL accumulation and scatter mapping.
1893        let _ = (head_k, head_v);
1894        let mut q_g = e.uninit(d_state * num_v * t)?;
1895        let mut k_g = e.uninit(d_state * num_v * t)?;
1896        let mut v_g = e.uninit(d_state * num_v * t)?;
1897        e.ssm_conv1d_gdn(&qkv_mixed, la.ssm_conv1d.float_data(), &mut q_g, &mut k_g, &mut v_g,
1898                         conv_dim, t, d_conv, d_state, num_v, num_k, key_dim)?;
1899        // L2-norm q,k per (head_dim) row — rows are contiguous d_state in q_g.
1900        let mut q_l2 = e.uninit(d_state * num_v * t)?;
1901        e.l2_norm(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
1902        let mut k_l2 = e.uninit(d_state * num_v * t)?;
1903        e.l2_norm(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
1904        let v_gd = v_g;
1905
1906        // beta = sigmoid(beta_raw) ; g_log = a * softplus(alpha + dt). Both need [num_v, T] layout
1907        // (g[t*num_v + h]). beta_raw/alpha are [T, num_v] token-major == that layout already.
1908        let mut beta = e.uninit(t * num_v)?;
1909        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
1910        // gdn_glog expects alpha [H,T] with alpha[t*H+h] and dt_bias/a [H] — matches token-major [T,num_v].
1911        let mut g_log = e.uninit(t * num_v)?;
1912        e.gdn_glog(&alpha, la.ssm_dt.float_data(), la.ssm_a.float_data(), &mut g_log, num_v, t)?;
1913
1914        // GDN scan (A4: gdn_scan_prefill dispatches chunked WY under MEMRA_GDN_CHUNKED)
1915        let state_in = e.zeros(d_state * d_state * num_v)?;  // zero state (prefill)
1916        let mut state_out = e.zeros(d_state * d_state * num_v)?;
1917        let mut o = e.uninit(d_state * num_v * t)?;
1918        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)?;
1919
1920        // gated RMSNorm: dst = RMSNorm(o, ssm_norm[head_v]) * silu(z). o is [d_state, num_v, T];
1921        // rows of head_v=d_state, nrows = num_v*T. z must match row layout: z is [T, value_dim] token-major
1922        // = [T, num_v*head_v]; per (t, vh) the head_v slice is contiguous -> rows align as (t*num_v+vh).
1923        // o rows are (t*num_v+vh) too. Good.
1924        let mut gn = e.uninit(d_state * num_v * t)?;
1925        e.gated_rmsnorm(&o, la.ssm_norm.float_data(), &z, &mut gn, d_state, num_v * t, eps)?;
1926
1927        // ssm_out projection: gn is [d_state, num_v, T] = [value_dim, T] viewed token-major as [T, value_dim]?
1928        // gn layout: (t*num_v+vh)*d_state + i  == token t, then (vh,i) = channel vh*d_state+i. That's
1929        // token-major [T, value_dim]. linear wants [T, in=value_dim]. Good.
1930        let out = e.matmul(&la.ssm_out, &gn, t)?;
1931        Ok(out)
1932    }
1933}
1934
1935impl HybridModel {
1936    /// MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd].
1937    /// Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.
1938    ///
1939    /// `il` is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a
1940    /// different 860160-byte block than the same expert of layer 7).
1941    ///
1942    /// Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER).
1943    /// Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache
1944    /// (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3).
1945    /// Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from `self`.
1946    pub fn moe_ffn_il(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16)
1947               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1948        Self::moe_ffn(e, m, z, t, &self.cfg, il, self.max_moe_block())
1949    }
1950
1951    /// Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the
1952    /// t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes
1953    /// (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.
1954    pub fn moe_ffn_il_zq8(&self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
1955                          zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16)
1956               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1957        Self::moe_ffn_inner(e, m, z, zq8, t, &self.cfg, il, self.max_moe_block())
1958    }
1959
1960    /// MoE FFN (EDGE-1), source-/model-agnostic. z: [T, n_embd] (already post-attention-normed).
1961    /// Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn. Shared by the hybrid
1962    /// (qwen35moe, shared expert present) and the dense-attention MoE (OLMoE, no shared expert) paths;
1963    /// `cfg.moe` supplies the dims and the optional shexp fields decide whether step 3 runs.
1964    ///
1965    /// `il` is the layer index — the residency-cache key prefix. `max_block` is the global max expert
1966    /// stride (fixed cache-slot size); pass `self.max_moe_block()`.
1967    pub(crate) fn moe_ffn(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
1968                          cfg: &ModelConfig, il: u16, max_block: usize)
1969               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1970        Self::moe_ffn_inner(e, m, z, None, t, cfg, il, max_block)
1971    }
1972
1973    #[allow(clippy::too_many_arguments)]
1974    pub(crate) fn moe_ffn_inner(
1975        e: &Engine,
1976        m: &MoeWeights,
1977        z: &CudaSlice<f32>,
1978        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
1979        t: usize,
1980        cfg: &ModelConfig,
1981        il: u16,
1982        max_block: usize,
1983    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1984        let worker_io = crate::spill_pread::worker_enabled();
1985        let epoch_lfu = std::env::var_os("MEMRA_MOE_LFU_DECAY").is_some();
1986        if Engine::moe_cache_enabled() && (worker_io || epoch_lfu) {
1987            e.with_moe_cache(max_block, |cache, _| {
1988                cache.begin_forward_epoch(il, t);
1989                if worker_io {
1990                    cache.begin_worker_scope();
1991                }
1992                Ok(())
1993            })?;
1994        }
1995        // A2: Expert-grouped dispatch for prefill (T>1). MEMRA_MOE_GROUPED=1 routes here.
1996        if t > 1 && std::env::var("MEMRA_MOE_GROUPED").is_ok() {
1997            let grouped_out = Self::moe_ffn_grouped(e, m, z, t, cfg, il, max_block)?;
1998            // MEMRA_MOE_GATE: byte-identity comparison vs sequential path.
1999            // KNOWN t>1 MISMATCH maxdiff ~3.4e-4 (deterministic, 5x bit-identical 2026-07-05): the
2000            // sequential arm routes resident experts through the dev_q8 dp4a path (q8_1-quantized z
2001            // and act rows) while grouped stays f32-dequant qmatvec — a quantize-path difference,
2002            // not a bug (per-stage: act q8-vs-f32 ~4-9e-3 abs on |act|<=3, down-only ~1-3e-4; the
2003            // q8_1 activation-quantize error class). MEMRA_MOE_Q8=0 restores BYTE-IDENTICAL.
2004            if std::env::var("MEMRA_MOE_GATE").is_ok() {
2005                let seq_out = Self::moe_ffn_sequential(e, m, z, t, cfg, il, max_block)?;
2006                let g_host = e.dtoh(&grouped_out)?;
2007                let s_host = e.dtoh(&seq_out)?;
2008                let g_bytes: &[u8] = unsafe { std::slice::from_raw_parts(g_host.as_ptr() as *const u8, g_host.len() * 4) };
2009                let s_bytes: &[u8] = unsafe { std::slice::from_raw_parts(s_host.as_ptr() as *const u8, s_host.len() * 4) };
2010                if g_bytes == s_bytes {
2011                    if il == 0 { println!("moe-gate il={il} t={t} BYTE-IDENTICAL (first layer only printed)"); }
2012                } else {
2013                    let diffs = g_host.iter().zip(s_host.iter()).enumerate()
2014                        .filter(|(_, (a, b))| a != b).count();
2015                    let maxdiff = g_host.iter().zip(s_host.iter())
2016                        .map(|(a, b)| (a - b).abs()).fold(0.0f32, f32::max);
2017                    panic!("moe-gate il={il} t={t} MISMATCH: {diffs}/{} elems differ, maxdiff={maxdiff:.6e}", g_host.len());
2018                }
2019            }
2020            return Ok(grouped_out);
2021        }
2022        Self::moe_ffn_sequential_zq8(e, m, z, zq8, t, cfg, il, max_block)
2023    }
2024
2025    /// Sequential (per-token) MoE FFN -- the original path. Factored out for the gate comparison.
2026    pub(crate) fn moe_ffn_sequential(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
2027                          cfg: &ModelConfig, il: u16, max_block: usize)
2028               -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2029        Self::moe_ffn_sequential_zq8(e, m, z, None, t, cfg, il, max_block)
2030    }
2031
2032    /// Append the host-visible router selection for one layer/forward when calibration tracing is
2033    /// enabled. Both sequential and expert-grouped prefill must call this after routing so the
2034    /// trace is independent of the dispatch optimization selected for the forward.
2035    fn trace_moe_routes(il: u16, t: usize, sel_all: &[u32], weights: &[f32])
2036                        -> Result<(), Box<dyn std::error::Error>> {
2037        use std::io::Write as _;
2038        if let Ok(path) = std::env::var("MEMRA_MOE_TRACE") {
2039            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2040            let ids: Vec<String> = sel_all.iter().map(|s| s.to_string()).collect();
2041            writeln!(f, "{} {} {}", il, t, ids.join(","))?;
2042        }
2043        if let Ok(path) = std::env::var("MEMRA_MOE_WEIGHT_TRACE") {
2044            let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?;
2045            let pairs: Vec<String> = sel_all.iter().zip(weights)
2046                .map(|(expert, weight)| format!("{expert}:{weight:.9}"))
2047                .collect();
2048            writeln!(f, "{} {} {}", il, t, pairs.join(","))?;
2049        }
2050        Ok(())
2051    }
2052
2053    /// Append the f32 input to one MoE layer for offline, layerwise calibration. This diagnostic
2054    /// intentionally performs a DtoH copy and is therefore disabled unless an explicit fresh trace
2055    /// directory is supplied. Each layer owns one payload file; index.jsonl records byte offsets so
2056    /// a validator can prove request/layer coverage before the trace is used for pruning or healing.
2057    fn trace_moe_input(e: &Engine, il: u16, t: usize, n_embd: usize, z: &CudaSlice<f32>)
2058                       -> Result<(), Box<dyn std::error::Error>> {
2059        use std::io::Write as _;
2060        let Ok(dir) = std::env::var("MEMRA_MOE_INPUT_TRACE_DIR") else { return Ok(()) };
2061        let host = e.dtoh(z)?;
2062        if host.len() != t * n_embd {
2063            return Err(format!(
2064                "MoE input trace shape mismatch at layer {il}: got {} values, expected {}x{}",
2065                host.len(), t, n_embd
2066            ).into());
2067        }
2068        let bytes = unsafe {
2069            std::slice::from_raw_parts(
2070                host.as_ptr().cast::<u8>(), host.len() * std::mem::size_of::<f32>()
2071            )
2072        };
2073        let state = MOE_INPUT_TRACE_WRITER.get_or_init(|| std::sync::Mutex::new(None));
2074        let mut state = state.lock().map_err(|_| "MoE input trace writer lock is poisoned")?;
2075        if state.is_none() {
2076            let dir = std::path::PathBuf::from(&dir);
2077            std::fs::create_dir_all(&dir)?;
2078            let index = std::fs::OpenOptions::new().create(true).append(true)
2079                .open(dir.join("index.jsonl"))?;
2080            *state = Some(MoeInputTraceWriter {
2081                dir,
2082                index,
2083                payloads: std::collections::HashMap::new(),
2084            });
2085        }
2086        let writer = state.as_mut().unwrap();
2087        if writer.dir != std::path::Path::new(&dir) {
2088            return Err("MEMRA_MOE_INPUT_TRACE_DIR changed after capture started".into());
2089        }
2090        let file_name = format!("layer-{il:03}.f32");
2091        if !writer.payloads.contains_key(&il) {
2092            let payload = std::fs::OpenOptions::new().create(true).append(true)
2093                .open(writer.dir.join(&file_name))?;
2094            let offset = payload.metadata()?.len();
2095            writer.payloads.insert(il, (payload, offset));
2096        }
2097        let (payload, offset) = writer.payloads.get_mut(&il).unwrap();
2098        let row_offset = *offset;
2099        payload.write_all(bytes)?;
2100        *offset += bytes.len() as u64;
2101        writeln!(
2102            writer.index,
2103            "{{\"format\":\"memra-moe-input-trace-v1\",\"layer\":{il},\"tokens\":{t},\
2104             \"hidden_size\":{n_embd},\"file\":\"{file_name}\",\"offset\":{row_offset},\
2105             \"payload_bytes\":{}}}",
2106            bytes.len()
2107        )?;
2108        Ok(())
2109    }
2110
2111    #[allow(clippy::too_many_arguments)]
2112    pub(crate) fn moe_ffn_sequential_zq8(
2113        e: &Engine,
2114        m: &MoeWeights,
2115        z: &CudaSlice<f32>,
2116        zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>,
2117        t: usize,
2118        cfg: &ModelConfig,
2119        il: u16,
2120        max_block: usize,
2121    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2122        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
2123        let moe = cfg.moe.as_ref().unwrap();
2124        let n_embd = cfg.n_embd as usize;          // 2048 (gate/up in_f, down out_f)
2125        let n_expert = moe.expert_count as usize;  // 256
2126        let n_used = moe.expert_used_count as usize; // 8
2127        let n_ff_exp = moe.expert_ff_length as usize; // 512 (gate/up out_f, down in_f)
2128
2129        // verify the HostExps dims match cfg (catches a wrong-file / transpose mixup)
2130        debug_assert_eq!(m.gate_exps.in_f, n_embd);
2131        debug_assert_eq!(m.gate_exps.out_f, n_ff_exp);
2132        debug_assert_eq!(m.down_exps.in_f, n_ff_exp);  // down is TRANSPOSED: in=512
2133        debug_assert_eq!(m.down_exps.out_f, n_embd);   //                     out=2048
2134        debug_assert_eq!(m.gate_exps.n_expert, n_expert);
2135
2136        let use_cache = Engine::moe_cache_enabled();
2137        let uniform_experts = m.has_uniform_expert_layout();
2138        let moe_q8 = uniform_experts && moe_q8_enabled()
2139            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2140            && q8_expert_supported(m.down_exps.qtype);
2141        // Experimental secondary backend: complete experts already resident in the SLRU stay on
2142        // CUDA; any expert missing one or more projections runs from the original host GGUF bytes
2143        // through llama.cpp CPU quant dots. Small-t only keeps decode and speculative verification
2144        // in the same numeric/dispatch class; real prefill remains on the established GPU path and
2145        // seeds the residency cache. An explicit library path is the build/runtime gate, so naked
2146        // commands and CI have no llama.cpp or OpenMP dependency.
2147        let cpu_expert_requested = crate::cpu_experts::configured();
2148        if cpu_expert_requested && (cfg.hy3.is_none() || cfg.m3.is_some()) {
2149            return Err(std::io::Error::other(
2150                "MEMRA_CPU_EXPERT_LIB is experimental and currently gated to Hy3",
2151            )
2152            .into());
2153        }
2154        let cpu_hybrid = cpu_expert_requested && t < PRIME_MIN_T && m.dev_exps.is_none();
2155        // A dynamic SLRU changes which complete experts run on CUDA versus llama.cpp CPU dots.
2156        // Those backends are each deterministic but are different numeric configurations, so a
2157        // later prefill eviction can change greedy output. Freeze after the first real prefill;
2158        // decode/spec reads the fixed resident set, while later prefill misses use transient GPU
2159        // staging below and cannot change backend assignment.
2160        let freeze_cpu_residency = cpu_expert_requested
2161            && std::env::var("MEMRA_CPU_EXPERT_FREEZE_CACHE").as_deref() == Ok("1");
2162        let caller_warms_before_freeze = std::env::var("MEMRA_CPU_EXPERT_FREEZE_WARMUP_TOKENS")
2163            .ok()
2164            .and_then(|value| value.parse::<usize>().ok())
2165            .is_some_and(|tokens| tokens > 0);
2166        if cpu_hybrid && freeze_cpu_residency && !caller_warms_before_freeze {
2167            e.freeze_moe_cache();
2168        }
2169        let cache_frozen = use_cache && e.moe_cache_frozen();
2170        let cache_dispatch = use_cache && (!cache_frozen || cpu_hybrid);
2171
2172        // 1. ROUTER: logits = ffn_gate_inp @ z  -> [T, 256]. gate_inp is F32 -> cuBLASLt, whose
2173        // reductions are n-DEPENDENT (lt_ndep probe: m=1 vs m=2 col0 differs every bit). At
2174        // small t (spec verify, 2..15) that shifts router logits vs the T=1 decode chain ->
2175        // top-k WEIGHTS (and at tie margins the SELECTION) differ -> verify != decode. Route
2176        // small-t through per-column m=1 calls (decode-exact contract); real prefill keeps the
2177        // batched GEMM.
2178        let logits = if t < PRIME_MIN_T {
2179            // t == 1 included since 2026-07-10 (was cuBLAS gemvx, 3.1% + adjacent of the depth
2180            // decode map): decode and verify now route through the SAME kernel — the
2181            // verify==decode router parity holds by construction instead of by FP-order luck.
2182            if crate::router_kernel_on() {
2183                // MEMRA_ROUTER_KERNEL=1: in-house router GEMV (battery-gated numeric config —
2184                // top-k discontinuity means FP-order changes can flip routing; oracle arbitrates).
2185                e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2186                              m.gate_exps.n_expert, t)?
2187            } else {
2188                e.matmul_decode_exact(&m.gate_inp, z, t)?
2189            }
2190        } else if crate::router_prefill_exact_on() && crate::router_kernel_on() {
2191            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): the cuBLASLt router GEMM
2192            // is m-DEPENDENT — probed on the Ornith-35B router weight, rows [0,19) of an m=65
2193            // call differ from the m=19 call by 3.9e-3 while the same probe on the lm_head /
2194            // wq MMQ+f16 weights is BIT-IDENTICAL across m (research/concat-prime-exact-20260802,
2195            // gemm-razor-router-o35b.log vs gemm-razor-o35b.log). Because the router feeds a
2196            // top-k DISCONTINUITY, that perturbation reorders ties and at ~16% of (layer,token)
2197            // pairs changes the selected expert SET — so a request's own prefill routing depended
2198            // on how many OTHER requests' tokens shared its concat batch (cross-request prime
2199            // batching, worker.rs task #13). The in-house router GEMV computes one row per
2200            // (expert, token) block with a fixed per-row reduction order and is m-INVARIANT
2201            // (same probe: BIT-IDENTICAL, gemm-razor-router-gemv-o35b.log), so routing prefill
2202            // through it makes a session's routing a function of its OWN tokens alone — the
2203            // serving isolation contract at the prime level. MEMRA_ROUTER_PREFILL_EXACT=0 reverts
2204            // to the batched cuBLASLt GEMM (numeric-config rollback seam).
2205            e.router_gemv(m.gate_inp.float_data(), z, cfg.n_embd as usize,
2206                          m.gate_exps.n_expert, t)?
2207        } else {
2208            e.matmul(&m.gate_inp, z, t)?
2209        };
2210
2211        // LAUNCH-STRUCTURE STAGE 3 (2026-07-05, MEMRA_MOE_DEV default ON, =0 rollback): ZERO-DtoH
2212        // device-dispatch when this layer's expert blocks are ALL cache-resident. The fused
2213        // router's sel/w stay ON DEVICE; the expert weight pointers come from the per-layer
2214        // device table of fixed slot addresses; gate/up/silu + down/fma run as the same TWO
2215        // launches per token as gdec. Removes the per-layer router DtoH + stream sync — the
2216        // per-token host stall that dominated the 35B decode wall after stages 1+2.
2217        // BIT-IDENTITY: the router kernel is selection-exact vs the host oracle (kernel-check
2218        // tie gate) and the _dev matvec twins reproduce the gdec kernels' exact FP chains; the
2219        // only difference is where sel/w/pointers are READ from (device instead of params).
2220        // Residency: one-shot PREWARM force-admits the layer while free slots cover it
2221        // (MEMRA_MOE_PREWARM=0 -> organic residency, dev path fires when the SLRU fills).
2222        // Any non-resident layer falls through to host routing + the gdec/sequential path.
2223        // FITS-VRAM RESIDENT EXPERTS (2026-07-06): the layer's expert slabs are device-resident
2224        // (load-time decision) — fire the zero-DtoH dev path unconditionally with the prebuilt
2225        // pointer row. No cache, no dispatch, no residency check: the llama full-offload regime
2226        // (it measured 169.55 vs the cache path's 28.5 on the local 35B — the residency-gate
2227        // all-or-nothing fallback was the 6x). BIT-IDENTITY: same _dev kernels, same math; only
2228        // the pointer table's provenance differs (slab base+stride vs SLRU slot addresses).
2229        // MoE PREFILL PAIR-BATCH (2026-07-06, the 16x pp hole): t>1 on resident experts — ONE
2230        // launch per proj covers ALL (token,expert) pairs (grid.y=pair, warp-per-row), replacing
2231        // the per-expert loop (256 experts x 3-4 launches x tiny m_e). Scatter is slot-ordered
2232        // per token (the sequential-axpy bit-identity class). Requires q8-supported qtypes +
2233        // resident slabs. MEMRA_MOE_PAIRS=0 rollback.
2234        // t >= PRIME_MIN_T only (2026-07-06 exactness fix, verify-probe proof): spec VERIFY
2235        // batches (t = 2..K+2) previously rode these pairs kernels while T=1 decode rode the
2236        // dev_q8 loop — different FP chains, verify-T2 logit maxdiff 2.6e-1 vs eager -> greedy
2237        // flips at tight margins -> 35B real-prompt spec self-consistency FAIL (the 27B
2238        // "verify must be kernel-DISPATCH-identical to decode" lesson, MoE edition). Small-t
2239        // now rides the dev loop below (same kernels per token as decode); pairs serves real
2240        // prefill (t >= 16, where spec never verifies).
2241        // sigmoid-router archs (M3, Hy3) must NOT enter the pairs/dev arms: those route via the
2242        // fused SOFTMAX device router (moe_router_topk) — silently wrong experts (the M3
2243        // gate-MISMATCH 74602-vs-92 lesson, 2026-07-07). Host sigmoid routing below is correct.
2244        // Per-expert macro-scales (compressed-tensors NVFP4 ST class, e.g. unsloth 35B-A3B):
2245        // the device-dispatch kernels (pairs/dev) do NOT fold them — those checkpoints must
2246        // ride the macro-aware sequential/staged paths below or every expert output is off by
2247        // its global scale (~3e4x, measured garbage 2026-07-16). GGUF experts: macros None.
2248        let no_exp_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2249            && m.down_exps.macros.is_none();
2250        if cfg.sigmoid_router().is_none() && cfg.m3.is_none() && cfg.hy3.is_none()
2251            && no_exp_macros
2252            && t >= PRIME_MIN_T && m.dev_exps.is_some() && moe_q8_enabled()
2253            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
2254            && q8_expert_supported(m.down_exps.qtype)
2255            && std::env::var("MEMRA_MOE_PAIRS").map(|v| v != "0").unwrap_or(true)
2256            && std::env::var("MEMRA_MOE_STATS").is_err() {
2257            return Self::moe_ffn_pairs(e, m, z, &logits, t, cfg);
2258        }
2259
2260        // t < PRIME_MIN_T: moe_ffn_dev loops tokens serially (1 launch-pair per token) — the
2261        // decode path AND the spec-verify path (dispatch parity = exactness; see pairs gate
2262        // above). Serial launches are fine at t<=10 (K+2); real prefill never lands here.
2263        // moe_ffn_dev routes via the FUSED SOFTMAX device router (moe_router_topk) — sigmoid
2264        // routing (M3, Hy3: +expert bias) has no device kernel yet, so those arches must NOT
2265        // enter the dev arms: with MOE_CACHE=1 M3 silently routed softmax = wrong experts
2266        // (gate MISMATCH 74602 vs 92, caught 2026-07-07). Host sigmoid path below is correct.
2267        // macro-carrying experts are handled inside the dev path now (epilogue fold + w-scale);
2268        // pairs/gdec/csr keep their macro gates until their kernels grow the fold.
2269        let dev_ok = uniform_experts && cfg.m3.is_none() && cfg.hy3.is_none();
2270        // Observation modes must route through the host-visible selection below. Otherwise a fully
2271        // resident layer returns through device dispatch before its trace/stats row is recorded,
2272        // silently biasing calibration toward only non-resident layers on large-VRAM machines.
2273        let observe_routes = std::env::var("MEMRA_MOE_STATS").is_ok()
2274            || std::env::var("MEMRA_MOE_TRACE").is_ok()
2275            || std::env::var("MEMRA_MOE_WEIGHT_TRACE").is_ok()
2276            || std::env::var("MEMRA_MOE_INPUT_TRACE_DIR").is_ok();
2277        if dev_ok && t < PRIME_MIN_T && m.dev_exps.is_some() && n_used <= 8 && moe_dev_enabled()
2278            && !observe_routes {
2279            return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2280        }
2281        if dev_ok && use_cache && n_used <= 8 && moe_dev_enabled()
2282            && !observe_routes {
2283            let row_ok = e.with_moe_cache(max_block, |c, eng| {
2284                if moe_prewarm_enabled() { c.prewarm_layer(il, m, eng)?; }
2285                Ok(c.layer_dev_row(il, n_expert, eng)?.is_some())
2286            })?;
2287            if row_ok {
2288                return Self::moe_ffn_dev(e, m, z, zq8, &logits, t, cfg, il, max_block);
2289            }
2290        }
2291
2292        // Per-token (sel[8], w[8]) — sigmoid host oracle (M3/Hy3), else fused-router/softmax.
2293        let (sel_all, w_all, routed_cpu_input) = if let Some(sig) = cfg.sigmoid_router() {
2294            if cpu_hybrid {
2295                let (sel, w, input) = Self::moe_route_sigmoid_with_input(
2296                    e,
2297                    &logits,
2298                    z,
2299                    t,
2300                    n_expert,
2301                    n_used,
2302                    m.exp_probs_b.as_deref(),
2303                    sig,
2304                    m.active_experts.as_deref(),
2305                )?;
2306                (sel, w, Some(input))
2307            } else {
2308                let (sel, w) = Self::moe_route_cfg(
2309                    e,
2310                    &logits,
2311                    t,
2312                    n_expert,
2313                    n_used,
2314                    m.exp_probs_b.as_deref(),
2315                    Some(sig),
2316                    m.active_experts.as_deref(),
2317                )?;
2318                (sel, w, None)
2319            }
2320        } else {
2321            let (sel, w) = Self::moe_route_cfg(
2322                e,
2323                &logits,
2324                t,
2325                n_expert,
2326                n_used,
2327                None,
2328                None,
2329                m.active_experts.as_deref(),
2330            )?;
2331            (sel, w, None)
2332        };
2333
2334        // MEMRA_MOE_TRACE=<path>: append one line per (layer, step) with the selected expert ids —
2335        // offline analysis derives the decode working set + step-to-step reuse (the go/no-go
2336        // measurement for resident-expert tiering; see rig5090.jsonl 2026-07-07 pinned-tier row).
2337        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
2338        Self::trace_moe_input(e, il, t, n_embd, z)?;
2339
2340        // Worker reads ordinarily overlap NVMe with the current expert but leave their H2D copies
2341        // demand-serialized on the compute stream. Hy3 decode already has a safe owner-thread
2342        // boundary here: sigmoid routing read `logits` back to the host, which synchronized all
2343        // earlier-layer compute. Submit the complete selected set first, then reserve only victims
2344        // outside that set and queue its H2D copies on the copy stream. Dispatch inserts an event
2345        // wait for each pending block, so later copies can overlap the earlier expert kernels while
2346        // the selected weights and numerical kernels remain unchanged. Restrict this experiment to
2347        // T=1; batched forwards can have token-local consumers still in flight between selections.
2348        // The CPU backend reads mmap-backed model bytes directly. Positioned-read worker buffers
2349        // are CUDA write-combined staging allocations and are intentionally not CPU-compute inputs;
2350        // do not spend NVMe bandwidth filling them for a layer whose misses go to CPU.
2351        let worker_disk_prefetch =
2352            cache_dispatch && crate::spill_pread::worker_enabled() && !cpu_hybrid;
2353        let promote_worker_h2d =
2354            t == 1 && worker_disk_prefetch && crate::spill_pread::copy_h2d_enabled();
2355        if promote_worker_h2d {
2356            let mut selected_blocks = Vec::with_capacity(n_used * 3);
2357            for &ex in sel_all.iter().take(n_used) {
2358                let ex = ex as u16;
2359                selected_blocks.extend([
2360                    BlockId::new(il, PROJ_GATE, ex),
2361                    BlockId::new(il, PROJ_UP, ex),
2362                    BlockId::new(il, PROJ_DOWN, ex),
2363                ]);
2364            }
2365            for &ex in sel_all.iter().take(n_used) {
2366                Self::moe_prefetch_disk_expert(e, il, ex as usize, m, max_block, &selected_blocks)?;
2367            }
2368            e.with_moe_cache(max_block, |cache, eng| {
2369                cache.promote_worker_reads_at_safe_boundary(
2370                    &selected_blocks,
2371                    &selected_blocks,
2372                    eng,
2373                )?;
2374                Ok(())
2375            })?;
2376        }
2377
2378        // MEMRA_MOE_STATS: per-layer routing stats for the A2 (expert-grouped prefill) baseline —
2379        // per-token expert-id entropy, active-expert coverage, tokens-per-expert group sizes.
2380        if t > 1 && std::env::var("MEMRA_MOE_STATS").is_ok() {
2381            let mut cnt = vec![0u32; n_expert];
2382            for &s in sel_all.iter() { cnt[s as usize] += 1; }
2383            let total = sel_all.len() as f64;
2384            let mut h = 0.0f64;
2385            let mut active = 0usize;
2386            for &c in &cnt { if c > 0 { active += 1; let p = c as f64 / total; h -= p * p.log2(); } }
2387            let maxc = cnt.iter().copied().max().unwrap_or(0);
2388            println!("moe-stats il={} t={} assignments={} active={}/{} entropy={:.3}b (max {:.3}b) mean_tok_per_active={:.2} max_tok_per_expert={}",
2389                     il, t, sel_all.len(), active, n_expert, h, (n_expert as f64).log2(), total / active.max(1) as f64, maxc);
2390        }
2391
2392        // LAUNCH-STRUCTURE STAGE 2 (2026-07-05): moe_out memset elision on the gdec path.
2393        // moe_down8_fma_f32 FULLY overwrites its token row (dst[o] = the in-kernel FMA chain that
2394        // starts at 0.0f — numerically the axpy-into-zeroed-row chain), so when the grouped-decode
2395        // path fires the upfront `e.zeros(t*n_embd)` memset is pure launch churn. Allocate uninit
2396        // when gdec CAN fire (any t — decode t=1 AND the spec verify t=K+1 route here per token)
2397        // and lazily zero ONLY the row of a token that falls through to the sequential axpy loop.
2398        // BIT-IDENTITY: unchanged — every row is either fully overwritten (gdec) or
2399        // zeroed-then-accumulated exactly as before (fallback).
2400        let gdec_may_fire = uniform_experts && use_cache && n_used <= 8 && gdec_enabled();
2401        let mut moe_out = if gdec_may_fire {
2402            e.uninit(t * n_embd)?
2403        } else {
2404            e.zeros(t * n_embd)?
2405        };
2406        // The router readback above already established a host boundary. Copy each small-t hidden
2407        // row once so CPU miss experts can start while the owner thread queues resident GPU work.
2408        let cpu_input = if cpu_hybrid {
2409            Some(routed_cpu_input.ok_or("CPU expert routing did not return the MoE input")?)
2410        } else {
2411            None
2412        };
2413
2414        // GPU scratch: one slot per proj, big enough for ONE expert (default stage-every-token path).
2415        // STAGE 2: LAZY — allocated only if the no-cache staging path actually runs (under
2416        // MEMRA_MOE_CACHE they were 3 dead ~1MB alloc_zeros + memset + free per layer per token,
2417        // measured ~123 memsets/token of the decode wall).
2418        let g_len = m.gate_exps.max_expert_bytes();  // 860160 for the uniform 35B gate
2419        let u_len = m.up_exps.max_expert_bytes();    // 860160 for the uniform 35B up
2420        let d_len = m.down_exps.max_expert_bytes();  // 1114112 for the uniform 35B down
2421        let mut scratch_g: Option<CudaSlice<u8>> = None;
2422        let mut scratch_u: Option<CudaSlice<u8>> = None;
2423        let mut scratch_d: Option<CudaSlice<u8>> = None;
2424        // `max_block` (the GLOBAL max expert stride across all layers) is passed in — the cache slots
2425        // are FIXED-ADDRESS and must fit any layer's block (UD/dynamic GGUFs vary quant per layer).
2426
2427        // EDGE-1 §C.2/C.3: the optional pipeline queues the next selected expert's MISS blocks on
2428        // the copy stream before launching the current expert's compute. Pending slots stay invisible
2429        // to cache hits until dispatch inserts the completion-event wait; current gate/up/down ids are
2430        // protected from eviction. This changes scheduling only, never the GGUF bytes or GEMM path.
2431        let page_window = moe_page_prefetch_window();
2432
2433        // 2. PER TOKEN: routed-expert loop. The ONE dispatch change vs Stage-1: a resident slot
2434        //    (cache HIT, no H2D) OR a staged slot (MISS) feeds the SAME unchanged qmatvec_view.
2435        for tok in 0..t {
2436            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
2437            let w = &w_all[tok * n_used..(tok + 1) * n_used];
2438            let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);  // CudaView<f32>
2439            let mut tok_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
2440
2441            // STAGE-2 GROUPED DECODE (2026-07-04, MEMRA_MOE_GDEC default ON, =0 rollback): fold
2442            // this token's whole routed-expert FFN (8x gate/up/silu + 8x down/axpy = 40 launches)
2443            // into TWO launches via expert-pointer indirection over the fixed-address cache slots.
2444            // Fires only when ALL 3*n_used blocks are ALREADY cache-resident (pure-HIT: zero
2445            // memcpy, zero admission, so no slot can move under the collected pointers) — any
2446            // miss falls through to the sequential loop below, which admits as before. In steady
2447            // state on a fully-resident rig every token-layer takes the grouped path.
2448            // BIT-IDENTITY: each in-kernel dot reproduces qmatvec_f32's exact reduction; SiLU is
2449            // silu_mul_f32's exact expression; the down accumulation is a slot-ordered
2450            // __fmaf_rn chain == the sequential axpy_f32 chain (MEMRA_MOE_GDEC_GATE compares).
2451            // cfg.m3: the grouped kernels' fused epilogues are plain SiLU — M3's swigluoai must
2452            // NOT take them until the kernels grow the clamped variant. NVFP4 experts carry
2453            // per-expert macro-scales the fused kernels don't fold — those fall through too.
2454            let no_macros = m.gate_exps.macros.is_none() && m.up_exps.macros.is_none()
2455                && m.down_exps.macros.is_none();
2456            if gdec_may_fire && moe_q8 && cfg.m3.is_none() && no_macros {
2457                if tok_q8.is_none() {
2458                    tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2459                }
2460                let (zq, zd) = tok_q8.as_ref().unwrap();
2461                if Self::moe_gdec_token_q8(e, m, il, max_block, zq, zd, sel, w,
2462                                           &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2463                    continue;
2464                }
2465            } else if gdec_may_fire && cfg.m3.is_none() && no_macros
2466                && Self::moe_gdec_token(e, m, il, max_block, &zt, sel, w,
2467                                        &mut moe_out, tok, n_embd, n_ff_exp, n_used)? {
2468                continue;
2469            }
2470
2471            // STAGE 2 memset-elision invariant: moe_out was allocated UNINIT when gdec could fire.
2472            // This token fell through to the sequential axpy loop, which ACCUMULATES — zero its row
2473            // first (row-sized memset, replaces the old full-buffer zeros; other rows are gdec-owned).
2474            if gdec_may_fire {
2475                let mut row = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2476                e.memset_zeros_view(&mut row)?;
2477            }
2478
2479            // Fiddler-style split at whole-expert granularity. Partial residency is deliberately a
2480            // CPU assignment: transferring the missing projection would reintroduce the PCIe/NVMe
2481            // stall this path exists to remove, while mixing projections would require another
2482            // activation round-trip. Weight addresses remain valid until this worker is joined at
2483            // the bottom of the token scope.
2484            let mut cpu_mask = vec![false; sel.len()];
2485            let cpu_worker = if let Some(host_input) = cpu_input.as_ref() {
2486                let gpu_resident = if use_cache {
2487                    e.with_moe_cache(max_block, |cache, _| {
2488                        Ok(sel
2489                            .iter()
2490                            .map(|&expert| {
2491                                let expert = expert as u16;
2492                                [PROJ_GATE, PROJ_UP, PROJ_DOWN]
2493                                    .into_iter()
2494                                    .filter(|&projection| {
2495                                        cache
2496                                            .resident(BlockId::new(il, projection, expert))
2497                                            .is_some()
2498                                    })
2499                                    .count()
2500                            })
2501                            .collect::<Vec<_>>())
2502                    })?
2503                } else {
2504                    vec![0; sel.len()]
2505                };
2506                let mut cpu_selected = Vec::new();
2507                for (index, (&expert, &route_weight)) in sel.iter().zip(w).enumerate() {
2508                    if gpu_resident[index] != 3 {
2509                        cpu_mask[index] = true;
2510                        crate::cpu_experts::record_incomplete_gpu_residency(gpu_resident[index]);
2511                        let expert = expert as usize;
2512                        cpu_selected.push((expert, route_weight));
2513                    }
2514                }
2515                if crate::cpu_experts::predictor_enabled() {
2516                    // Fire-and-forget lookahead: the predictor worker scores layers il+1..
2517                    // from this layer's MoE input and prefetches predicted-and-missing
2518                    // experts into the companion RAM cache. Never blocks this thread.
2519                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2520                    crate::cpu_experts::predictor_submit(il, row);
2521                }
2522                if cpu_selected.is_empty() {
2523                    None
2524                } else {
2525                    let row = &host_input[tok * n_embd..(tok + 1) * n_embd];
2526                    let job = crate::cpu_experts::prepare_job(m, il, &cpu_selected, row)
2527                        .map_err(std::io::Error::other)?;
2528                    Some(crate::cpu_experts::submit(job).map_err(std::io::Error::other)?)
2529                }
2530            } else {
2531                None
2532            };
2533
2534            let worker_window = worker_disk_prefetch
2535                .then(worker_prefetch_window)
2536                .unwrap_or(0);
2537            for (j, &ex) in sel.iter().enumerate() {
2538                if cpu_mask[j] {
2539                    continue;
2540                }
2541                let ex = ex as usize;
2542                for next in page_prefetch_positions(j, sel.len(), page_window) {
2543                    Self::moe_prefetch_host_expert(sel[next] as usize, m);
2544                }
2545                let keep = [
2546                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_GATE, ex as u16),
2547                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_UP, ex as u16),
2548                    crate::moe_cache::BlockId::new(il, crate::moe_cache::PROJ_DOWN, ex as u16),
2549                ];
2550                if worker_disk_prefetch && worker_window > 0 {
2551                    for next in worker_prefetch_positions(j, sel.len(), worker_window) {
2552                        Self::moe_prefetch_disk_expert(
2553                            e,
2554                            il,
2555                            sel[next] as usize,
2556                            m,
2557                            max_block,
2558                            &keep,
2559                        )?;
2560                    }
2561                } else if cache_dispatch
2562                    && !cpu_hybrid
2563                    && moe_prefetch_enabled()
2564                    && j + 1 < sel.len()
2565                {
2566                    let next = sel[j + 1] as usize;
2567                    Self::moe_prefetch_expert(e, il, next, m, max_block, &keep)?;
2568                }
2569                let [gate_q8, up_q8, down_q8] = [moe_q8; 3];
2570                if cache_dispatch && (gate_q8 || up_q8 || down_q8) {
2571                    // dp4a EXPERT PATH (MEMRA_MOE_Q8): quantize z-row once per token. Mixed expert
2572                    // layouts stay on the metadata-aware f32 path.
2573                    if (gate_q8 || up_q8) && tok_q8.is_none() {
2574                        tok_q8 = Some(e.quantize_q8_1_view(&zt, 1, n_embd)?);
2575                    }
2576                    let gate = if gate_q8 {
2577                        let (zq, zd) = tok_q8.as_ref().unwrap();
2578                        Self::moe_cached_gemm_q8(e, il, PROJ_GATE, ex, m, max_block, zq, zd)?
2579                    } else {
2580                        Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?
2581                    };
2582                    let up = if up_q8 {
2583                        let (zq, zd) = tok_q8.as_ref().unwrap();
2584                        Self::moe_cached_gemm_q8(e, il, PROJ_UP, ex, m, max_block, zq, zd)?
2585                    } else {
2586                        Self::moe_cached_gemm(e, il, PROJ_UP, ex, m, max_block, &zt)?
2587                    };
2588                    let mut act = e.uninit(n_ff_exp)?;
2589                    Self::ffn_act_scaled(
2590                        e,
2591                        cfg,
2592                        &gate,
2593                        &up,
2594                        m.gate_exps.macro_scale(ex),
2595                        m.up_exps.macro_scale(ex),
2596                        &mut act,
2597                        n_ff_exp,
2598                    )?;
2599                    let y = if down_q8 {
2600                        let (aq2, ad2) = e.quantize_q8_1(&act, 1, n_ff_exp)?;
2601                        Self::moe_cached_gemm_q8(e, il, PROJ_DOWN, ex, m, max_block, &aq2, &ad2)?
2602                    } else {
2603                        let actv = act.slice(0..n_ff_exp);
2604                        Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?
2605                    };
2606                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2607                    // down-proj macro folds into the accumulate weight (1.0 for non-macro archs).
2608                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2609                } else if cache_dispatch {
2610                    // SLRU residency cache: per-projection, dispatch the block (HIT => resident slot,
2611                    // MISS => staged slot) then run the SAME unchanged qmatvec_view from that slot.
2612                    // The bytes the kernel reads are byte-for-byte the same GGUF block (§B.3); the
2613                    // only difference between HIT and MISS is whether the memcpy_htod ran.
2614                    let gate = Self::moe_cached_gemm(e, il, PROJ_GATE, ex, m, max_block, &zt)?;
2615                    let up   = Self::moe_cached_gemm(e, il, PROJ_UP,   ex, m, max_block, &zt)?;
2616                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
2617                    Self::ffn_act_scaled(e, cfg, &gate, &up,
2618                        m.gate_exps.macro_scale(ex), m.up_exps.macro_scale(ex), &mut act, n_ff_exp)?;
2619                    let actv = act.slice(0..n_ff_exp);
2620                    let y = Self::moe_cached_gemm(e, il, PROJ_DOWN, ex, m, max_block, &actv)?;
2621                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2622                    // down-proj macro folds into the accumulate weight (post-matmul linear scale).
2623                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2624                } else if cache_frozen {
2625                    // A later prompt prime must not change the CPU/GPU assignment frozen after the
2626                    // first prime. Reuse every fixed resident projection directly and stage only a
2627                    // true miss through the ordinary scratch slot. This preserves the established
2628                    // f32-dequant numeric path while avoiding a full-bank reread on every prime.
2629                    let gate = Self::moe_frozen_gemm(
2630                        e,
2631                        il,
2632                        PROJ_GATE,
2633                        ex,
2634                        m,
2635                        max_block,
2636                        &zt,
2637                        &mut scratch_g,
2638                        g_len,
2639                    )?;
2640                    let up = Self::moe_frozen_gemm(
2641                        e,
2642                        il,
2643                        PROJ_UP,
2644                        ex,
2645                        m,
2646                        max_block,
2647                        &zt,
2648                        &mut scratch_u,
2649                        u_len,
2650                    )?;
2651                    let mut act = e.uninit(n_ff_exp)?;
2652                    Self::ffn_act_scaled(
2653                        e,
2654                        cfg,
2655                        &gate,
2656                        &up,
2657                        m.gate_exps.macro_scale(ex),
2658                        m.up_exps.macro_scale(ex),
2659                        &mut act,
2660                        n_ff_exp,
2661                    )?;
2662                    let actv = act.slice(0..n_ff_exp);
2663                    let y = Self::moe_frozen_gemm(
2664                        e,
2665                        il,
2666                        PROJ_DOWN,
2667                        ex,
2668                        m,
2669                        max_block,
2670                        &actv,
2671                        &mut scratch_d,
2672                        d_len,
2673                    )?;
2674                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2675                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2676                } else {
2677                    // Stage-1: stage gate/up/down for expert `ex` into the scratch slots, then GEMM.
2678                    // Lazy scratch: first no-cache expert allocates the 3 slots (uninit — stage_expert
2679                    // fully overwrites the byte range the GEMM reads).
2680                    if scratch_g.is_none() {
2681                        scratch_g = Some(e.alloc_u8_uninit(g_len)?);
2682                        scratch_u = Some(e.alloc_u8_uninit(u_len)?);
2683                        scratch_d = Some(e.alloc_u8_uninit(d_len)?);
2684                    }
2685                    let (sg, su, sd) = (scratch_g.as_mut().unwrap(), scratch_u.as_mut().unwrap(),
2686                                        scratch_d.as_mut().unwrap());
2687                    let gl = m.gate_exps.expert_layout(ex);
2688                    let ul = m.up_exps.expert_layout(ex);
2689                    let dl = m.down_exps.expert_layout(ex);
2690                    e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
2691                    let gate = e.qmatvec_view(sg, 0..gl.len, &zt, 1,
2692                        m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
2693
2694                    e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
2695                    let up = e.qmatvec_view(su, 0..ul.len, &zt, 1,
2696                        m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
2697
2698                    let mut act = e.uninit(n_ff_exp)?;  // activation fully overwrites
2699                    Self::ffn_act_scaled(e, cfg, &gate, &up,
2700                        m.gate_exps.macro_scale(ex), m.up_exps.macro_scale(ex), &mut act, n_ff_exp)?;
2701
2702                    e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
2703                    let actv = act.slice(0..n_ff_exp);
2704                    let y = e.qmatvec_view(sd, 0..dl.len, &actv, 1,
2705                        m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?;
2706
2707                    let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2708                    e.axpy_into(&y, w[j] * m.down_exps.macro_scale(ex), &mut dst, n_embd)?;
2709                }
2710            }
2711            if let Some(worker) = cpu_worker {
2712                let cpu_output = worker.wait().map_err(std::io::Error::other)?;
2713                let cpu_output = e.htod(&cpu_output)?;
2714                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
2715                e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
2716            }
2717            if cpu_hybrid && !cache_frozen && cpu_expert_profile_admit_enabled() {
2718                for (j, &ex) in sel.iter().enumerate() {
2719                    if cpu_mask[j] {
2720                        Self::moe_profile_admit_expert(e, il, ex as usize, m, max_block)?;
2721                    }
2722                }
2723            }
2724        }
2725
2726        // 3. SHARED EXPERT (ALWAYS-ON, no routing) on the SAME z — qwen35moe only. OLMoE and most
2727        //    vanilla MoE have NO shared expert (the shexp tensors are absent / `None`); skip it then.
2728        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
2729        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
2730        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
2731            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
2732        {
2733            let n_ff_sh = gate_shexp.out_features();  // 512
2734            // Q8 TRUNK-FUSION (decode t=1): gate_shexp+up_shexp are Q8_0 same-shape on the 35B —
2735            // ONE fused2 launch (also folds the two per-matmul re-quantizes of z into one).
2736            // Bit-identical per (tensor,row); falls back to the two matmul calls when ineligible.
2737            // Small-t (spec verify 2..15) rides matmul_decode_exact so shexp FP chains match the
2738            // t==1 decode chain per column (cuBLASLt n-dependence + dp4a-vs-mmvq class); real
2739            // prefill keeps the batched matmul. Activation routes through ffn_act (SiLU for
2740            // softmax archs, clamped swigluoai for M3 — identical to silu_mul when cfg.m3 is None).
2741            let verify_t = t > 1 && t < PRIME_MIN_T;
2742            let (sg_gate, sg_up) = if t == 1 {
2743                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
2744                    Some(pair) => pair,
2745                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
2746                }
2747            } else if verify_t {
2748                (e.matmul_decode_exact(gate_shexp, z, t)?, e.matmul_decode_exact(up_shexp, z, t)?)
2749            } else {
2750                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)   // [T, 512] each
2751            };
2752            let mut sa = e.uninit(t * n_ff_sh)?;  // activation fully overwrites
2753            Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
2754            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
2755                     else { e.matmul(down_shexp, &sa, t)? };     // [T, n_embd]
2756
2757            // shexp gate: qwen35moe sigmoid-gates via ffn_gate_inp_shexp (1-D ne=[n_embd] ->
2758            // out_f=1); M3 has no gate tensor -> weight 1.0. Decode + verify ride the fused
2759            // sigmoid-dot kernel (one fold order for both chains; kills the per-layer
2760            // cuBLASLt m=1 splitK GEMM — 40x/step, ~10% of the H100 q35 decode step).
2761            // SERVE ISOLATION (lane/concat-prime-exact, 2026-08-02): PREFILL rides it too.
2762            // This out_f=1 cuBLASLt GEMV is the SECOND m-dependent op in the trunk (probed:
2763            // rows [0,19) move by 1.07e-4 between m=74 and m=75 while sigmoid_dot_rows is
2764            // BIT-IDENTICAL — allw-shexpgate-o35b.log). The gate multiplies the shared
2765            // expert's contribution into every token's residual, so under cross-request
2766            // concat prefill a session's hidden state depended on its co-arrivals' token
2767            // count. MEMRA_ROUTER_PREFILL_EXACT=0 reverts to the batched cuBLASLt linear.
2768            let g = match &m.gate_inp_shexp {
2769                Some(gate_inp_shexp) => {
2770                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
2771                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
2772                    } else {
2773                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
2774                        let mut g = e.uninit(t)?;  // sigmoid fully overwrites
2775                        e.sigmoid(&gs, &mut g, t)?;
2776                        g
2777                    }
2778                }
2779                None => e.htod(&vec![1.0f32; t])?,
2780            };
2781            // moe_out[r, :] += sh[r, :] * g[r]   (per-token scalar gate; g=1 ungated)
2782            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
2783        }
2784
2785        Ok(moe_out)
2786    }
2787
2788    /// Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every
2789    /// token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.
2790    pub fn stage1_h2d_per_token(&self) -> u64 {
2791        use crate::hybrid::Ffn;
2792        let n_used = self.cfg.moe.as_ref().map(|m| m.expert_used_count as u64).unwrap_or(0);
2793        let mut bytes = 0u64;
2794        for l in self.layers.iter() {
2795            if let Ffn::Moe(m) = &l.ffn {
2796                bytes += n_used * (m.gate_exps.max_expert_bytes() + m.up_exps.max_expert_bytes()
2797                                   + m.down_exps.max_expert_bytes()) as u64;
2798            }
2799        }
2800        bytes
2801    }
2802
2803    /// Largest expert block (bytes) across ALL MoE layers + the MTP head — the fixed cache slot size.
2804    /// UD/dynamic GGUFs quant different layers differently, so `expert_stride` varies per layer; the
2805    /// residency cache slots are fixed-address and must fit any block, so size to this global max.
2806    pub(crate) fn max_moe_block(&self) -> usize {
2807        use crate::hybrid::Ffn;
2808        let mut mx = 0usize;
2809        let mut scan = |ffn: &Ffn| {
2810            if let Ffn::Moe(m) = ffn {
2811                mx = mx.max(m.gate_exps.max_expert_bytes())
2812                       .max(m.up_exps.max_expert_bytes())
2813                       .max(m.down_exps.max_expert_bytes());
2814            }
2815        };
2816        for l in self.layers.iter() { scan(&l.ffn); }
2817        if let Some(mtp) = self.mtp.as_ref() { scan(&mtp.ffn); }
2818        mx
2819    }
2820
2821    /// Exact retained projection lengths for cache class sizing. Pruned ids keep router positions
2822    /// but have no bytes and therefore consume no residency slot.
2823    pub(crate) fn moe_cache_block_sizes(&self) -> Vec<usize> {
2824        use crate::hybrid::Ffn;
2825        let mut sizes = Vec::new();
2826        let mut scan = |ffn: &Ffn| {
2827            let Ffn::Moe(m) = ffn else { return };
2828            for ex in 0..m.gate_exps.n_expert {
2829                if m.active_experts.as_ref().is_some_and(|active| !active[ex]) {
2830                    continue;
2831                }
2832                for exps in [&m.gate_exps, &m.up_exps, &m.down_exps] {
2833                    let len = exps.expert_layout(ex).len;
2834                    if len > 0 {
2835                        sizes.push(len);
2836                    }
2837                }
2838            }
2839        };
2840        for layer in &self.layers {
2841            scan(&layer.ffn);
2842        }
2843        if let Some(mtp) = &self.mtp {
2844            scan(&mtp.ffn);
2845        }
2846        sizes
2847    }
2848
2849    /// Persist the frozen residency set so a later process can restage it directly and skip
2850    /// the profiling warmup. Plain text: a versioned header binding slot geometry, then one
2851    /// `layer proj ex` triple per line. A mismatched or stale profile is rejected at load
2852    /// (header check) or degrades to fewer restaged blocks (per-id checks); either way the
2853    /// post-freeze argmax gate still validates the serving assignment.
2854    pub fn save_cpu_expert_residency_profile(
2855        &self,
2856        e: &Engine,
2857        path: &std::path::Path,
2858    ) -> Result<(), Box<dyn std::error::Error>> {
2859        let Some(ids) = e.export_moe_residency() else {
2860            return Err("no MoE residency cache to persist".into());
2861        };
2862        let mut body = format!(
2863            "memra-freeze-profile v1 max_block={} blocks={}\n",
2864            self.max_moe_block(),
2865            ids.len()
2866        );
2867        for (layer, proj, ex) in &ids {
2868            body.push_str(&format!("{layer} {proj} {ex}\n"));
2869        }
2870        let tmp = path.with_extension("tmp");
2871        std::fs::write(&tmp, body)?;
2872        std::fs::rename(&tmp, path)?;
2873        println!(
2874            "[moe-cache] freeze profile saved: {} blocks -> {}",
2875            ids.len(),
2876            path.display()
2877        );
2878        Ok(())
2879    }
2880
2881    /// Restage a saved freeze profile and freeze immediately, skipping the profiling warmup.
2882    /// Returns false (leaving the cache untouched for a normal warmup) when the profile is
2883    /// missing or its header does not match this model's slot geometry.
2884    pub fn restore_cpu_expert_residency_profile(
2885        &self,
2886        e: &Engine,
2887        path: &std::path::Path,
2888    ) -> Result<bool, Box<dyn std::error::Error>> {
2889        use crate::hybrid::Ffn;
2890        use crate::moe_cache::BlockId;
2891        let Ok(content) = std::fs::read_to_string(path) else {
2892            return Ok(false);
2893        };
2894        let mut lines = content.lines();
2895        let Some(header) = lines.next() else { return Ok(false) };
2896        let expected = format!("memra-freeze-profile v1 max_block={}", self.max_moe_block());
2897        if !header.starts_with(&expected) {
2898            println!(
2899                "[moe-cache] freeze profile ignored (geometry mismatch): {}",
2900                path.display()
2901            );
2902            return Ok(false);
2903        }
2904        let mut by_layer: std::collections::HashMap<u16, Vec<BlockId>> =
2905            std::collections::HashMap::new();
2906        for line in lines {
2907            let mut fields = line.split_whitespace();
2908            let (Some(layer), Some(proj), Some(ex)) =
2909                (fields.next(), fields.next(), fields.next())
2910            else {
2911                continue;
2912            };
2913            let (Ok(layer), Ok(proj), Ok(ex)) =
2914                (layer.parse::<u16>(), proj.parse::<u8>(), ex.parse::<u16>())
2915            else {
2916                continue;
2917            };
2918            by_layer
2919                .entry(layer)
2920                .or_default()
2921                .push(BlockId::new(layer, proj, ex));
2922        }
2923        let requested: usize = by_layer.values().map(Vec::len).sum();
2924        if requested == 0 {
2925            return Ok(false);
2926        }
2927        let max_block = self.max_moe_block();
2928        let mut restaged = 0usize;
2929        let mut stage_layer = |layer_index: u16,
2930                               ffn: &Ffn|
2931         -> Result<(), Box<dyn std::error::Error>> {
2932            let Ffn::Moe(m) = ffn else { return Ok(()) };
2933            let Some(ids) = by_layer.get(&layer_index) else {
2934                return Ok(());
2935            };
2936            e.with_moe_cache(max_block, |cache, eng| {
2937                for id in ids {
2938                    if cache.restage_block(*id, m, eng)? {
2939                        restaged += 1;
2940                    }
2941                }
2942                Ok(())
2943            })
2944        };
2945        for (index, layer) in self.layers.iter().enumerate() {
2946            stage_layer(index as u16, &layer.ffn)?;
2947        }
2948        if let Some(mtp) = self.mtp.as_ref() {
2949            stage_layer(u16::MAX, &mtp.ffn)?;
2950        }
2951        e.freeze_moe_cache();
2952        println!(
2953            "[moe-cache] freeze profile restored: {restaged}/{requested} blocks restaged from {}",
2954            path.display()
2955        );
2956        Ok(true)
2957    }
2958
2959    /// Freeze the heterogeneous CPU/GPU split after the caller's discarded profile warmup.
2960    pub fn freeze_cpu_expert_residency(
2961        &self,
2962        e: &Engine,
2963    ) -> Result<(), Box<dyn std::error::Error>> {
2964        e.freeze_moe_cache();
2965        Ok(())
2966    }
2967
2968    /// FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the
2969    /// standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows
2970    /// the model's activation exactly.
2971    pub fn ffn_act(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
2972               act: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2973        Self::ffn_act_scaled(e, cfg, gate, up, 1.0, 1.0, act, n)
2974    }
2975
2976    /// ffn_act with per-tensor post-matmul macro-scales folded in (gs/us == 1.0 -> identical
2977    /// float ops to ffn_act; used by the ModelOpt NVFP4 expert path where each expert tensor
2978    /// carries a `weight_scale_2`).
2979    #[allow(clippy::too_many_arguments)]
2980    pub(crate) fn ffn_act_scaled(e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
2981               gs: f32, us: f32, act: &mut CudaSlice<f32>, n: usize)
2982               -> Result<(), Box<dyn std::error::Error>> {
2983        if let Some(m3) = cfg.m3.as_ref() {
2984            return e.swigluoai_mul_scaled(gate, up, gs, us, m3.swiglu_alpha, m3.swiglu_limit, act, n);
2985        }
2986        if gs == 1.0 && us == 1.0 { return e.silu_mul(gate, up, act, n); }
2987        e.silu_mul_scaled(gate, up, gs, us, act, n)
2988    }
2989
2990    /// Routing for the whole batch: returns (sel [T*n_used] expert ids, w [T*n_used] renorm weights),
2991    /// token-major. Default = the Stage-1 host path (dtoh logits, softmax-256, stable DESC top-k,
2992    /// renorm). MEMRA_FUSED_ROUTER = the device kernel (§A) which reproduces the same numerics; we
2993    /// still dtoh the tiny [T,n_used] sel/w buffers (64 B/token vs 1 KB/token) — the host loop
2994    /// indexes HostExps.bytes on the CPU to choose the DMA source (§A.2 output staging).
2995    fn moe_route(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
2996                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
2997        Self::moe_route_cfg(e, logits, t, n_expert, n_used, None, None, None)
2998    }
2999
3000    /// DeepSeek-V3-class sigmoid routing (MiniMax-M3, Hy3), host oracle. Reference:
3001    /// M3/Hy3 modeling code — scores = sigmoid(logits); selection over scores + expert bias
3002    /// (M3 `e_score_correction_bias` / Hy3 `expert_bias`, both surfaced as `exp_probs_b`);
3003    /// weights = un-biased scores of the selected experts, sum-normalized when `route_norm`,
3004    /// x scaling factor (M3 routed_scaling_factor 2.0 / Hy3 router_scaling_factor 2.826).
3005    /// `sig` = (scaling_factor, route_norm) from `cfg.sigmoid_router()`; softmax archs pass
3006    /// None -> the qwen35moe/OLMoE path below.
3007    fn moe_route_cfg(e: &Engine, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize,
3008                     bias: Option<&[f32]>, sig: Option<(f32, bool)>, active: Option<&[bool]>)
3009                 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3010        if let Some((sf, route_norm)) = sig {
3011            // sigmoid routing. Host path only for now (fused-router kernel is softmax-top-k).
3012            let lg = e.dtoh(logits)?;
3013            return Self::moe_route_sigmoid_host(
3014                &lg, t, n_expert, n_used, bias, sf, route_norm, active,
3015            );
3016        }
3017        // LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router DEFAULT ON (MEMRA_FUSED_ROUTER=0
3018        // rollback) via the single-sync pinned readback — softmax arch only; the M3 sigmoid arm
3019        // above returns before this (host path until a sigmoid fused-router kernel exists).
3020        if active.is_none() && !matches!(std::env::var("MEMRA_FUSED_ROUTER").as_deref(), Ok("0")) {
3021            return e.moe_router_topk_host(logits, t, n_expert, n_used);
3022        }
3023        // Host oracle (the §D bit-identity reference).
3024        let lg = e.dtoh(logits)?;   // [T*n_expert] host
3025        let mut sel = vec![0u32; t * n_used];
3026        let mut w_out = vec![0f32; t * n_used];
3027        for tok in 0..t {
3028            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3029            // softmax over ALL n_expert (stable: subtract max)
3030            let maxl = row.iter().enumerate()
3031                .filter(|(i, _)| active.is_none_or(|mask| mask[*i]))
3032                .map(|(_, &x)| x).fold(f32::NEG_INFINITY, f32::max);
3033            let mut probs = vec![0f32; n_expert];
3034            let mut den = 0f32;
3035            for i in 0..n_expert {
3036                if active.is_some_and(|mask| !mask[i]) { continue; }
3037                let x = (row[i] - maxl).exp(); probs[i] = x; den += x;
3038            }
3039            for p in probs.iter_mut() { *p /= den; }
3040            // stable DESC sort: prob DESC, ascending-index tiebreak.
3041            let mut idx: Vec<usize> = (0..n_expert)
3042                .filter(|&i| active.is_none_or(|mask| mask[i])).collect();
3043            idx.sort_by(|&a, &b| probs[b].total_cmp(&probs[a]).then(a.cmp(&b)));
3044            let sl = &idx[..n_used];
3045            let mut wv: Vec<f32> = sl.iter().map(|&i| probs[i]).collect();
3046            let mut ws: f32 = wv.iter().sum();
3047            ws = ws.max(6.103515625e-5_f32);  // F16 smallest normal, clamp BEFORE divide
3048            for x in wv.iter_mut() { *x /= ws; }
3049            for j in 0..n_used {
3050                sel[tok * n_used + j] = sl[j] as u32;
3051                w_out[tok * n_used + j] = wv[j];
3052            }
3053        }
3054        Ok((sel, w_out))
3055    }
3056
3057    #[allow(clippy::too_many_arguments)]
3058    fn moe_route_sigmoid_with_input(
3059        e: &Engine,
3060        logits: &CudaSlice<f32>,
3061        input: &CudaSlice<f32>,
3062        t: usize,
3063        n_expert: usize,
3064        n_used: usize,
3065        bias: Option<&[f32]>,
3066        (sf, route_norm): (f32, bool),
3067        active: Option<&[bool]>,
3068    ) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
3069        let (lg, input) = e.dtoh_pair(logits, input)?;
3070        let (sel, w) =
3071            Self::moe_route_sigmoid_host(&lg, t, n_expert, n_used, bias, sf, route_norm, active)?;
3072        Ok((sel, w, input))
3073    }
3074
3075    /// Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after
3076    /// residency freeze: the worker filters against a static snapshot of the frozen HBM set.
3077    /// Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias,
3078    /// active mask, prebuilt projection descriptors) so no model reference escapes.
3079    pub fn start_moe_prefetch_predictor(
3080        &self,
3081        e: &Engine,
3082        cfg: &ModelConfig,
3083    ) -> Result<(), Box<dyn std::error::Error>> {
3084        use crate::hybrid::Ffn;
3085        let Some(sig) = cfg.sigmoid_router() else {
3086            return Err("prefetch predictor requires a sigmoid-router arch".into());
3087        };
3088        let resident: std::collections::HashSet<(u16, u8, u16)> = e
3089            .export_moe_residency()
3090            .ok_or("prefetch predictor needs the frozen MoE residency cache")?
3091            .into_iter()
3092            .collect();
3093        let mut layers = Vec::new();
3094        for (index, layer) in self.layers.iter().enumerate() {
3095            let Ffn::Moe(m) = &layer.ffn else { continue };
3096            let crate::model::GpuTensor::Float { data, .. } = &m.gate_inp else { continue };
3097            let router = e.dtoh(data)?;
3098            let n_expert = m.gate_exps.n_expert;
3099            let n_embd = m.gate_exps.in_f;
3100            if router.len() != n_embd * n_expert {
3101                continue;
3102            }
3103            let build = |exps: &crate::model::HostExps| {
3104                (0..n_expert)
3105                    .map(|expert| crate::cpu_experts::predictor_projection(exps, expert))
3106                    .collect::<Vec<_>>()
3107            };
3108            layers.push((index as u16, crate::cpu_experts::PredictLayerInit {
3109                router,
3110                bias: m.exp_probs_b.clone(),
3111                active: m.active_experts.clone(),
3112                n_embd,
3113                n_used: cfg
3114                    .moe
3115                    .as_ref()
3116                    .map(|moe| moe.expert_used_count as usize)
3117                    .ok_or("prefetch predictor requires MoE config")?,
3118                sig,
3119                weights_n_expert: n_expert,
3120                gate: build(&m.gate_exps),
3121                up: build(&m.up_exps),
3122                down: build(&m.down_exps),
3123            }));
3124        }
3125        crate::cpu_experts::start_prefetch_predictor(layers, resident)
3126            .map_err(|error| error.into())
3127    }
3128
3129    /// Crate-visible sigmoid-routing oracle for the prefetch predictor: identical selection
3130    /// math to the runtime router, applied to host-computed lookahead logits.
3131    #[allow(clippy::too_many_arguments)]
3132    pub(crate) fn moe_route_sigmoid_host_public(
3133        logits: &[f32],
3134        t: usize,
3135        n_expert: usize,
3136        n_used: usize,
3137        bias: Option<&[f32]>,
3138        sf: f32,
3139        route_norm: bool,
3140        active: Option<&[bool]>,
3141    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3142        Self::moe_route_sigmoid_host(logits, t, n_expert, n_used, bias, sf, route_norm, active)
3143    }
3144
3145    #[allow(clippy::too_many_arguments)]
3146    fn moe_route_sigmoid_host(
3147        lg: &[f32],
3148        t: usize,
3149        n_expert: usize,
3150        n_used: usize,
3151        bias: Option<&[f32]>,
3152        sf: f32,
3153        route_norm: bool,
3154        active: Option<&[bool]>,
3155    ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
3156        if lg.len() != t * n_expert {
3157            return Err(format!(
3158                "sigmoid router logits length mismatch: got {}, expected {}",
3159                lg.len(),
3160                t * n_expert,
3161            )
3162            .into());
3163        }
3164        let mut sel = vec![0u32; t * n_used];
3165        let mut w_out = vec![0f32; t * n_used];
3166        for tok in 0..t {
3167            let row = &lg[tok * n_expert..(tok + 1) * n_expert];
3168            let scores: Vec<f32> = row.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect();
3169            // selection score = sigmoid + bias; weight = plain sigmoid.
3170            let selsc: Vec<f32> = match bias {
3171                Some(b) => scores.iter().zip(b).map(|(s, bb)| s + bb).collect(),
3172                None => scores.clone(),
3173            };
3174            let mut idx: Vec<usize> = (0..n_expert)
3175                .filter(|&i| active.is_none_or(|mask| mask[i]))
3176                .collect();
3177            idx.sort_by(|&a, &b| selsc[b].total_cmp(&selsc[a]).then(a.cmp(&b)));
3178            let sl = &idx[..n_used];
3179            let mut wv: Vec<f32> = sl.iter().map(|&i| scores[i]).collect();
3180            if route_norm {
3181                let ws: f32 = wv.iter().sum::<f32>().max(1e-20);
3182                for x in wv.iter_mut() {
3183                    *x = *x / ws * sf;
3184                }
3185            } else {
3186                for x in wv.iter_mut() {
3187                    *x *= sf;
3188                }
3189            }
3190            for j in 0..n_used {
3191                sel[tok * n_used + j] = sl[j] as u32;
3192                w_out[tok * n_used + j] = wv[j];
3193            }
3194        }
3195        Ok((sel, w_out))
3196    }
3197
3198    /// LAUNCH-STRUCTURE STAGE 3: the ZERO-DtoH fully-resident MoE FFN. Caller guarantees the
3199    /// layer's device pointer row exists (checked under the cache lock). Router top-k runs on
3200    /// device; sel/w are consumed by the `_dev` matvec twins directly; NOTHING crosses PCIe.
3201    /// Same numerics as the fused-router + gdec chain (kernel-level bit-identity, see the
3202    /// MoE PREFILL PAIR-BATCH: host routing (sel/w like the sequential path), then 5 launches
3203    /// TOTAL per layer (quantize z, gate-pairs, up-pairs, silu, act-quantize, down-pairs,
3204    /// scatter) regardless of T or expert count. Bit-identity class: per (pair,row) dot =
3205    /// qmatvec_expert_q8 order; per-token accumulation slot-ordered (scatter kernel).
3206    fn moe_ffn_pairs(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, logits: &CudaSlice<f32>,
3207                     t: usize, cfg: &ModelConfig)
3208                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3209        let moe = cfg.moe.as_ref().unwrap();
3210        let n_embd = cfg.n_embd as usize;
3211        let n_expert = moe.expert_count as usize;
3212        let n_used = moe.expert_used_count as usize;
3213        let n_ff_exp = moe.expert_ff_length as usize;
3214        let dev = m.dev_exps.as_ref().unwrap();
3215        // WALL-GAP ARC: interleaved gate/up slab strides (see moe_ffn_dev).
3216        let (rbg_d, rbu_d) = if dev.gu_il {
3217            let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3218        } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3219
3220        let (sel_all, w_all) = Self::moe_route(e, logits, t, n_expert, n_used)?;
3221        let n_pairs = t * n_used;
3222        // pair arrays: pair p = (token p/n_used, slot p%n_used) — ALREADY slot-ordered per token,
3223        // so the CSR is trivial: tok_pair_off[tok] = tok*n_used, ids identity.
3224        let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
3225        let pair_ex:  Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
3226        let pair_w:   Vec<f32> = w_all.clone();
3227        let tok_off:  Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
3228        let tok_ids:  Vec<i32> = (0..n_pairs as i32).collect();
3229        let pt = e.htod_i32(&pair_tok)?;
3230        let px = e.htod_i32(&pair_ex)?;
3231        let pw = e.htod(&pair_w)?;
3232        let toff = e.htod_i32(&tok_off)?;
3233        let tids = e.htod_i32(&tok_ids)?;
3234
3235        // z quantized ONCE for all tokens; gate/up pair matvecs; silu; act quantize; down; scatter.
3236        // EXPERT-MAJOR CSR (rung 2): pairs grouped by expert -> the kernel reuses each weight
3237        // row across the expert's token group (llama-MMQ's core win). Host grouping is O(pairs).
3238        let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
3239        for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
3240        let mut ex_ids: Vec<i32> = Vec::new();
3241        let mut ex_off: Vec<i32> = vec![0];
3242        let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
3243        for (ex, list) in by_ex.iter().enumerate() {
3244            if list.is_empty() { continue; }
3245            ex_ids.push(ex as i32);
3246            ex_pairs.extend_from_slice(list);
3247            ex_off.push(ex_pairs.len() as i32);
3248        }
3249        let n_active = ex_ids.len();
3250        let exi = e.htod_i32(&ex_ids)?;
3251        let exo = e.htod_i32(&ex_off)?;
3252        let exp_d = e.htod_i32(&ex_pairs)?;
3253        let _ = &px;   // pair-major twin keeps it; em path uses CSR
3254
3255        // INT8-MMA EXPERT MMQ (MEMRA_MOE_MMA=1, opt-in): the m16n8k16.s8 tensor-core analog of the
3256        // _dec dp4a kernel (cu/mmq_iq_experts.cu). Same CSR grouping; per-expert matvec runs as a
3257        // 128x128-tile int8 MMA GEMM over the expert's token group. Weight IQ nibbles decode to int8
3258        // at tile-load + per-32 float scale; activation is q8_1_mmq (D4, same quant class as dp4a).
3259        // FP-ORDER differs from dp4a (MMA reduction) — logits SHIFT, gated on argmax/spec/closeness,
3260        // NOT byte-identity (like the W4A8 path). Requires IQ3_S/IQ4_XS + in_f % 256 == 0.
3261        // t >= 16 (GEMM_M-class rule): the MMA tile needs token volume (crossover ~200 tok/expert;
3262        // microbench: dp4a wins at tiny groups). ALSO an exactness requirement — spec verify
3263        // batches (t=2..K+2) must ride the dp4a path whose FP order matches the T=1 decode chain,
3264        // else K=1 self-consistency FAILs (caught 2026-07-06: MMA at T=2 flipped a verify argmax).
3265        // DEFAULT ON (2026-07-06, third flip — this time with the real culprit fixed): the
3266        // "MMA prime breaks spec" failure was the ROUTER's cuBLASLt n-dependence (d994271),
3267        // not MMA's own FP order — both this and the k-quant arms were innocent suspects whose
3268        // margin shifts surfaced the router bug. With the router decode-exact at verify t, the
3269        // full battery is green with MMA on (spec p1/p2/p3 PASS, raw K=1..8 PASS, argmax MATCH,
3270        // pp6257 2862 = 2.1x dec). t>=16 floor still required: verify batches must ride dp4a
3271        // (dispatch parity with the T=1 decode chain). MEMRA_MOE_MMA=0 rollback;
3272        // MEMRA_MOE_MMA_T overrides the floor (bisect seam).
3273        static MMA_T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
3274        let mma_t = *MMA_T.get_or_init(|| {
3275            std::env::var("MEMRA_MOE_MMA_T").ok().and_then(|v| v.parse().ok()).unwrap_or(16)
3276        });
3277        let use_mma = std::env::var("MEMRA_MOE_MMA").map(|v| v != "0").unwrap_or(true)
3278            && t >= mma_t
3279            && q8_expert_dec_supported(m.gate_exps.qtype) && q8_expert_dec_supported(m.up_exps.qtype)
3280            && q8_expert_dec_supported(m.down_exps.qtype)
3281            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3282        // GROUPED f16 LANE admission (MEMRA_MOE_F16G, rounds 46-49): its own door, no longer a
3283        // subset of the MMQ arm — q35's k-quant stragglers (Q6_K/Q4_K down, one Q3_K gate/up
3284        // layer) fail q8_expert_dec_supported but dequant fine to f16, so f16g must be able to
3285        // take a layer the MMQ arm would reject. Same t >= mma_t floor as MMA: decode and
3286        // spec-verify batches must ride the dp4a path whose FP order matches the T=1 chain.
3287        // MODE 2 (sm_120a naked default, lane/f16g-default-rearb 2026-08-02): every layer
3288        // whose three projections pass f16g_proj_ok rides the sk visitor with direct tile
3289        // loaders — with IQ4_XS/IQ3_S direct coverage the sk arm beats the int8-MMA MMQ
3290        // tiles on the IQ-bank models too (q35 board-2048 +33.9%, KAT pp512 +46.7%).
3291        // AUTO-KQUANT (mode 3, MEMRA_MOE_F16G=3, lane/q4k-expert-prefill): admit f16g ONLY
3292        // where the MMA arm can't take the layer's QTYPES — the k-quant expert class whose
3293        // baseline is the per-pair _em fallback (Ornith-35B Q4_K: board-2048 3.14x). Its
3294        // MMA-capable carve-out (IQ3_S/IQ4_XS/Q4_0 banks on MMQ) was priced pre-IQ-direct;
3295        // it survives as the rollback seam. Keyed on qtype capability, NOT use_mma, so
3296        // MEMRA_MOE_MMA=0 stays a pure dp4a rollback seam.
3297        let mma_capable = q8_expert_dec_supported(m.gate_exps.qtype)
3298            && q8_expert_dec_supported(m.up_exps.qtype)
3299            && q8_expert_dec_supported(m.down_exps.qtype)
3300            && n_embd % 256 == 0 && n_ff_exp % 256 == 0;
3301        let f16g_mode = crate::moe_f16g_mode();
3302        let f16g = f16g_mode != 0 && t >= mma_t
3303            && (f16g_mode != 3 || !mma_capable)
3304            && f16g_proj_ok(m.gate_exps.qtype, n_embd)
3305            && f16g_proj_ok(m.up_exps.qtype, n_embd)
3306            && f16g_proj_ok(m.down_exps.qtype, n_ff_exp);
3307        if use_mma || f16g {
3308            // GROUPED f16 LANE (MEMRA_MOE_F16G, rounds 46-49, experimental door): dequant
3309            // the active experts ONCE per projection to f16 and run one grouped f16 GEMM over
3310            // the CSR groups. f16-mirror numeric class — argmax/spec gated before promotion.
3311            // NOTE: the pair-gather kernel needs pair p ordered EXPERT-MAJOR (ex_pairs order);
3312            // the y rows come back in that same CSR order, so gate/up/down all stay pair-major
3313            // in ex_pairs order — but moe_pairs_silu_mul and the scatter consume PAIR-ID order.
3314            // We therefore gather activations per ex_pairs and scatter y back through ex_pairs.
3315            let y_down = if f16g {
3316                // CSR order end-to-end: gather z rows by the pair's TOKEN (pair p's token is
3317                // p / n_used — the trivial CSR above), silu in CSR order (elementwise), one
3318                // permute at the very end back to pair-id order for the scatter.
3319                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
3320                let csr_tok_d = e.htod_i32(&csr_tok)?;
3321                let (z_f16, z_s) = e.moe_f16g_act(z, Some(&csr_tok_d), n_embd, n_pairs)?;
3322                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
3323                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3324                                              m.gate_exps.qtype, rbg_d)?;
3325                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
3326                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
3327                                              m.up_exps.qtype, rbu_d)?;
3328                let act_csr = e.moe_pairs_silu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
3329                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
3330                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
3331                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
3332                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
3333                e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?
3334            } else {
3335            // gate/up: activation = z, token-major over t tokens; pair_tok gathers the routed row.
3336            let z_scr = e.mmq_iq_quantize_act(z, n_embd, t)?;
3337            let gate = e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3338                                        n_embd, n_ff_exp, n_active, n_pairs, t,
3339                                        m.gate_exps.qtype, rbg_d)?;
3340            let up = e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
3341                                      n_embd, n_ff_exp, n_active, n_pairs, t,
3342                                      m.up_exps.qtype, rbu_d)?;
3343            // down: activation = silu(gate)*up, pair-major [n_pairs, n_ff_exp]; pair_tok =
3344            // identity. FUSED ACT-EPILOGUE (default on): one launch computes the activation in
3345            // registers and writes ONLY the quantized scratch — the two-pass chain
3346            // (moe_pairs_silu_mul writes act f32, mmq_iq_quantize_act re-reads it) is the
3347            // MEMRA_MOE_FUSE_ACTQ=0 rollback. Scratch bytes are BYTE-IDENTICAL (kernel-check).
3348            let a_scr = if crate::moe_fuse_actq_on() {
3349                e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 0)?
3350            } else {
3351                let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3352                e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
3353            };
3354            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3355            let pself = e.htod_i32(&pair_self)?;
3356            e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
3357                             n_ff_exp, n_embd, n_active, n_pairs, n_pairs,
3358                             m.down_exps.qtype, m.down_exps.row_bytes)?
3359            };
3360            let mut moe_out = e.uninit(t * n_embd)?;
3361            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3362            if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3363                (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3364            {
3365                let n_ff_sh = gate_shexp.out_features();
3366                let sg_gate = e.matmul(gate_shexp, z, t)?;
3367                let sg_up = e.matmul(up_shexp, z, t)?;
3368                let mut sa = e.uninit(t * n_ff_sh)?;
3369                Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3370                let sh = e.matmul(down_shexp, &sa, t)?;
3371                // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
3372                // SERVE ISOLATION (lane/concat-prime-exact): the fused sigmoid-dot is the
3373                // m-INVARIANT form — see the sequential arm's note. This is the PAIRS arm,
3374                // i.e. the one real prefill actually takes on a resident-expert MoE model,
3375                // so the concat-prime isolation fix has to land here as well.
3376                let g = match &m.gate_inp_shexp {
3377                    Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3378                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3379                    }
3380                    Some(gate_inp_shexp) => {
3381                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3382                        let mut g = e.uninit(t)?;
3383                        e.sigmoid(&gs, &mut g, t)?;
3384                        g
3385                    }
3386                    None => e.htod(&vec![1.0f32; t])?,
3387                };
3388                e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3389            }
3390            return Ok(moe_out);
3391        }
3392
3393        // DECODE-ONCE MMQ (rung 3, MEMRA_MOE_DEC=1 default-on): dequant each weight group once per
3394        // (row,group) then dp4a across the expert's tokens. _em re-decoded per token (NEUTRAL).
3395        let dec = std::env::var("MEMRA_MOE_DEC").map(|v| v != "0").unwrap_or(true);
3396        let matvec = |proj, exi: &_, exo: &_, exp_d: &_, pt: &_, aq: &_, ad: &_,
3397                      inf, outf, qtype, rb| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3398            // _dec's decode-once extractors are IQ-only; k-quant expert layers take the _em dot path.
3399            let dec = dec && q8_expert_dec_supported(qtype);
3400            if dec { e.moe_pairs_matvec_q8_dec(&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3401                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3402            else   { e.moe_pairs_matvec_q8_em (&dev.ptr_row, proj, exi, exo, exp_d, pt, aq, ad,
3403                                               inf, outf, n_expert, n_active, n_pairs, qtype, rb) }
3404        };
3405        let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3406        let gate = matvec(0, &exi, &exo, &exp_d, &pt, &zq, &zd,
3407                          n_embd, n_ff_exp, m.gate_exps.qtype, rbg_d)?;
3408        let up = matvec(1, &exi, &exo, &exp_d, &pt, &zq, &zd,
3409                        n_embd, n_ff_exp, m.up_exps.qtype, rbu_d)?;
3410        let act = e.moe_pairs_silu_mul(&gate, &up, n_pairs * n_ff_exp)?;
3411        let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3412        // down consumes PAIR-major activation rows: pair_tok = identity.
3413        let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
3414        let pself = e.htod_i32(&pair_self)?;
3415        let y_down = matvec(2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
3416                            n_ff_exp, n_embd, m.down_exps.qtype, m.down_exps.row_bytes)?;
3417        let mut moe_out = e.uninit(t * n_embd)?;   // scatter fully overwrites per (token,col)
3418        e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
3419
3420        // SHARED EXPERT epilogue — same as the other paths.
3421        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
3422        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
3423        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3424            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3425        {
3426            let n_ff_sh = gate_shexp.out_features();
3427            let sg_gate = e.matmul(gate_shexp, z, t)?;
3428            let sg_up = e.matmul(up_shexp, z, t)?;
3429            let mut sa = e.uninit(t * n_ff_sh)?;
3430            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3431            let sh = e.matmul(down_shexp, &sa, t)?;
3432            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
3433            // m-INVARIANT fused sigmoid-dot under router_prefill_exact (serve isolation,
3434            // lane/concat-prime-exact) — every shexp-gate arm shares the same form so a
3435            // dispatch choice cannot change bits.
3436            let g = match &m.gate_inp_shexp {
3437                Some(gate_inp_shexp) if crate::router_prefill_exact_on() => {
3438                    e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3439                }
3440                Some(gate_inp_shexp) => {
3441                    let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3442                    let mut g = e.uninit(t)?;
3443                    e.sigmoid(&gs, &mut g, t)?;
3444                    g
3445                }
3446                None => e.htod(&vec![1.0f32; t])?,
3447            };
3448            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3449        }
3450        Ok(moe_out)
3451    }
3452
3453    /// kernel headers); the shared-expert epilogue is byte-identical to moe_ffn_sequential's.
3454    #[allow(clippy::too_many_arguments)]
3455    #[allow(clippy::too_many_arguments)]
3456    fn moe_ffn_dev(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>,
3457                   zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, logits: &CudaSlice<f32>,
3458                   t: usize, cfg: &ModelConfig, il: u16, max_block: usize)
3459                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3460        let moe = cfg.moe.as_ref().unwrap();
3461        let n_embd = cfg.n_embd as usize;
3462        let n_expert = moe.expert_count as usize;
3463        let n_used = moe.expert_used_count as usize;
3464        let n_ff_exp = moe.expert_ff_length as usize;
3465
3466        // device top-k: sel [t, n_used] i32, w [t, n_used] f32 — stays on device.
3467        let (sel_d, mut w_d) = e.moe_router_topk(logits, t, n_expert, n_used)?;
3468        // Down-projection macro fold (compressed-tensors NVFP4 artifacts): one tiny launch,
3469        // skipped entirely for macro-free experts (every k-quant GGUF).
3470        if m.has_macros {
3471            e.moe_w_scale_by_expert(&mut w_d, &sel_d, &m.dev_macros, n_expert, t * n_used)?;
3472        }
3473
3474        // moe_out rows are FULLY overwritten by moe_down8_fma_dev — uninit (stage-2 rule).
3475        let mut moe_out = e.uninit(t * n_embd)?;
3476
3477        // RESIDENT-EXPERTS arm: the pointer row comes from the load-time slab (no cache, no
3478        // lock). Same kernels/loop as the SLRU arm below — only the row's provenance differs.
3479        if let Some(dev) = m.dev_exps.as_ref() {
3480            // WALL-GAP ARC: interleaved gate/up slab (MEMRA_MOE_GU_IL) -> both projections use
3481            // the combined stride; up's base is offset in the ptr table. Down unchanged.
3482            let (rbg_d, rbu_d) = if dev.gu_il {
3483                let sxx = m.gate_exps.row_bytes + m.up_exps.row_bytes; (sxx, sxx)
3484            } else { (m.gate_exps.row_bytes, m.up_exps.row_bytes) };
3485            let q8 = moe_q8_enabled()
3486                && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3487                && q8_expert_supported(m.down_exps.qtype);
3488            // SMALL-M ROWS ARM (MEMRA_SPEC_M2, lane/spec-m2): batch the verify token loop —
3489            // ONE batched z-quantize + ONE gate_up rows launch + ONE act quantize + ONE down
3490            // rows launch (4 launches/layer, was 4t). BIT-IDENTICAL per token to the serial
3491            // loop below (rows twins = the _v/w8h2v per-token programs on a grid.z token axis;
3492            // quantize_q8_1 is per-32-block row-independent). Gated to the AUTO kernel modes —
3493            // a custom MEMRA_MOE_DEVQ8_GU/DOWN diagnostic run keeps the serial loop so the
3494            // dispatched kernel stays exactly the env-selected one — and to the w8h2v shape
3495            // (n_ff_exp==512, n_used<=8), the same contract the AUTO down dispatch keys on.
3496            let rows_arm = q8 && t > 1 && crate::spec::spec_m2()
3497                && n_ff_exp == 512 && n_used <= 8
3498                && std::env::var("MEMRA_MOE_DEVQ8_GU").map(|v| v.is_empty() || v == "v").unwrap_or(true)
3499                && std::env::var("MEMRA_MOE_DEVQ8_DOWN").map(|v| v.is_empty() || v == "w8h2v").unwrap_or(true);
3500            // CSR EXPERT-DEDUP gate_up (verify-cost target #1, DEFAULT ON 2026-07-10): the
3501            // owner-scan kernel serves every (token, slot) pair of an expert from ONE block,
3502            // deduping the 38-40% duplicated weight-stream+decode the overlap probe measured
3503            // (gate_up 55.0 -> 39.7us/launch; +1.3-2.1% spec e2e p2, +0.6-1.7% p3, all K).
3504            // Bit-identical to the _rows twins (explicit-intrinsic accumulate — the ULP/fmad
3505            // lesson; =2 byte-compare verified zero diffs). down stays on _rows (CSR down
3506            // measured 23.5 -> 37.5us: 16-group rows can't amortize the serial pair loop).
3507            // MEMRA_MOE_CSR=0 rollback; =2 runs BOTH paths and byte-compares (debug).
3508            let csr_mode = std::env::var("MEMRA_MOE_CSR").ok()
3509                .and_then(|v| v.parse::<i32>().ok()).unwrap_or(1);
3510            let csr_qt = |qt: i32| qt == crate::QT_IQ4_XS || qt == crate::QT_IQ3_S;
3511            let csr_arm = rows_arm && csr_mode > 0 && t <= 10
3512                && csr_qt(m.gate_exps.qtype) && csr_qt(m.up_exps.qtype)
3513                && csr_qt(m.down_exps.qtype);
3514            if csr_arm {
3515                if csr_mode == 2 {
3516                    static ENGAGED: std::sync::Once = std::sync::Once::new();
3517                    ENGAGED.call_once(|| eprintln!("[memra] moe CSR byte-compare mode ON (t={t})"));
3518                }
3519                let n_pairs = t * n_used;
3520                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3521                let act = e.moe_gate_up_silu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, n_pairs,
3522                                                         n_embd, n_ff_exp, n_used, n_expert,
3523                                                         m.gate_exps.qtype, m.up_exps.qtype,
3524                                                         rbg_d, rbu_d)?;
3525                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
3526                // down stays on the _rows twin — BOTH CSR down variants measured negative
3527                // (v1 serial pairs 23.5->37.5us; v2 warp-parallel+SMEM -14% e2e, K=8 -37%):
3528                // 16-group rows have too little decode to amortize any dedup structure.
3529                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
3530                                            t, n_ff_exp, n_embd, n_used, n_expert,
3531                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
3532                if csr_mode == 2 {
3533                    // DEBUG BYTE-COMPARE: run the _rows twins on the same inputs, diff bits.
3534                    let act_r = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
3535                                                                n_embd, n_ff_exp, n_used, n_expert,
3536                                                                m.gate_exps.qtype, m.up_exps.qtype,
3537                                                                rbg_d, rbu_d, &m.dev_macros)?;
3538                    let mut out_r = e.uninit(t * n_embd)?;
3539                    let (aq2r, ad2r) = e.quantize_q8_1(&act_r, n_pairs, n_ff_exp)?;
3540                    e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2r, &ad2r, &mut out_r,
3541                                                t, n_ff_exp, n_embd, n_used, n_expert,
3542                                                m.down_exps.qtype, m.down_exps.row_bytes)?;
3543                    let (a1, a2) = (e.dtoh(&act)?, e.dtoh(&act_r)?);
3544                    let (o1, o2) = (e.dtoh(&moe_out)?, e.dtoh(&out_r)?);
3545                    let ba = a1.iter().zip(&a2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3546                    let bo = o1.iter().zip(&o2).filter(|(x, y)| x.to_bits() != y.to_bits()).count();
3547                    if ba + bo > 0 {
3548                        eprintln!("[csr-check] il={il} t={t} ACT diffs={ba}/{} OUT diffs={bo}/{}",
3549                                  a1.len(), o1.len());
3550                        // First 4 differing ACT elements: (pair, o, csr, rows) + that pair's expert
3551                        let sel_h = e.dtoh_i32(&sel_d)?;
3552                        let mut shown = 0;
3553                        for (i, (x, y)) in a1.iter().zip(&a2).enumerate() {
3554                            if x.to_bits() != y.to_bits() && shown < 4 {
3555                                let (p, o) = (i / n_ff_exp, i % n_ff_exp);
3556                                let ex = sel_h[p];
3557                                let npx = sel_h.iter().filter(|&&v| v == ex).count();
3558                                eprintln!("  ACT p={p} ex={ex} np={npx} o={o} csr={x:e} rows={y:e}");
3559                                shown += 1;
3560                            }
3561                        }
3562                        std::process::exit(3);
3563                    }
3564                }
3565            } else if rows_arm {
3566                // TEMP PROBE (MEMRA_MOE_OVERLAP=1): cross-token expert-activation overlap at
3567                // verify — sizes the CSR dedup win (unique experts vs t*n_used pairs).
3568                if std::env::var("MEMRA_MOE_OVERLAP").as_deref() == Ok("1") {
3569                    use std::sync::atomic::{AtomicU64, Ordering};
3570                    static PAIRS: AtomicU64 = AtomicU64::new(0);
3571                    static UNIQ: AtomicU64 = AtomicU64::new(0);
3572                    static CALLS: AtomicU64 = AtomicU64::new(0);
3573                    let sel_h = e.dtoh_i32(&sel_d)?;
3574                    let mut u: Vec<i32> = sel_h.clone(); u.sort_unstable(); u.dedup();
3575                    PAIRS.fetch_add(sel_h.len() as u64, Ordering::Relaxed);
3576                    UNIQ.fetch_add(u.len() as u64, Ordering::Relaxed);
3577                    let c = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
3578                    if c % 480 == 0 {
3579                        let p = PAIRS.load(Ordering::Relaxed); let q = UNIQ.load(Ordering::Relaxed);
3580                        eprintln!("[overlap] calls={c} pairs={p} unique={q} ratio={:.3} (t={t})",
3581                                  q as f64 / p as f64);
3582                    }
3583                }
3584                let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3585                let act = e.moe_gate_up_silu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
3586                                                          n_embd, n_ff_exp, n_used, n_expert,
3587                                                          m.gate_exps.qtype, m.up_exps.qtype,
3588                                                          rbg_d, rbu_d, &m.dev_macros)?;
3589                let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
3590                e.moe_down8_fma_dev_q8_rows(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out,
3591                                            t, n_ff_exp, n_embd, n_used, n_expert,
3592                                            m.down_exps.qtype, m.down_exps.row_bytes)?;
3593            } else {
3594            for tok in 0..t {
3595                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
3596                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
3597                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
3598                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3599                if q8 {
3600                    let (zq, zd) = match (t, zq8) {
3601                        (1, Some((q, d))) => (q.clone(), d.clone()),
3602                        _ => e.quantize_q8_1_view(&zt, 1, n_embd)?,
3603                    };
3604                    let act = e.moe_gate_up_silu8_dev_q8(&dev.ptr_row, &selt, &zq, &zd,
3605                                                         n_embd, n_ff_exp, n_used, n_expert,
3606                                                         m.gate_exps.qtype, m.up_exps.qtype,
3607                                                         rbg_d, rbu_d, &m.dev_macros)?;
3608                    let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
3609                    e.moe_down8_fma_dev_q8(&dev.ptr_row, &selt, &wt, &aq2, &ad2, &mut dst,
3610                                           n_ff_exp, n_embd, n_used, n_expert,
3611                                           m.down_exps.qtype, m.down_exps.row_bytes)?;
3612                } else {
3613                    let act = e.moe_gate_up_silu8_dev(&dev.ptr_row, &selt, &zt, n_embd, n_ff_exp,
3614                                                      n_used, n_expert,
3615                                                      m.gate_exps.qtype, m.up_exps.qtype,
3616                                                      rbg_d, rbu_d, &m.dev_macros)?;
3617                    e.moe_down8_fma_dev(&dev.ptr_row, &selt, &wt, &act, &mut dst,
3618                                        n_ff_exp, n_embd, n_used, n_expert,
3619                                        m.down_exps.qtype, m.down_exps.row_bytes)?;
3620                }
3621            }
3622            }
3623        } else {
3624        // Launch under the cache lock: the row borrow lives as long as the closure, and the
3625        // lock covers only launch ISSUE (µs), same policy as moe_cached_gemm.
3626        // Q8 ARM PARITY (2026-07-06): the SLRU arm ran the f32-dequant _dev kernels only —
3627        // 80us/launch vs the q8 twins' 15us on the SAME shapes (fixed-build profile: 228
3628        // f32 launches = 36ms of the 64-tok window). Same q8 gate + kernels as the resident
3629        // arm above; MEMRA_MOE_Q8=0 restores the byte-identical f32 path.
3630        let q8 = moe_q8_enabled()
3631            && q8_expert_supported(m.gate_exps.qtype) && q8_expert_supported(m.up_exps.qtype)
3632            && q8_expert_supported(m.down_exps.qtype);
3633        e.with_moe_cache(max_block, |c, eng| {
3634            let row = c.layer_dev_row(il, n_expert, eng)?
3635                .ok_or("moe_ffn_dev: layer row vanished under the lock")?;
3636            for tok in 0..t {
3637                let zt = z.slice(tok * n_embd..(tok + 1) * n_embd);
3638                let selt = sel_d.slice(tok * n_used..(tok + 1) * n_used);
3639                let wt = w_d.slice(tok * n_used..(tok + 1) * n_used);
3640                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3641                if q8 {
3642                    let (zq, zd) = match (t, zq8) {
3643                        (1, Some((q, d))) => (q.clone(), d.clone()),
3644                        _ => eng.quantize_q8_1_view(&zt, 1, n_embd)?,
3645                    };
3646                    let act = eng.moe_gate_up_silu8_dev_q8(row, &selt, &zq, &zd,
3647                                                           n_embd, n_ff_exp, n_used, n_expert,
3648                                                           m.gate_exps.qtype, m.up_exps.qtype,
3649                                                           m.gate_exps.row_bytes, m.up_exps.row_bytes,
3650                                                           &m.dev_macros)?;
3651                    let (aq2, ad2) = eng.quantize_q8_1(&act, n_used, n_ff_exp)?;
3652                    eng.moe_down8_fma_dev_q8(row, &selt, &wt, &aq2, &ad2, &mut dst,
3653                                             n_ff_exp, n_embd, n_used, n_expert,
3654                                             m.down_exps.qtype, m.down_exps.row_bytes)?;
3655                } else {
3656                    let act = eng.moe_gate_up_silu8_dev(row, &selt, &zt, n_embd, n_ff_exp,
3657                                                        n_used, n_expert,
3658                                                        m.gate_exps.qtype, m.up_exps.qtype,
3659                                                        m.gate_exps.row_bytes, m.up_exps.row_bytes,
3660                                                        &m.dev_macros)?;
3661                    eng.moe_down8_fma_dev(row, &selt, &wt, &act, &mut dst,
3662                                          n_ff_exp, n_embd, n_used, n_expert,
3663                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
3664                }
3665            }
3666            // instrumentation parity with the host paths (3 blocks/expert-slot, all hits).
3667            c.hits += (t * 3 * n_used) as u64;
3668            Ok(())
3669        })?;
3670        }
3671
3672        // SHARED EXPERT epilogue — byte-identical to moe_ffn_sequential step 3 (incl. its Q8
3673        // TRUNK-FUSION arm: fused2 is bit-identical to the two matmul calls per (tensor,row)).
3674        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
3675        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
3676        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
3677            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
3678        {
3679            let n_ff_sh = gate_shexp.out_features();
3680            // verify-t (2..15) decode-exact arm: this fn now serves the spec verify batches
3681            // (pairs gate moved to t>=PRIME_MIN_T), so the shexp chain must match t==1 per col.
3682            let verify_t = t > 1 && t < PRIME_MIN_T;
3683            let (sg_gate, sg_up) = if t == 1 {
3684                match e.matmul_q8_fused2_x(gate_shexp, up_shexp, z)? {
3685                    Some(pair) => pair,
3686                    None => (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?),
3687                }
3688            } else if verify_t {
3689                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): the shexp gate+up pair
3690                // rides one shared quantize + one fused2 batched launch instead of two
3691                // decode-exact calls. Bit-identical per (tensor,token,row) — see spec_fused_t.
3692                let mut fused = None;
3693                if crate::spec::spec_fused_t() && (2..=4).contains(&t)
3694                    && e.uses_q8_1_fast(gate_shexp) && e.uses_q8_1_fast(up_shexp) {
3695                    let (zq, zd) = e.quantize_q8_1(z, t, n_embd)?;
3696                    fused = e.matmul_q8_fused2_t(gate_shexp, up_shexp, &zq, &zd, t)?;
3697                }
3698                match fused {
3699                    Some(pair) => pair,
3700                    None => (e.matmul_decode_exact(gate_shexp, z, t)?,
3701                             e.matmul_decode_exact(up_shexp, z, t)?),
3702                }
3703            } else {
3704                (e.matmul(gate_shexp, z, t)?, e.matmul(up_shexp, z, t)?)
3705            };
3706            let mut sa = e.uninit(t * n_ff_sh)?;  // silu_mul fully overwrites
3707            e.silu_mul(&sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
3708            let sh = if verify_t { e.matmul_decode_exact(down_shexp, &sa, t)? }
3709                     else { e.matmul(down_shexp, &sa, t)? };
3710            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
3711            // Same fused sigmoid-dot as moe_ffn_sequential step 3 (byte-identity contract
3712            // between the two arms; prefill keeps the batched cuBLASLt linear).
3713            let g = match &m.gate_inp_shexp {
3714                Some(gate_inp_shexp) => {
3715                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
3716                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
3717                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
3718                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
3719                    } else {
3720                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
3721                        let mut g = e.uninit(t)?;
3722                        e.sigmoid(&gs, &mut g, t)?;
3723                        g
3724                    }
3725                }
3726                None => e.htod(&vec![1.0f32; t])?,
3727            };
3728            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
3729        }
3730
3731        Ok(moe_out)
3732    }
3733
3734    /// STAGE-2 GROUPED DECODE (2026-07-04): run ONE token's whole routed-expert FFN in TWO
3735    /// launches when every one of its 3*n_used blocks is ALREADY cache-resident. Returns
3736    /// Ok(true) if the grouped path ran (caller skips the sequential loop for this token);
3737    /// Ok(false) on ANY miss (caller falls through — the sequential loop stages/admits as
3738    /// before, so the NEXT occurrence takes the grouped path). Pointer safety: cache slots are
3739    /// fixed-address for the engine's lifetime and the pure-HIT path performs no admission, so
3740    /// the collected raw pointers cannot move between collection and launch (single-threaded
3741    /// decode; the lock is held only for collection, launches are stream-ordered after any
3742    /// prior same-stream staging writes).
3743    #[allow(clippy::too_many_arguments)]
3744    /// q8 twin of moe_gdec_token (dp4a arc): same residency check + 2-launch shape; the fused
3745    /// kernels consume the pre-quantized z-row and re-quantize act per slot batch.
3746    #[allow(clippy::too_many_arguments)]
3747    fn moe_gdec_token_q8(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
3748                      zq: &CudaSlice<i8>, zd: &CudaSlice<f32>, sel: &[u32], w: &[f32],
3749                      moe_out: &mut CudaSlice<f32>, tok: usize,
3750                      n_embd: usize, n_ff_exp: usize, n_used: usize)
3751                      -> Result<bool, Box<dyn std::error::Error>> {
3752        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
3753        use cudarc::driver::DevicePtr;
3754        let ptrs = e.with_moe_cache(max_block, |c, eng| {
3755            let mut g = [0u64; 8];
3756            let mut u = [0u64; 8];
3757            let mut d = [0u64; 8];
3758            for (j, &ex) in sel.iter().enumerate() {
3759                let ex = ex as u16;
3760                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
3761                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
3762                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
3763                else { return Ok(None); };
3764                let __s = eng.stream();
3765                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
3766                let (pu, _e1) = c.slot(su).device_ptr(&__s);
3767                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
3768                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
3769            }
3770            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
3771                for &ex in sel {
3772                    let ex = ex as u16;
3773                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
3774                        c.note_profile_hit(BlockId::new(il, proj, ex));
3775                    }
3776                }
3777            }
3778            c.hits += (3 * n_used) as u64;
3779            Ok(Some((g, u, d)))
3780        })?;
3781        let Some((g, u, d)) = ptrs else { return Ok(false) };
3782        let mut wv = [0f32; 8];
3783        wv[..n_used].copy_from_slice(w);
3784        let act = e.moe_gate_up_silu8_q8(crate::WPtr8(g), crate::WPtr8(u), zq, zd,
3785                                         n_embd, n_ff_exp, n_used,
3786                                         m.gate_exps.qtype, m.up_exps.qtype,
3787                                         m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
3788        // per-slot act quantize: [n_used, n_ff] rows in one quantize launch.
3789        let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
3790        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3791        e.moe_down8_fma_q8(crate::WPtr8(d), crate::F32x8(wv), &aq2, &ad2, &mut dst,
3792                           n_ff_exp, n_embd, n_used,
3793                           m.down_exps.qtype, m.down_exps.row_bytes)?;
3794        Ok(true)
3795    }
3796
3797    fn moe_gdec_token(e: &Engine, m: &MoeWeights, il: u16, max_block: usize,
3798                      zt: &cudarc::driver::CudaView<f32>, sel: &[u32], w: &[f32],
3799                      moe_out: &mut CudaSlice<f32>, tok: usize,
3800                      n_embd: usize, n_ff_exp: usize, n_used: usize)
3801                      -> Result<bool, Box<dyn std::error::Error>> {
3802        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
3803        use cudarc::driver::DevicePtr;
3804        // One lock hold: residency-check all 3*n_used blocks, collect raw slot pointers.
3805        let ptrs = e.with_moe_cache(max_block, |c, eng| {
3806            let mut g = [0u64; 8];
3807            let mut u = [0u64; 8];
3808            let mut d = [0u64; 8];
3809            for (j, &ex) in sel.iter().enumerate() {
3810                let ex = ex as u16;
3811                let (Some(sg), Some(su), Some(sd)) = (c.resident(BlockId::new(il, PROJ_GATE, ex)),
3812                                                      c.resident(BlockId::new(il, PROJ_UP,   ex)),
3813                                                      c.resident(BlockId::new(il, PROJ_DOWN, ex)))
3814                else { return Ok(None); };
3815                let __s = eng.stream();
3816                let (pg, _e0) = c.slot(sg).device_ptr(&__s);
3817                let (pu, _e1) = c.slot(su).device_ptr(&__s);
3818                let (pd, _e2) = c.slot(sd).device_ptr(&__s);
3819                g[j] = pg as u64; u[j] = pu as u64; d[j] = pd as u64;
3820            }
3821            if cpu_expert_profile_admit_enabled() && !c.is_frozen() {
3822                for &ex in sel {
3823                    let ex = ex as u16;
3824                    for proj in [PROJ_GATE, PROJ_UP, PROJ_DOWN] {
3825                        c.note_profile_hit(BlockId::new(il, proj, ex));
3826                    }
3827                }
3828            }
3829            c.hits += (3 * n_used) as u64; // instrumentation parity with dispatch()
3830            Ok(Some((g, u, d)))
3831        })?;
3832        let Some((g, u, d)) = ptrs else { return Ok(false) };
3833        let mut wv = [0f32; 8];
3834        wv[..n_used].copy_from_slice(w);
3835        // 2 launches: (gate+up+silu) x8, then (down + slot-ordered FMA accumulate) x8.
3836        let act = e.moe_gate_up_silu8(crate::WPtr8(g), crate::WPtr8(u), zt,
3837                                      n_embd, n_ff_exp, n_used,
3838                                      m.gate_exps.qtype, m.up_exps.qtype,
3839                                      m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
3840        let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
3841        e.moe_down8_fma_into(crate::WPtr8(d), crate::F32x8(wv), &act, &mut dst,
3842                             n_ff_exp, n_embd, n_used,
3843                             m.down_exps.qtype, m.down_exps.row_bytes)?;
3844        Ok(true)
3845    }
3846
3847    /// EDGE-1 §B.3: dispatch one expert projection through the SLRU cache, then run the SAME
3848    /// `qmatvec_view` from whichever slot it landed in (resident HIT or staged MISS). `x` is the
3849    /// sliced activation row. `proj` selects the gate/up/down HostExps tensor. Returns y = W_expert @ x.
3850    /// q8 twin of moe_cached_gemm: same dispatch/slot mechanics, dp4a expert kernel.
3851    fn moe_cached_gemm_q8(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
3852                          max_block: usize, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
3853                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3854        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
3855        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
3856        let layout = exps.expert_layout(ex);
3857        let id = BlockId::new(il, proj, ex as u16);
3858        let source = exps.expert_source(ex);
3859        e.with_moe_cache(max_block, |c, eng| {
3860            let slot = c.dispatch_source(id, source, eng)?;
3861            let DispatchSlot::Resident(sl) = slot;
3862            let buf = c.slot(sl);
3863            eng.qmatvec_expert_q8(buf, 0..layout.len, aq, ad, 1, exps.in_f, exps.out_f,
3864                                  layout.qtype, layout.row_bytes)
3865        })
3866    }
3867
3868    fn moe_cached_gemm(e: &Engine, il: u16, proj: u8, ex: usize, m: &MoeWeights,
3869                       max_block: usize, x: &cudarc::driver::CudaView<f32>)
3870                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3871        use crate::moe_cache::{BlockId, DispatchSlot, PROJ_GATE, PROJ_UP};
3872        let exps = match proj { PROJ_GATE => &m.gate_exps, PROJ_UP => &m.up_exps, _ => &m.down_exps };
3873        let layout = exps.expert_layout(ex);
3874        let id = BlockId::new(il, proj, ex as u16);
3875        let source = exps.expert_source(ex);
3876        // dispatch under the lock (lookup/admit/memcpy-issue), then resolve the slot and GEMM.
3877        e.with_moe_cache(max_block, |c, eng| {
3878            let slot = c.dispatch_source(id, source, eng)?;
3879            // resolve the device buffer for this slot; the GEMM is enqueued on the compute stream
3880            // (the same stream the memcpy was issued on, so ordering holds without extra sync).
3881            let DispatchSlot::Resident(sl) = slot;
3882            let buf = c.slot(sl);
3883            eng.qmatvec_view(buf, 0..layout.len, x, 1, exps.in_f, exps.out_f,
3884                             layout.qtype, layout.row_bytes)
3885        })
3886    }
3887
3888    /// Populate the warmup cache for an expert whose current-token output intentionally ran on
3889    /// CPU. No GEMM is launched and callers invoke this only after the CPU result has completed,
3890    /// so the current forward's backend assignment and output remain unchanged.
3891    fn moe_profile_admit_expert(
3892        e: &Engine,
3893        il: u16,
3894        ex: usize,
3895        m: &MoeWeights,
3896        max_block: usize,
3897    ) -> Result<(), Box<dyn std::error::Error>> {
3898        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3899        e.with_moe_cache(max_block, |cache, eng| {
3900            for (proj, exps) in [
3901                (PROJ_GATE, &m.gate_exps),
3902                (PROJ_UP, &m.up_exps),
3903                (PROJ_DOWN, &m.down_exps),
3904            ] {
3905                let id = BlockId::new(il, proj, ex as u16);
3906                let _ = cache.dispatch_source(id, exps.expert_source(ex), eng)?;
3907            }
3908            Ok(())
3909        })
3910    }
3911
3912    /// Read a projection from the immutable residency set when present; otherwise use one
3913    /// transient slot without admitting or evicting anything. Used only by post-freeze prefill.
3914    #[allow(clippy::too_many_arguments)]
3915    fn moe_frozen_gemm(
3916        e: &Engine,
3917        il: u16,
3918        proj: u8,
3919        ex: usize,
3920        m: &MoeWeights,
3921        max_block: usize,
3922        x: &cudarc::driver::CudaView<f32>,
3923        scratch: &mut Option<CudaSlice<u8>>,
3924        scratch_len: usize,
3925    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3926        use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP};
3927        let exps = match proj {
3928            PROJ_GATE => &m.gate_exps,
3929            PROJ_UP => &m.up_exps,
3930            _ => &m.down_exps,
3931        };
3932        let layout = exps.expert_layout(ex);
3933        let id = BlockId::new(il, proj, ex as u16);
3934        if let Some(output) = e.with_moe_cache(max_block, |cache, eng| {
3935            let Some(slot) = cache.resident(id) else {
3936                return Ok(None);
3937            };
3938            let buf = cache.slot(slot);
3939            Ok(Some(eng.qmatvec_view(
3940                buf,
3941                0..layout.len,
3942                x,
3943                1,
3944                exps.in_f,
3945                exps.out_f,
3946                layout.qtype,
3947                layout.row_bytes,
3948            )?))
3949        })? {
3950            return Ok(output);
3951        }
3952        if scratch.is_none() {
3953            *scratch = Some(e.alloc_u8_uninit(scratch_len)?);
3954        }
3955        let scratch = scratch.as_mut().unwrap();
3956        e.stage_expert(exps.expert_bytes(ex), scratch, 0)?;
3957        e.qmatvec_view(
3958            scratch,
3959            0..layout.len,
3960            x,
3961            1,
3962            exps.in_f,
3963            exps.out_f,
3964            layout.qtype,
3965            layout.row_bytes,
3966        )
3967    }
3968
3969    fn moe_prefetch_expert(
3970        e: &Engine,
3971        il: u16,
3972        ex: usize,
3973        m: &MoeWeights,
3974        max_block: usize,
3975        keep: &[crate::moe_cache::BlockId],
3976    ) -> Result<(), Box<dyn std::error::Error>> {
3977        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3978        e.with_moe_cache(max_block, |c, eng| {
3979            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
3980                                 (PROJ_DOWN, &m.down_exps)] {
3981                let id = BlockId::new(il, proj, ex as u16);
3982                let _ = c.prefetch_source(id, exps.expert_source(ex), keep, eng)?;
3983            }
3984            Ok(())
3985        })
3986    }
3987
3988    /// Worker-mode disk lookahead for grouped prefill. Memory sources are deliberately skipped so
3989    /// selecting `worker` changes only storage scheduling here; all CUDA work stays in dispatch.
3990    fn moe_prefetch_disk_expert(e: &Engine, il: u16, ex: usize, m: &MoeWeights,
3991                                max_block: usize, keep: &[crate::moe_cache::BlockId])
3992                                -> Result<(), Box<dyn std::error::Error>> {
3993        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
3994        e.with_moe_cache(max_block, |c, eng| {
3995            for (proj, exps) in [(PROJ_GATE, &m.gate_exps), (PROJ_UP, &m.up_exps),
3996                                 (PROJ_DOWN, &m.down_exps)] {
3997                let source = exps.expert_source(ex);
3998                if let crate::model::ExpertSource::Disk { .. } = &source {
3999                    let id = BlockId::new(il, proj, ex as u16);
4000                    let _ = c.prefetch_source(id, source, keep, eng)?;
4001                }
4002            }
4003            Ok(())
4004        })
4005    }
4006
4007    #[inline]
4008    fn moe_prefetch_host_expert(ex: usize, m: &MoeWeights) {
4009        let _ = m.gate_exps.prefetch_expert_pages(ex);
4010        let _ = m.up_exps.prefetch_expert_pages(ex);
4011        let _ = m.down_exps.prefetch_expert_pages(ex);
4012    }
4013}
4014
4015// ================================================================================================
4016// A2: EXPERT-GROUPED MoE PREFILL (MEMRA_MOE_GROUPED=1). Resident-case prototype.
4017//
4018// Instead of the per-token loop (T * 8 experts * 3 projections = 12024 individual m=1 matvecs),
4019// this groups tokens by expert and runs ONE matmul per active expert per projection at m=m_e.
4020// On a 501-token prefill with ~170 active experts, that's ~510 matmuls (vs 12024).
4021//
4022// EXACTNESS: per-token accumulation across its 8 experts is reordered (grouped processes experts
4023// in expert-id order, not the router's top-k order). To preserve bit-identity with the sequential
4024// loop, we use an 8-SLOT scheme: expert outputs are scattered into slots keyed by the token's
4025// top-k position (0..7), then reduced in that fixed order. This makes the f32 addition order
4026// identical to the per-token loop regardless of expert processing order.
4027//
4028// Memory: T * 8 * n_embd * 4 = 501 * 8 * 2048 * 4 = ~32 MB (slot buffer). Fine on 96GB.
4029// ================================================================================================
4030
4031impl HybridModel {
4032    /// A2 expert-grouped MoE FFN (prefill path, MEMRA_MOE_GROUPED=1). Same semantics as moe_ffn:
4033    /// z [T, n_embd] -> moe_out [T, n_embd]. BIT-IDENTICAL to moe_ffn when using the slot scheme.
4034    pub(crate) fn moe_ffn_grouped(e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize,
4035                                  cfg: &ModelConfig, il: u16, _max_block: usize)
4036                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4037        let moe = cfg.moe.as_ref().unwrap();
4038        let n_embd = cfg.n_embd as usize;
4039        let n_expert = moe.expert_count as usize;
4040        let n_used = moe.expert_used_count as usize;
4041        let n_ff_exp = moe.expert_ff_length as usize;
4042
4043        // 1. ROUTER (identical to moe_ffn).
4044        let logits = e.matmul(&m.gate_inp, z, t)?;
4045        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4046            Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4047                                m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4048        } else {
4049            Self::moe_route_cfg(e, &logits, t, n_expert, n_used,
4050                                None, None, m.active_experts.as_deref())?
4051        };
4052        Self::trace_moe_routes(il, t, &sel_all, &w_all)?;
4053
4054        // 2. BUILD PER-EXPERT TOKEN LISTS (host-side grouping).
4055        // For each expert e, we need: which tokens use it, their positions in z, their top-k
4056        // slot index (for bit-identical accumulation), and their weights.
4057        struct ExpertGroup {
4058            tok_indices: Vec<i32>,   // indices into z rows (0..T-1)
4059            slot_indices: Vec<i32>,  // top-k slot (0..n_used-1) for that token-expert pair
4060            weights: Vec<f32>,       // renormalized weight for that token-expert pair
4061        }
4062        let mut groups: Vec<ExpertGroup> = (0..n_expert).map(|_| ExpertGroup {
4063            tok_indices: Vec::new(), slot_indices: Vec::new(), weights: Vec::new(),
4064        }).collect();
4065
4066        for tok in 0..t {
4067            for j in 0..n_used {
4068                let ex = sel_all[tok * n_used + j] as usize;
4069                let w = w_all[tok * n_used + j];
4070                groups[ex].tok_indices.push(tok as i32);
4071                groups[ex].slot_indices.push(j as i32);
4072                groups[ex].weights.push(w);
4073            }
4074        }
4075
4076        // 3. ALLOCATE SLOT BUFFER: [T, n_used, n_embd] f32, zero-initialized.
4077        // Each token's 8 expert contributions land in their respective slots.
4078        let mut slot_buf = e.zeros(t * n_used * n_embd)?;
4079        let mut wbuf = e.zeros(t * n_used)?;  // [T, n_used] weight buffer for FMA reduce
4080
4081        // Expert weight dimensions (used in both cache and staging paths).
4082        let g_len = m.gate_exps.max_expert_bytes();
4083        let u_len = m.up_exps.max_expert_bytes();
4084        let d_len = m.down_exps.max_expert_bytes();
4085        let use_cache = Engine::moe_cache_enabled();
4086        let max_block = _max_block;
4087
4088        // GPU scratch for staging (only allocated when NOT using cache).
4089        let (mut scratch_g, mut scratch_u, mut scratch_d) = if !use_cache {
4090            (Some(e.alloc_u8(g_len)?), Some(e.alloc_u8(u_len)?), Some(e.alloc_u8(d_len)?))
4091        } else {
4092            (None, None, None)
4093        };
4094
4095        // 4. PER ACTIVE EXPERT: gather, compute, scatter.
4096        // Processing ORDER: DESCENDING m_e (biggest token batches first) — the concluded winner
4097        // (rig5090 2026-07-04, the ascending-id arm and its MEMRA_MOE_ORDER seam removed): desc is
4098        // a first-forward win at partial cache capacity — the hot (big-m_e) experts are admitted
4099        // to the SLRU before the small-m_e tail can pollute it, so residency converges in ONE
4100        // forward instead of several: auto-cache T=501 126.9 -> 169.9 tok/s (1.34x), cap512
4101        // 119.6 -> 160.8 (and kills the rep-to-rep bimodal); wash (<2%) at cap64 pure-spill and
4102        // at long prompts where every expert stages regardless. Order is FREE to change without
4103        // breaking the byte-identity gate: the slot scheme pins each token's accumulation order
4104        // regardless of expert processing order (the whole point of the slots).
4105        let mut order: Vec<usize> =
4106            (0..n_expert).filter(|&ex| !groups[ex].tok_indices.is_empty()).collect();
4107        order.sort_by(|&a, &b| groups[b].tok_indices.len()
4108            .cmp(&groups[a].tok_indices.len()).then(a.cmp(&b)));
4109        let mut m_dist: Vec<usize> = Vec::new();  // for stats
4110        let page_window = moe_page_prefetch_window();
4111        let worker_disk_prefetch = use_cache && crate::spill_pread::worker_enabled();
4112        if worker_disk_prefetch {
4113            if let Some(first) = grouped_worker_prefetch_position(order.len(), None) {
4114                Self::moe_prefetch_disk_expert(e, il, order[first], m, max_block, &[])?;
4115            }
4116        }
4117        for (order_pos, &ex) in order.iter().enumerate() {
4118            for next in page_prefetch_positions(order_pos, order.len(), page_window) {
4119                Self::moe_prefetch_host_expert(order[next], m);
4120            }
4121            if worker_disk_prefetch {
4122                if let Some(next) = grouped_worker_prefetch_position(order.len(), Some(order_pos)) {
4123                    use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4124                    let keep = [
4125                        BlockId::new(il, PROJ_GATE, ex as u16),
4126                        BlockId::new(il, PROJ_UP, ex as u16),
4127                        BlockId::new(il, PROJ_DOWN, ex as u16),
4128                    ];
4129                    Self::moe_prefetch_disk_expert(e, il, order[next], m, max_block, &keep)?;
4130                }
4131            }
4132            let grp = &groups[ex];
4133            let m_e = grp.tok_indices.len();
4134            m_dist.push(m_e);
4135            let gl = m.gate_exps.expert_layout(ex);
4136            let ul = m.up_exps.expert_layout(ex);
4137            let dl = m.down_exps.expert_layout(ex);
4138
4139            // Upload index/weight arrays to device. The down-proj per-expert macro-scale
4140            // (ModelOpt weight_scale_2) folds into the scatter weights — post-matmul linear,
4141            // same fold as the sequential loop's `w[j] * macro_scale(ex)`. 1.0 for GGUF experts.
4142            let tok_idx_d = e.htod_i32(&grp.tok_indices)?;
4143            let slot_idx_d = e.htod_i32(&grp.slot_indices)?;
4144            let dmac = m.down_exps.macro_scale(ex);
4145            let weight_d = if dmac == 1.0 { e.htod(&grp.weights)? } else {
4146                let scaled: Vec<f32> = grp.weights.iter().map(|&w| w * dmac).collect();
4147                e.htod(&scaled)?
4148            };
4149
4150            // GATHER: collect m_e activation rows from z into a contiguous buffer.
4151            let mut gathered = e.zeros(m_e * n_embd)?;
4152            e.gather_rows(z, &tok_idx_d, &mut gathered, n_embd, m_e)?;
4153            let gv = gathered.slice(0..m_e * n_embd);
4154
4155            // Compute gate/up/down matmuls -- two paths: cache-resident or host-staged.
4156            let y = if use_cache {
4157                use crate::moe_cache::{BlockId, PROJ_GATE, PROJ_UP, PROJ_DOWN};
4158                // CACHE PATH: dispatch through MOE cache, get device-resident buffer, GEMM at m=m_e.
4159                let gate = e.with_moe_cache(max_block, |c, eng| {
4160                    let id = BlockId::new(il, PROJ_GATE, ex as u16);
4161                    let slot = c.dispatch_source(id, m.gate_exps.expert_source(ex), eng)?;
4162                    let buf = c.buf(slot);
4163                    eng.qmatvec_view(buf, 0..gl.len, &gv, m_e,
4164                        m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4165                })?;
4166                let up = e.with_moe_cache(max_block, |c, eng| {
4167                    let id = BlockId::new(il, PROJ_UP, ex as u16);
4168                    let slot = c.dispatch_source(id, m.up_exps.expert_source(ex), eng)?;
4169                    let buf = c.buf(slot);
4170                    eng.qmatvec_view(buf, 0..ul.len, &gv, m_e,
4171                        m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4172                })?;
4173                // SiLU-MUL activation (per-expert macro-scales folded).
4174                let mut act = e.zeros(m_e * n_ff_exp)?;
4175                Self::ffn_act_scaled(e, cfg, &gate, &up,
4176                    m.gate_exps.macro_scale(ex), m.up_exps.macro_scale(ex), &mut act, m_e * n_ff_exp)?;
4177                let actv = act.slice(0..m_e * n_ff_exp);
4178                e.with_moe_cache(max_block, |c, eng| {
4179                    let id = BlockId::new(il, PROJ_DOWN, ex as u16);
4180                    let slot = c.dispatch_source(id, m.down_exps.expert_source(ex), eng)?;
4181                    let buf = c.buf(slot);
4182                    eng.qmatvec_view(buf, 0..dl.len, &actv, m_e,
4183                        m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4184                })?
4185            } else {
4186                // STAGING PATH: H2D the expert blocks into scratch buffers, then GEMM.
4187                let sg = scratch_g.as_mut().unwrap();
4188                let su = scratch_u.as_mut().unwrap();
4189                let sd = scratch_d.as_mut().unwrap();
4190                e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4191                e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4192                e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4193                let gate = e.qmatvec_view(sg, 0..gl.len, &gv, m_e,
4194                    m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)?;
4195                let up = e.qmatvec_view(su, 0..ul.len, &gv, m_e,
4196                    m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)?;
4197                // SiLU-MUL activation (per-expert macro-scales folded).
4198                let mut act = e.zeros(m_e * n_ff_exp)?;
4199                Self::ffn_act_scaled(e, cfg, &gate, &up,
4200                    m.gate_exps.macro_scale(ex), m.up_exps.macro_scale(ex), &mut act, m_e * n_ff_exp)?;
4201                let actv = act.slice(0..m_e * n_ff_exp);
4202                e.qmatvec_view(sd, 0..dl.len, &actv, m_e,
4203                    m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)?
4204            };
4205
4206            // SCATTER into slot buffer: each row goes to slot_buf[tok, slot, :].
4207            e.scatter_slot(&y, &tok_idx_d, &slot_idx_d, &weight_d,
4208                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4209        }
4210
4211        // 5. REDUCE SLOTS: sum the 8 slots per token into the final moe_out.
4212        let mut moe_out = e.zeros(t * n_embd)?;
4213        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, t)?;
4214
4215        // STATS: print m-distribution when MEMRA_MOE_STATS is set.
4216        if std::env::var("MEMRA_MOE_STATS").is_ok() && !m_dist.is_empty() {
4217            m_dist.sort_unstable();
4218            let active = m_dist.len();
4219            let mean = m_dist.iter().sum::<usize>() as f64 / active as f64;
4220            let median = m_dist[active / 2];
4221            let max_m = *m_dist.last().unwrap();
4222            let min_m = m_dist[0];
4223            let above16 = m_dist.iter().filter(|&&x| x >= 16).count();
4224            println!("moe-grouped il={il} t={t} active={active}/{n_expert} \
4225                      m_e: min={min_m} median={median} mean={mean:.1} max={max_m} \
4226                      above_gemm_threshold(>=16)={above16}/{active}");
4227        }
4228
4229        // 6. SHARED EXPERT (same as moe_ffn — untouched).
4230        // gate_inp_shexp is OPTIONAL: qwen35moe gates the shared expert (sigmoid(gate_inp) x sh);
4231        // MiniMax-M3 (DeepSeek-V3 class) has NO shexp gate — the shared expert adds directly.
4232        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4233            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4234        {
4235            let n_ff_sh = gate_shexp.out_features();
4236            let sg_gate = e.matmul(gate_shexp, z, t)?;
4237            let sg_up = e.matmul(up_shexp, z, t)?;
4238            let mut sa = e.zeros(t * n_ff_sh)?;
4239            Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, t * n_ff_sh)?;
4240            let sh = e.matmul(down_shexp, &sa, t)?;
4241            // shexp gate: qwen35moe sigmoid-gates; M3 has no gate tensor -> weight 1.0.
4242            // Fused sigmoid-dot below PRIME_MIN_T — one fold order with the sequential and
4243            // dev decode arms (dispatch choice must not change bits).
4244            let g = match &m.gate_inp_shexp {
4245                Some(gate_inp_shexp) => {
4246                    // router_prefill_exact_on(): the fused sigmoid-dot is m-INVARIANT and
4247                    // serves EVERY t (serve isolation, lane/concat-prime-exact).
4248                    if t < PRIME_MIN_T || crate::router_prefill_exact_on() {
4249                        e.sigmoid_dot_rows(z, gate_inp_shexp.float_data(), n_embd, t)?
4250                    } else {
4251                        let gs = e.linear(z, gate_inp_shexp.float_data(), t, n_embd, 1)?;
4252                        let mut g = e.uninit(t)?;
4253                        e.sigmoid(&gs, &mut g, t)?;
4254                        g
4255                    }
4256                }
4257                None => e.htod(&vec![1.0f32; t])?,
4258            };
4259            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, t)?;
4260        }
4261
4262        Ok(moe_out)
4263    }
4264
4265    /// Lane-3 M2: cross-stream MoE for lockstep decode. Routes all m stream rows in one
4266    /// batch, executes fully-HBM-resident experts through the grouped gather/GEMM/scatter
4267    /// machinery at m_e>1 (weight reads amortized across streams), and assigns any expert
4268    /// with a missing projection to that row's CPU companion call (whole-expert granularity,
4269    /// same rule as the sequential frozen path). Slot-pinned accumulation keeps each row's
4270    /// expert-sum order identical to the sequential path.
4271    pub(crate) fn moe_ffn_lockstep(
4272        &self,
4273        e: &Engine,
4274        m: &MoeWeights,
4275        zbatch: &CudaSlice<f32>,
4276        mrows: usize,
4277        il: u16,
4278        max_block: usize,
4279    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4280        use crate::moe_cache::{BlockId, PROJ_DOWN, PROJ_GATE, PROJ_UP};
4281        let cfg = &self.cfg;
4282        let moe = cfg.moe.as_ref().unwrap();
4283        let n_embd = cfg.n_embd as usize;
4284        let n_expert = moe.expert_count as usize;
4285        let n_used = moe.expert_used_count as usize;
4286        let n_ff_exp = moe.expert_ff_length as usize;
4287
4288        let logits = e.matmul(&m.gate_inp, zbatch, mrows)?;
4289        let (sel_all, w_all) = if let Some(sig) = cfg.sigmoid_router() {
4290            Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4291                                m.exp_probs_b.as_deref(), Some(sig), m.active_experts.as_deref())?
4292        } else {
4293            Self::moe_route_cfg(e, &logits, mrows, n_expert, n_used,
4294                                None, None, m.active_experts.as_deref())?
4295        };
4296        Self::trace_moe_routes(il, mrows, &sel_all, &w_all)?;
4297
4298        // Residency split at whole-expert granularity against the (frozen) cache.
4299        let resident_expert: Vec<bool> = e.with_moe_cache(max_block, |c, _| {
4300            Ok((0..n_expert)
4301                .map(|ex| {
4302                    [PROJ_GATE, PROJ_UP, PROJ_DOWN].into_iter().all(|p| {
4303                        c.resident(BlockId::new(il, p, ex as u16)).is_some()
4304                    })
4305                })
4306                .collect())
4307        })?;
4308
4309        struct Group {
4310            rows: Vec<i32>,
4311            slots: Vec<i32>,
4312            weights: Vec<f32>,
4313        }
4314        let mut groups: std::collections::HashMap<usize, Group> = Default::default();
4315        let mut cpu_rows: Vec<Vec<(usize, f32)>> = vec![Vec::new(); mrows];
4316        let mut cpu_by_expert: std::collections::HashMap<usize, Vec<(usize, f32)>> =
4317            Default::default();
4318        for row in 0..mrows {
4319            for j in 0..n_used {
4320                let ex = sel_all[row * n_used + j] as usize;
4321                let w = w_all[row * n_used + j];
4322                if resident_expert[ex] {
4323                    let group = groups.entry(ex).or_insert_with(|| Group {
4324                        rows: Vec::new(),
4325                        slots: Vec::new(),
4326                        weights: Vec::new(),
4327                    });
4328                    group.rows.push(row as i32);
4329                    group.slots.push(j as i32);
4330                    group.weights.push(w);
4331                } else {
4332                    crate::cpu_experts::record_incomplete_gpu_residency(0);
4333                    cpu_rows[row].push((ex, w));
4334                    cpu_by_expert.entry(ex).or_default().push((row, w));
4335                }
4336            }
4337        }
4338
4339        // CPU tickets first: reads/compute overlap the GPU grouped work below. Experts routed
4340        // by >=2 streams go through the multi-row ABI (weight decode amortized across rows);
4341        // each row's remaining experts stay one ordinary per-row call. Contribution FP-sum
4342        // order per row differs from the sequential single-call chunk — part of the
4343        // documented lockstep numeric class.
4344        let host_rows = e.dtoh(zbatch)?;
4345        let rows_ok = crate::cpu_experts::rows_supported();
4346        enum CpuPart {
4347            Single { row: usize },
4348            Rows { rows: Vec<usize> },
4349        }
4350        let mut tickets: Vec<(CpuPart, crate::cpu_experts::CpuExpertTicket)> = Vec::new();
4351        let mut rows_served: std::collections::HashSet<(usize, usize)> = Default::default();
4352        if rows_ok {
4353            let mut shared: Vec<(usize, Vec<(usize, f32)>)> = cpu_by_expert
4354                .into_iter()
4355                .filter(|(_, rows)| rows.len() >= 2)
4356                .collect();
4357            shared.sort_by_key(|(ex, _)| *ex);
4358            for (ex, mut row_weights) in shared {
4359                row_weights.sort_by_key(|(row, _)| *row);
4360                let inputs: Vec<(&[f32], f32)> = row_weights
4361                    .iter()
4362                    .map(|&(row, w)| (&host_rows[row * n_embd..(row + 1) * n_embd], w))
4363                    .collect();
4364                let job = crate::cpu_experts::prepare_rows_job(m, ex, &inputs)
4365                    .map_err(std::io::Error::other)?;
4366                for &(row, _) in &row_weights {
4367                    rows_served.insert((row, ex));
4368                }
4369                tickets.push((
4370                    CpuPart::Rows {
4371                        rows: row_weights.iter().map(|&(row, _)| row).collect(),
4372                    },
4373                    crate::cpu_experts::submit_rows(job).map_err(std::io::Error::other)?,
4374                ));
4375            }
4376        }
4377        for (row, selected) in cpu_rows.iter().enumerate() {
4378            let leftover: Vec<(usize, f32)> = selected
4379                .iter()
4380                .copied()
4381                .filter(|&(ex, _)| !rows_served.contains(&(row, ex)))
4382                .collect();
4383            if leftover.is_empty() {
4384                continue;
4385            }
4386            let host_row = &host_rows[row * n_embd..(row + 1) * n_embd];
4387            let job = crate::cpu_experts::prepare_job(m, il, &leftover, host_row)
4388                .map_err(std::io::Error::other)?;
4389            tickets.push((
4390                CpuPart::Single { row },
4391                crate::cpu_experts::submit(job).map_err(std::io::Error::other)?,
4392            ));
4393        }
4394
4395        let mut slot_buf = e.zeros(mrows * n_used * n_embd)?;
4396        let mut wbuf = e.zeros(mrows * n_used)?;
4397        let mut order: Vec<usize> = groups.keys().copied().collect();
4398        order.sort_by(|&a, &b| {
4399            groups[&b].rows.len().cmp(&groups[&a].rows.len()).then(a.cmp(&b))
4400        });
4401        for &ex in &order {
4402            let group = &groups[&ex];
4403            let m_e = group.rows.len();
4404            let gl = m.gate_exps.expert_layout(ex);
4405            let ul = m.up_exps.expert_layout(ex);
4406            let dl = m.down_exps.expert_layout(ex);
4407            let row_idx_d = e.htod_i32(&group.rows)?;
4408            let slot_idx_d = e.htod_i32(&group.slots)?;
4409            let dmac = m.down_exps.macro_scale(ex);
4410            let weight_d = if dmac == 1.0 {
4411                e.htod(&group.weights)?
4412            } else {
4413                let scaled: Vec<f32> = group.weights.iter().map(|&w| w * dmac).collect();
4414                e.htod(&scaled)?
4415            };
4416            let mut gathered = e.zeros(m_e * n_embd)?;
4417            e.gather_rows(zbatch, &row_idx_d, &mut gathered, n_embd, m_e)?;
4418            let gv = gathered.slice(0..m_e * n_embd);
4419            let gate = e.with_moe_cache(max_block, |c, eng| {
4420                let slot = c
4421                    .resident(BlockId::new(il, PROJ_GATE, ex as u16))
4422                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4423                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..gl.len, &gv, m_e,
4424                    m.gate_exps.in_f, m.gate_exps.out_f, gl.qtype, gl.row_bytes)
4425            })?;
4426            let up = e.with_moe_cache(max_block, |c, eng| {
4427                let slot = c
4428                    .resident(BlockId::new(il, PROJ_UP, ex as u16))
4429                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4430                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..ul.len, &gv, m_e,
4431                    m.up_exps.in_f, m.up_exps.out_f, ul.qtype, ul.row_bytes)
4432            })?;
4433            let mut act = e.zeros(m_e * n_ff_exp)?;
4434            Self::ffn_act_scaled(e, cfg, &gate, &up,
4435                m.gate_exps.macro_scale(ex), m.up_exps.macro_scale(ex), &mut act, m_e * n_ff_exp)?;
4436            let actv = act.slice(0..m_e * n_ff_exp);
4437            let y = e.with_moe_cache(max_block, |c, eng| {
4438                let slot = c
4439                    .resident(BlockId::new(il, PROJ_DOWN, ex as u16))
4440                    .ok_or("lockstep resident expert vanished (cache not frozen?)")?;
4441                eng.qmatvec_view(c.buf(crate::moe_cache::DispatchSlot::Resident(slot)), 0..dl.len, &actv, m_e,
4442                    m.down_exps.in_f, m.down_exps.out_f, dl.qtype, dl.row_bytes)
4443            })?;
4444            e.scatter_slot(&y, &row_idx_d, &slot_idx_d, &weight_d,
4445                           &mut slot_buf, &mut wbuf, n_embd, n_used, m_e)?;
4446        }
4447        let mut moe_out = e.zeros(mrows * n_embd)?;
4448        e.reduce_slots(&slot_buf, &wbuf, &mut moe_out, n_embd, n_used, mrows)?;
4449
4450        // CPU contributions join BEFORE the shared expert (the sequential path's placement).
4451        let mut row_sums: Vec<Option<Vec<f32>>> = vec![None; mrows];
4452        for (part, ticket) in tickets {
4453            let cpu_output = ticket.wait().map_err(std::io::Error::other)?;
4454            let mut add_row = |row: usize, chunk: &[f32]| {
4455                let sum = row_sums[row].get_or_insert_with(|| vec![0.0f32; n_embd]);
4456                for (accumulator, value) in sum.iter_mut().zip(chunk) {
4457                    *accumulator += value;
4458                }
4459            };
4460            match part {
4461                CpuPart::Single { row } => add_row(row, &cpu_output),
4462                CpuPart::Rows { rows } => {
4463                    for (slot, row) in rows.into_iter().enumerate() {
4464                        add_row(row, &cpu_output[slot * n_embd..(slot + 1) * n_embd]);
4465                    }
4466                }
4467            }
4468        }
4469        for (row, sum) in row_sums.into_iter().enumerate() {
4470            let Some(sum) = sum else { continue };
4471            let cpu_output = e.htod(&sum)?;
4472            let mut dst = moe_out.slice_mut(row * n_embd..(row + 1) * n_embd);
4473            e.axpy_into(&cpu_output, 1.0, &mut dst, n_embd)?;
4474        }
4475
4476        if let (Some(gate_shexp), Some(up_shexp), Some(down_shexp)) =
4477            (&m.gate_shexp, &m.up_shexp, &m.down_shexp)
4478        {
4479            let n_ff_sh = gate_shexp.out_features();
4480            let sg_gate = e.matmul(gate_shexp, zbatch, mrows)?;
4481            let sg_up = e.matmul(up_shexp, zbatch, mrows)?;
4482            let mut sa = e.zeros(mrows * n_ff_sh)?;
4483            Self::ffn_act(e, cfg, &sg_gate, &sg_up, &mut sa, mrows * n_ff_sh)?;
4484            let sh = e.matmul(down_shexp, &sa, mrows)?;
4485            // lockstep rows ARE decode tokens: fused sigmoid-dot per row so batched serving
4486            // decode matches the single-sequence decode chain bit-for-bit.
4487            let g = match &m.gate_inp_shexp {
4488                Some(gate_inp_shexp) => {
4489                    e.sigmoid_dot_rows(zbatch, gate_inp_shexp.float_data(), n_embd, mrows)?
4490                }
4491                None => e.htod(&vec![1.0f32; mrows])?,
4492            };
4493            e.add_scaled_rows(&sh, &g, &mut moe_out, n_embd, mrows)?;
4494        }
4495
4496        Ok(moe_out)
4497    }
4498}
4499
4500// ============================ gemma4 (R8 verified wiring) ==================================
4501// Node-for-node vs llama.cpp src/models/gemma4.cpp:180-405 (HANDOVER "R8 VERIFIED WIRING").
4502// v0 bring-up: full attention everywhere (exact for prompts < sliding_window 1024 — R6 masking
4503// later), sdpa_naive (hd 512 has no FA stamp), sequential host-staged MoE (the perf arms grow
4504// gemma variants after the correctness gate).
4505impl HybridModel {
4506    /// Per-layer attention geometry (R5): (head_dim, n_kv, n_head, rope_base, scale, is_swa).
4507    pub(crate) fn gemma4_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
4508        let g = self.cfg.gemma4.as_ref().unwrap();
4509        let swa = g.swa_pattern[il];
4510        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
4511        // attention scale = 1.0 (llama gemma4.cpp:11 "Gemma4 uses self.scaling = 1.0" — q/k are
4512        // per-head rms-normed; NOT the 1/sqrt(hd) default. Bring-up bug: 1/sqrt(hd) left token-0
4513        // rows exact (softmax over one element) while every later position drifted).
4514        (hd, g.head_count_kv[il] as usize, self.cfg.n_head as usize,
4515         if swa { g.rope_base_swa } else { g.rope_base_global },
4516         1.0, swa)
4517    }
4518
4519    /// Suppress-token mask over t logits rows (tokenizer.ggml.suppress_tokens; no-op when the
4520    /// model ships none). NOT monotonic like softcap — must run before every argmax/sample, so
4521    /// every gemma4 logits tail (forward/prime/decode/dc/verify/e4b) calls this on device ld.
4522    fn gemma4_suppress(&self, e: &Engine, ld: &mut CudaSlice<f32>, t: usize)
4523                       -> Result<(), Box<dyn std::error::Error>> {
4524        if let Some((ids, n)) = self.gemma4_aux.as_ref().and_then(|a| a.suppress_d.as_ref()) {
4525            e.mask_ids_rows(ld, ids, *n, self.output.out_features(), t)?;
4526        }
4527        Ok(())
4528    }
4529
4530    /// gemma4 attention (R5 geometry, R7 weightless V-norm on the RAW K projection, R9 dual rope).
4531    /// `cache`: Some => PRIME mode — append the T post-rope K / normed V rows into the quantized
4532    /// KV cache (same per-row quantize math as the decode append) and advance len. Fresh-prompt
4533    /// only (v0): attends within `tokens` via the f32 sdpa.
4534    fn gemma4_attn_prime(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
4535                         h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize,
4536                         cache: Option<&mut Cache>)
4537                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4538        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
4539        let eps = self.cfg.rms_eps;
4540        let aux = self.gemma4_aux.as_ref().unwrap();
4541
4542        // quantize-once window: q/k/v share `h` — the MMQ D4 activation quantizes once
4543        // (h stays borrowed across the triple, so the cache key can't go stale).
4544        e.mmq_act_begin();
4545        let q0 = e.matmul(&fa.wq, h, t)?;   // [t, nh*hd]
4546        let k0 = e.matmul(&fa.wk, h, t)?;   // [t, nkv*hd]
4547        // globals ship no v_proj (wv := wk at load) — V input is the SAME projection output;
4548        // reuse k0 instead of re-running the identical matmul (K=V dedup, 5 layers).
4549        let v0 = if swa { e.matmul(&fa.wv, h, t)? } else { e.clone_dtod(&k0)? };
4550
4551        let mut q = e.uninit(t * nh * hd)?;
4552        let mut k = e.uninit(t * nkv * hd)?;
4553        // R7: V = weightless rms_norm of the raw projection; NEVER roped.
4554        let mut v = e.uninit(t * nkv * hd)?;
4555        // 31B glue lane: producers emit the bf16 FA operands (norm emits vb; rope emits qb/kb
4556        // post-rope) — kills 3 f32->bf16 converts + re-reads per layer. Bit-identical operands
4557        // (same __float2bfloat16); MEMRA_FA_EMIT=0 reverts to the convert-in-FA path.
4558        static EMIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4559        let emit = t >= 16 && crate::Engine::qkvnorm_w_on_prefill(nh * t + 2 * nkv * t, hd)
4560            && *EMIT.get_or_init(|| std::env::var("MEMRA_FA_EMIT").map(|s| s != "0").unwrap_or(true));
4561        let mut qb = e.alloc_uninit::<u8>(if emit { t * nh * hd * 2 } else { 1 })?;
4562        let mut kb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
4563        let mut vb = e.alloc_uninit::<u8>(if emit { t * nkv * hd * 2 } else { 1 })?;
4564        // f16-P/V door: emit V as f16 straight from the norm when this layer's FA consumer
4565        // reads f16 — kills the per-layer bf16->f16 re-encode (1.4GB/pass on 31B SWA).
4566        let v_f16 = emit && crate::fa_f16pv_on() && match hd {
4567            512 => true,
4568            256 => swa && crate::faw_hp_on() && nh % 2 == 0 && (nh / nkv) % 2 == 0,
4569            _ => false,
4570        };
4571        if emit {
4572            e.rms_norm_qkv_w4b(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
4573                               &aux.ones, &mut q, &mut k, &mut v, &mut vb,
4574                               hd, nh * t, nkv * t, eps, v_f16)?;
4575        } else {
4576            e.rms_norm_qkv(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
4577                           &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t, eps)?;
4578        }
4579
4580        let ff = if swa { None } else {
4581            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
4582        };
4583        if emit {
4584            e.rope_neox2_bf16e(&mut q, &mut k, &mut qb, &mut kb, pos_d, hd, hd, nh, nkv, t,
4585                               base, 1.0, ff)?;
4586        } else {
4587            e.rope_neox2(&mut q, &mut k, pos_d, hd, hd, nh, nkv, t, base, 1.0, ff)?;
4588        }
4589
4590        if let Some(cache) = cache {
4591            let kvl = cache.kv[il].as_mut().unwrap();
4592            assert_eq!(kvl.len, 0, "gemma4 prime is fresh-prompt only (v0)");
4593            e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
4594                                       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()))?;
4595            kvl.len += t;
4596        }
4597        let mut attn = e.zeros(t * nh * hd)?;
4598        // R6: SWA layers mask keys older than sliding_window once the prompt exceeds it
4599        // (windowed naive twin; fa windowed stamps later). Under the window, full attention
4600        // is exact — SWA rides fa_prefill (hd-256 stamp), the hd-512 globals stay naive.
4601        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
4602        if swa && t > win {
4603            if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
4604                if emit { e.fa_prefill_w_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
4605                                             scale, true, win, v_f16)?; }
4606                else { e.fa_prefill_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true,
4607                                      win)?; }
4608            } else {
4609                e.sdpa_naive_w(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true, win)?;
4610            }
4611        } else if hd == 256 && std::env::var("MEMRA_NOFA").is_err() {
4612            e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
4613        } else if hd == 512 && std::env::var("MEMRA_NOFA").is_err() {
4614            if emit { e.fa_prefill_hd512_pre(&qb, &kb, &vb, &mut attn, hd, nh, nkv, t, t,
4615                                             scale, true, v_f16)?; }
4616            else { e.fa_prefill_hd512(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?; }
4617        } else {
4618            e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, t, t, scale, true)?;
4619        }
4620        Ok(e.matmul(&fa.wo, &attn, t)?)
4621    }
4622
4623    /// Back-compat wrapper (pure prefill, no cache).
4624    fn gemma4_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
4625                   h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
4626                   -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4627        self.gemma4_attn_prime(e, fa, il, h, pos_d, t, None)
4628    }
4629
4630    /// gemma4 MoE with the expert input PRE-QUANTIZED (the q8z tail fusion). Caller guarantees
4631    /// the fast-arm conditions (resident dev slabs + dp4a qtypes) — decode t=1 and verify rows
4632    /// arms only, per-token kernel chains identical to the f32-input path (same quantize bytes:
4633    /// the q8z epilogue is quantize_q8_1 verbatim).
4634    fn gemma4_moe_q8(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
4635                     bits: &crate::hybrid::Gemma4MoeBits,
4636                     mq: &(CudaSlice<i8>, CudaSlice<f32>),
4637                     router_in: &CudaSlice<f32>, t: usize)
4638                     -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4639        let cfg = &self.cfg;
4640        let moe = cfg.moe.as_ref().unwrap();
4641        let n_embd = cfg.n_embd as usize;
4642        let n_expert = moe.expert_count as usize;
4643        let n_used = moe.expert_used_count as usize;
4644        let n_ff_exp = moe.expert_ff_length as usize;
4645        // router stays the two-launch pair: THREE fuse variants measured worse (serial-dot
4646        // -50%, 1024-thread warp-parallel -12% on topk sync overhead — jsonl 2026-07-11);
4647        // the pair's 12us is kernel time, not launch gaps.
4648        let logits = if crate::router_kernel_on() {
4649            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
4650        } else {
4651            e.matmul(&m.gate_inp, router_in, t)?
4652        };
4653        let dev = m.dev_exps.as_ref().unwrap();
4654        let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
4655                                                    &bits.per_expert_scale_d)?;
4656        let (zq, zd) = mq;
4657        if t == 1 {
4658            let selv = sel_d.slice(0..n_used);
4659            let wv = w_d.slice(0..n_used);
4660            let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, zq, zd,
4661                                                 n_embd, n_ff_exp, n_used, n_expert,
4662                                                 m.gate_exps.qtype, m.up_exps.qtype,
4663                                                 m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4664            let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4665            let mut moe_out = e.uninit(n_embd)?;
4666            e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
4667                                   &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
4668                                   n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
4669            return Ok(moe_out);
4670        }
4671        let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
4672        let act = if csr {
4673            e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, zq, zd, t * n_used,
4674                                           n_embd, n_ff_exp, n_used, n_expert,
4675                                           m.gate_exps.qtype, m.up_exps.qtype,
4676                                           m.gate_exps.row_bytes, m.up_exps.row_bytes)?
4677        } else {
4678            e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, zq, zd, t,
4679                                            n_embd, n_ff_exp, n_used, n_expert,
4680                                            m.gate_exps.qtype, m.up_exps.qtype,
4681                                            m.gate_exps.row_bytes, m.up_exps.row_bytes)?
4682        };
4683        let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
4684        let mut moe_out = e.uninit(t * n_embd)?;
4685        // down stays rows_g: the CSR dedup twin measured NEGATIVE at nsb=22 too (189.4 vs
4686        // 207.0 depth spec, bitwise-exact — jsonl 2026-07-10; qwen nsb=16 same verdict).
4687        e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
4688                                      n_ff_exp, n_embd, n_used, n_expert,
4689                                      m.down_exps.qtype, m.down_exps.row_bytes)?;
4690        Ok(moe_out)
4691    }
4692
4693    /// gemma4 MoE (R2 router prologue input supplied by caller, R3 per-expert output scale).
4694    /// Sequential host-staged v0 — softmax gating + renorm (moe_route, the qwen recipe), GELU
4695    /// experts, scale folded into the accumulate weight (post-matmul linear scale, exact fold).
4696    fn gemma4_moe(&self, e: &Engine, m: &crate::hybrid::MoeWeights,
4697                  bits: &crate::hybrid::Gemma4MoeBits, moe_in: &CudaSlice<f32>,
4698                  router_in: &CudaSlice<f32>, t: usize)
4699                  -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4700        let cfg = &self.cfg;
4701        let moe = cfg.moe.as_ref().unwrap();
4702        let n_embd = cfg.n_embd as usize;
4703        let n_expert = moe.expert_count as usize;
4704        let n_used = moe.expert_used_count as usize;
4705        let n_ff_exp = moe.expert_ff_length as usize;
4706
4707        // Router: the in-house GEMV for ALL small t (decode AND verify ride the same per-column
4708        // kernel — the cuBLASLt n-dependence flipped top-k at verify t on the 27B, d994271);
4709        // batched matmul only at real prefill.
4710        let logits = if t < PRIME_MIN_T && crate::router_kernel_on() {
4711            e.router_gemv(m.gate_inp.float_data(), router_in, n_embd, n_expert, t)?
4712        } else {
4713            e.matmul(&m.gate_inp, router_in, t)?
4714        };
4715
4716        // FAST SMALL-T ARM (decode t=1 AND spec verify t=2..15): device softmax-topk router,
4717        // then PER TOKEN the same fused gate_up GELU + down8 FMA launch pair over the resident
4718        // dev slabs — verify rides the EXACT decode kernel chain per token (dispatch-parity
4719        // law; the qwen "verify must be kernel-dispatch-identical to decode" lesson).
4720        if t < PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
4721            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
4722            && expert_dp4a_supported(m.down_exps.qtype)
4723            && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0") {
4724            let dev = m.dev_exps.as_ref().unwrap();
4725            let (sel_d, w_d) = e.moe_router_topk_scaled(&logits, t, n_expert, n_used,
4726                                                        &bits.per_expert_scale_d)?;
4727            if t == 1 {
4728                let (zq, zd) = e.quantize_q8_1(moe_in, 1, n_embd)?;
4729                let selv = sel_d.slice(0..n_used);
4730                let wv = w_d.slice(0..n_used);
4731                let act = e.moe_gate_up_gelu8_dev_q8(&dev.ptr_row, &selv, &zq, &zd,
4732                                                     n_embd, n_ff_exp, n_used, n_expert,
4733                                                     m.gate_exps.qtype, m.up_exps.qtype,
4734                                                     m.gate_exps.row_bytes, m.up_exps.row_bytes)?;
4735                let (aq2, ad2) = e.quantize_q8_1(&act, n_used, n_ff_exp)?;
4736                let mut moe_out = e.uninit(n_embd)?;
4737                e.moe_down8_fma_dev_q8(&dev.ptr_row, &selv, &wv, &aq2, &ad2,
4738                                       &mut moe_out.slice_mut(0..n_embd), n_ff_exp, n_embd,
4739                                       n_used, n_expert, m.down_exps.qtype, m.down_exps.row_bytes)?;
4740                return Ok(moe_out);
4741            }
4742            // VERIFY ROWS TWINS (t=2..15): ONE launch pair for all tokens; per (token,row,slot)
4743            // bodies are the t=1 kernels VERBATIM (bit-identical to the per-token loop).
4744            // gate_up rides the CSR owner-scan dedup when t <= 10 (each duplicated expert's
4745            // weight stream decoded once; the qwen CSR class). MEMRA_GEMMA_CSR=0 -> rows.
4746            let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
4747            let csr = t <= 10 && std::env::var("MEMRA_GEMMA_CSR").as_deref() != Ok("0");
4748            let act = if csr {
4749                e.moe_gate_up_gelu8_dev_q8_csr(&dev.ptr_row, &sel_d, &zq, &zd, t * n_used,
4750                                               n_embd, n_ff_exp, n_used, n_expert,
4751                                               m.gate_exps.qtype, m.up_exps.qtype,
4752                                               m.gate_exps.row_bytes, m.up_exps.row_bytes)?
4753            } else {
4754                e.moe_gate_up_gelu8_dev_q8_rows(&dev.ptr_row, &sel_d, &zq, &zd, t,
4755                                                n_embd, n_ff_exp, n_used, n_expert,
4756                                                m.gate_exps.qtype, m.up_exps.qtype,
4757                                                m.gate_exps.row_bytes, m.up_exps.row_bytes)?
4758            };
4759            let (aq2, ad2) = e.quantize_q8_1(&act, t * n_used, n_ff_exp)?;
4760            let mut moe_out = e.uninit(t * n_embd)?;
4761            e.moe_down8_fma_dev_q8_rows_g(&dev.ptr_row, &sel_d, &w_d, &aq2, &ad2, &mut moe_out, t,
4762                                          n_ff_exp, n_embd, n_used, n_expert,
4763                                          m.down_exps.qtype, m.down_exps.row_bytes)?;
4764            return Ok(moe_out);
4765        }
4766
4767        let (sel_all, mut w_all) = Self::moe_route(e, &logits, t, n_expert, n_used)?;
4768        for (i, &sx) in sel_all.iter().enumerate() {
4769            w_all[i] *= bits.per_expert_scale[sx as usize];
4770        }
4771
4772        // PREFILL PAIRS ARM (t >= 16): expert-major CSR over the resident slabs — ONE launch per
4773        // projection covers ALL (token,expert) pairs (the qwen pairs recipe, _em dot = expert_dot_g
4774        // which has the Q4_0 body; GELU pairs epilogue; R3 scale folded into pair_w).
4775        if t >= PRIME_MIN_T && m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
4776            && expert_dp4a_supported(m.gate_exps.qtype) && expert_dp4a_supported(m.up_exps.qtype)
4777            && expert_dp4a_supported(m.down_exps.qtype)
4778            && std::env::var("MEMRA_GEMMA_MOE_PAIRS").as_deref() != Ok("0") {
4779            let dev = m.dev_exps.as_ref().unwrap();
4780            let n_pairs = t * n_used;
4781            let pair_ex: Vec<i32> = sel_all.iter().map(|&x| x as i32).collect();
4782            let pair_tok: Vec<i32> = (0..n_pairs).map(|p| (p / n_used) as i32).collect();
4783            let tok_off: Vec<i32> = (0..=t).map(|tok| (tok * n_used) as i32).collect();
4784            let tok_ids: Vec<i32> = (0..n_pairs as i32).collect();
4785            let pt = e.htod_i32(&pair_tok)?;
4786            let pw = e.htod(&w_all)?;
4787            let toff = e.htod_i32(&tok_off)?;
4788            let tids = e.htod_i32(&tok_ids)?;
4789            let mut by_ex: Vec<Vec<i32>> = vec![Vec::new(); n_expert];
4790            for p in 0..n_pairs { by_ex[pair_ex[p] as usize].push(p as i32); }
4791            let mut ex_ids: Vec<i32> = Vec::new();
4792            let mut ex_off: Vec<i32> = vec![0];
4793            let mut ex_pairs: Vec<i32> = Vec::with_capacity(n_pairs);
4794            for (ex, list) in by_ex.iter().enumerate() {
4795                if list.is_empty() { continue; }
4796                ex_ids.push(ex as i32);
4797                ex_pairs.extend_from_slice(list);
4798                ex_off.push(ex_pairs.len() as i32);
4799            }
4800            let n_active = ex_ids.len();
4801            let exi = e.htod_i32(&ex_ids)?;
4802            let exo = e.htod_i32(&ex_off)?;
4803            let exp_d = e.htod_i32(&ex_pairs)?;
4804            // GROUPED f16 LANE (MEMRA_MOE_F16G=1, round 46 arc 2): dequant active experts to
4805            // f16 once per projection + one grouped f16 GEMM over the CSR groups; CSR order
4806            // end-to-end (gelu is elementwise), one row permute before the scatter. The
4807            // ragged down k (704) needs no padding here — cublas takes any k.
4808            // f16-mirror numeric class, argmax/spec gated. PER-MODEL default OFF (round 50):
4809            // the gelu class regressed g26 board-2048 prefill -8.3% under the round-49
4810            // Hopper default — see moe_f16g_gemma_on.
4811            if crate::moe_f16g_gemma_on()
4812                && f16g_proj_ok(m.gate_exps.qtype, n_embd)
4813                && f16g_proj_ok(m.up_exps.qtype, n_embd)
4814                && f16g_proj_ok(m.down_exps.qtype, n_ff_exp) {
4815                let csr_tok: Vec<i32> = ex_pairs.iter().map(|&p| p / n_used as i32).collect();
4816                let csr_tok_d = e.htod_i32(&csr_tok)?;
4817                let (z_f16, z_s) = e.moe_f16g_act(moe_in, Some(&csr_tok_d), n_embd, n_pairs)?;
4818                let g_csr = e.moe_f16_grouped(&dev.ptr_row, 0, n_expert, &exi, &ex_off, &exo,
4819                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4820                                              m.gate_exps.qtype, m.gate_exps.row_bytes)?;
4821                let u_csr = e.moe_f16_grouped(&dev.ptr_row, 1, n_expert, &exi, &ex_off, &exo,
4822                                              &z_f16, &z_s, n_embd, n_ff_exp, n_active, n_pairs,
4823                                              m.up_exps.qtype, m.up_exps.row_bytes)?;
4824                let act_csr = e.moe_pairs_gelu_mul(&g_csr, &u_csr, n_pairs * n_ff_exp)?;
4825                let (a_f16, a_s) = e.moe_f16g_act(&act_csr, None, n_ff_exp, n_pairs)?;
4826                let d_csr = e.moe_f16_grouped(&dev.ptr_row, 2, n_expert, &exi, &ex_off, &exo,
4827                                              &a_f16, &a_s, n_ff_exp, n_embd, n_active, n_pairs,
4828                                              m.down_exps.qtype, m.down_exps.row_bytes)?;
4829                let y_down = e.rows_permute(&d_csr, &exp_d, n_pairs, n_embd)?;
4830                let mut moe_out = e.uninit(t * n_embd)?;
4831                e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
4832                if std::env::var("MEMRA_F16G_DEBUG").is_ok() {
4833                    let scan = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
4834                    let (yd, mo) = (e.dtoh(&y_down)?, e.dtoh(&moe_out)?);
4835                    eprintln!("[f16g-debug] post-permute bad={} post-scatter bad={}",
4836                              scan(&yd), scan(&mo));
4837                }
4838                return Ok(moe_out);
4839            }
4840            // gate/up: int8-MMA expert GEMM (in_f = n_embd 2816 = 11x256 tiles ok); down keeps
4841            // the decode-once dp4a (in_f 704 fails the 256-superblock tiling).
4842            let mma = n_embd % 256 == 0
4843                && std::env::var("MEMRA_GEMMA_MOE_MMA").as_deref() != Ok("0");
4844            let (gate, up) = if mma {
4845                let z_scr = e.mmq_iq_quantize_act(moe_in, n_embd, t)?;
4846                (e.mmq_iq_experts(&dev.ptr_row, 0, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4847                                  n_embd, n_ff_exp, n_active, n_pairs, t,
4848                                  m.gate_exps.qtype, m.gate_exps.row_bytes)?,
4849                 e.mmq_iq_experts(&dev.ptr_row, 1, n_expert, &exi, &exo, &exp_d, &pt, &z_scr,
4850                                  n_embd, n_ff_exp, n_active, n_pairs, t,
4851                                  m.up_exps.qtype, m.up_exps.row_bytes)?)
4852            } else {
4853                let (zq, zd) = e.quantize_q8_1(moe_in, t, n_embd)?;
4854                (e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 0, &exi, &exo, &exp_d, &pt, &zq, &zd,
4855                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
4856                                           m.gate_exps.qtype, m.gate_exps.row_bytes)?,
4857                 e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 1, &exi, &exo, &exp_d, &pt, &zq, &zd,
4858                                           n_embd, n_ff_exp, n_expert, n_active, n_pairs,
4859                                           m.up_exps.qtype, m.up_exps.row_bytes)?)
4860            };
4861            let pair_self: Vec<i32> = (0..n_pairs as i32).collect();
4862            let pself = e.htod_i32(&pair_self)?;
4863            // DOWN through the int8-MMA expert GEMM (2026-07-31, g26 prefill lever): the
4864            // ragged k (n_ff_exp=704 on the 26B) rides a PADDED k-walk — in_f rounds up
4865            // to the 256-val superblock (768) while the act quantizer's zero padding
4866            // makes every padded-k product exactly zero (weight overread bytes multiply
4867            // zero int8 act values; the dev slab carries 144B tail slack for the OOB).
4868            // The old dp4a matvec was 11.3ms/call at m=T (the 0.07x prefill wall).
4869            // MEMRA_GEMMA_MOE_MMA=0 reverts down together with gate/up.
4870            // FUSED ACT-EPILOGUE (default on): gelu_tanh(gate)*up + D4 quantize in one
4871            // launch — no f32 act buffer (the fused kernel zero-pads the ragged tail
4872            // exactly like the two-pass quantizer). MEMRA_MOE_FUSE_ACTQ=0 rollback.
4873            // Scratch bytes are BYTE-IDENTICAL (kernel-check gated).
4874            let y_down = if mma {
4875                let in_pad = n_ff_exp.div_ceil(256) * 256;
4876                let a_scr = if crate::moe_fuse_actq_on() {
4877                    e.mmq_iq_fused_act_quant(&gate, &up, n_ff_exp, n_pairs, 1)?
4878                } else {
4879                    let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
4880                    e.mmq_iq_quantize_act(&act, n_ff_exp, n_pairs)?
4881                };
4882                e.mmq_iq_experts(&dev.ptr_row, 2, n_expert, &exi, &exo, &exp_d, &pself, &a_scr,
4883                                 in_pad, n_embd, n_active, n_pairs, n_pairs,
4884                                 m.down_exps.qtype, m.down_exps.row_bytes)?
4885            } else {
4886                let act = e.moe_pairs_gelu_mul(&gate, &up, n_pairs * n_ff_exp)?;
4887                let (aq2, ad2) = e.quantize_q8_1(&act, n_pairs, n_ff_exp)?;
4888                e.moe_pairs_matvec_q8_dec(&dev.ptr_row, 2, &exi, &exo, &exp_d, &pself, &aq2, &ad2,
4889                                          n_ff_exp, n_embd, n_expert, n_active, n_pairs,
4890                                          m.down_exps.qtype, m.down_exps.row_bytes)?
4891            };
4892            let mut moe_out = e.uninit(t * n_embd)?;
4893            e.moe_pairs_scatter(&y_down, &pw, &toff, &tids, &mut moe_out, t, n_embd)?;
4894            return Ok(moe_out);
4895        }
4896
4897        let g_len = m.gate_exps.expert_stride;
4898        let u_len = m.up_exps.expert_stride;
4899        let d_len = m.down_exps.expert_stride;
4900        // Resident dev slabs (fits-VRAM regime): read each expert straight from the device slab
4901        // at ex*stride — zero H2D, SAME qmatvec_view kernel/bytes as the staged path. Staging is
4902        // the spill fallback.
4903        let dev = m.dev_exps.as_ref().filter(|d| !d.gu_il);
4904        let (mut sg, mut su, mut sd) = if dev.is_some() { (None, None, None) } else {
4905            (Some(e.alloc_u8_uninit(g_len)?), Some(e.alloc_u8_uninit(u_len)?), Some(e.alloc_u8_uninit(d_len)?))
4906        };
4907        let mut moe_out = e.zeros(t * n_embd)?;
4908        for tok in 0..t {
4909            let sel = &sel_all[tok * n_used..(tok + 1) * n_used];
4910            let w = &w_all[tok * n_used..(tok + 1) * n_used];
4911            let zt = moe_in.slice(tok * n_embd..(tok + 1) * n_embd);
4912            for (j, &ex) in sel.iter().enumerate() {
4913                let ex = ex as usize;
4914                let gate = match dev {
4915                    Some(d) => e.qmatvec_view(&d.gate, ex * g_len..(ex + 1) * g_len, &zt, 1,
4916                        m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?,
4917                    None => {
4918                        let sg = sg.as_mut().unwrap();
4919                        e.stage_expert(m.gate_exps.expert_bytes(ex), sg, 0)?;
4920                        e.qmatvec_view(sg, 0..g_len, &zt, 1,
4921                            m.gate_exps.in_f, m.gate_exps.out_f, m.gate_exps.qtype, m.gate_exps.row_bytes)?
4922                    }
4923                };
4924                let up = match dev {
4925                    Some(d) => e.qmatvec_view(&d.up, ex * u_len..(ex + 1) * u_len, &zt, 1,
4926                        m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?,
4927                    None => {
4928                        let su = su.as_mut().unwrap();
4929                        e.stage_expert(m.up_exps.expert_bytes(ex), su, 0)?;
4930                        e.qmatvec_view(su, 0..u_len, &zt, 1,
4931                            m.up_exps.in_f, m.up_exps.out_f, m.up_exps.qtype, m.up_exps.row_bytes)?
4932                    }
4933                };
4934                let mut act = e.uninit(n_ff_exp)?;
4935                e.gelu_tanh_mul(&gate, &up, &mut act, n_ff_exp)?;
4936                let actv = act.slice(0..n_ff_exp);
4937                let y = match dev {
4938                    Some(d) => e.qmatvec_view(&d.down, ex * d_len..(ex + 1) * d_len, &actv, 1,
4939                        m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?,
4940                    None => {
4941                        let sd = sd.as_mut().unwrap();
4942                        e.stage_expert(m.down_exps.expert_bytes(ex), sd, 0)?;
4943                        e.qmatvec_view(sd, 0..d_len, &actv, 1,
4944                            m.down_exps.in_f, m.down_exps.out_f, m.down_exps.qtype, m.down_exps.row_bytes)?
4945                    }
4946                };
4947                let mut dst = moe_out.slice_mut(tok * n_embd..(tok + 1) * n_embd);
4948                e.axpy_into(&y, w[j], &mut dst, n_embd)?;
4949            }
4950        }
4951        Ok(moe_out)
4952    }
4953
4954    /// One gemma4 trunk layer (R8): x -> x_next.
4955    fn gemma4_layer(&self, e: &Engine, il: usize, layer: &crate::hybrid::HybridLayer,
4956                    x: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize)
4957                    -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4958        let n_embd = self.cfg.n_embd as usize;
4959        let eps = self.cfg.rms_eps;
4960
4961        let mut h = e.zeros(t * n_embd)?;
4962        e.rms_norm(x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4963        let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
4964        let o = self.gemma4_attn(e, fa, il, &h, pos_d, t)?;
4965        // gemma order: post_attention_norm applies to the ATTENTION OUTPUT, then the residual.
4966        let mut cur = e.zeros(t * n_embd)?;
4967        e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
4968        self.gemma4_layer_tail_add(e, layer, &cur, x, t)
4969    }
4970
4971    /// Everything after the attention output in a gemma4 layer: the residual add (cur + x ->
4972    /// attn_out) FUSED with the three attn_out norms, then shared FFN + router + MoE + combine +
4973    /// layer scale — shared verbatim by the prefill, decode and verify paths.
4974    fn gemma4_layer_tail_add(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
4975                             cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
4976                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4977        Ok(self.gemma4_layer_tail_add_n(e, layer, cur, x, t, None)?.0)
4978    }
4979
4980    /// tail_add with the NEXT layer's attn_norm fused into the closing add+scale (one launch
4981    /// produces both x_next and h_next — the cross-layer fusion). `next_norm` None = last layer.
4982    fn gemma4_layer_tail_add_n(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
4983                               cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
4984                               next_norm: Option<&CudaSlice<f32>>)
4985                               -> Result<(CudaSlice<f32>, Option<CudaSlice<f32>>), Box<dyn std::error::Error>> {
4986        let n_embd = self.cfg.n_embd as usize;
4987        let bits = layer.gemma4.as_ref().unwrap();
4988        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
4989        let mut xn = e.uninit(t * n_embd)?;
4990        match next_norm {
4991            Some(w) => {
4992                let mut hn = e.uninit(t * n_embd)?;
4993                e.add_scale_rms_norm(&sn, &attn_out, bits.layer_scale, w, &mut xn, &mut hn,
4994                                     n_embd, t, self.cfg.rms_eps)?;
4995                Ok((xn, Some(hn)))
4996            }
4997            None => {
4998                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
4999                Ok((xn, None))
5000            }
5001        }
5002    }
5003
5004    /// Tail core: attn_out = cur+x (fused with the 3 norms), both FFN branches, the combine
5005    /// norm — returns (sn, attn_out) for the closing add+scale variants.
5006    fn gemma4_layer_tail_core(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5007                              cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize)
5008                              -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5009        self.gemma4_layer_tail_core_pn(e, layer, cur, x, t, None, false)
5010    }
5011
5012    /// tail_core with an optional PRE-NORM fold (glue-fusion lane): `pre_norm = Some(wa)`
5013    /// means `cur` is the RAW attention output and the dense entry runs
5014    /// rms(cur, wa) + residual-add + ffn_norm as ONE launch (E4B's post-attn norm fold).
5015    /// `defer_post_norm`: dense-arm exit returns RAW f0 (ffn_down output) instead of
5016    /// sn = rms(f0, post_ffw) — the caller fuses the post-norm into its residual emit
5017    /// (rms_pre_add_q8_1, E4B glue wave 5). MoE arm ignores it.
5018    fn gemma4_layer_tail_core_pn(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5019                                 cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5020                                 pre_norm: Option<&CudaSlice<f32>>, defer_post_norm: bool)
5021                                 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5022        let n_embd = self.cfg.n_embd as usize;
5023        let eps = self.cfg.rms_eps;
5024        let bits = layer.gemma4.as_ref().unwrap();
5025
5026        // DENSE gemma4 variants (31B/E4B — no MoE, no parallel branch): attn_out = cur + x
5027        // fused with the single ffn_norm; GELU_PAR ffn; post_ffw_norm.
5028        let Some(mbits) = bits.moe_bits.as_ref() else {
5029            let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
5030            else { panic!("gemma4 dense layer without Dense ffn") };
5031            let mut attn_out = e.uninit(t * n_embd)?;
5032            let mut zsh = e.uninit(t * n_embd)?;
5033            // wave-2: with the pre-norm fold active the entry ALSO emits zsh q8_1 — the
5034            // t=1 fused2 gate/up consume the pair with no standalone quantize launch.
5035            let mut zpair: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5036            match pre_norm {
5037                Some(wa) if t == 1 => {
5038                    zpair = Some(e.rms_pre_add_rms_norm_q8z(cur, wa, x,
5039                                                            bits.ffn_norm.float_data(),
5040                                                            &mut attn_out, &mut zsh,
5041                                                            n_embd, t, eps)?);
5042                }
5043                Some(wa) => e.rms_pre_add_rms_norm(cur, wa, x, bits.ffn_norm.float_data(),
5044                                                   &mut attn_out, &mut zsh, n_embd, t, eps)?,
5045                None => e.add_rms_norm(cur, x, bits.ffn_norm.float_data(), &mut attn_out,
5046                                       &mut zsh, n_embd, t, eps)?,
5047            }
5048            let n_ff = ffn_gate.out_features();
5049            // FFN persistent slab (counter-barrier form) FALSIFIED here 2026-07-14
5050            // (falsification #7, jsonl row): PDL glue already hides the launch boundaries
5051            // it fused; barrier + worst-segment occupancy net −0.3% (31B depth) / −2.3%
5052            // (E4B spec). Down's act dependency is all-to-all, so sentinel sync cannot
5053            // rescue segment C — the megakernel front is closed for the dense tail.
5054            let (gate, up) = if t == 1 {
5055                let (zq, zd) = match zpair {
5056                    Some(p) => p,
5057                    None => e.quantize_q8_1(&zsh, 1, n_embd)?,
5058                };
5059                match e.matmul_q4_fused2(ffn_gate, ffn_up, &zq, &zd)? {
5060                    Some(p) => p,
5061                    None => (e.matmul_pre(ffn_gate, &zq, &zd, &zsh, 1)?,
5062                             e.matmul_pre(ffn_up, &zq, &zd, &zsh, 1)?),
5063                }
5064            } else {
5065                // BATCHED FUSED2 (DEFAULT ON 2026-07-13, MEMRA_F2B=0 seam): one segmented
5066                // launch for the verify's gate+up — the up segment's blocks fill SMs as
5067                // the gate segment drains (the launch-tail mechanism behind the b-tier
5068                // plateau; first positive after six falsified in-kernel variants).
5069                static F2B: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5070                let f2b = *F2B.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
5071                let fused = if f2b {
5072                    let (zq, zd) = e.quantize_q8_1(&zsh, t, n_embd)?;
5073                    e.matmul_q4_fused2_batched(ffn_gate, ffn_up, &zq, &zd, t)?
5074                } else { None };
5075                match fused {
5076                    Some(p) => p,
5077                    None => {
5078                        // quantize-once window: gate/up share `zsh` (borrowed across the pair).
5079                        e.mmq_act_begin();
5080                        (e.matmul(ffn_gate, &zsh, t)?, e.matmul(ffn_up, &zsh, t)?)
5081                    }
5082                }
5083            };
5084            let mut act = e.uninit(t * n_ff)?;
5085            // act quantize folds into the GELU epilogue (bit-identical q8_1 rounding);
5086            // ffn_down rides matmul_pre — one quantize launch fewer per layer.
5087            let f0 = if e.uses_q8_1_fast(ffn_down) {
5088                let upv = e.view(&up, t * n_ff);
5089                let up_all = upv.slice(0..t * n_ff);
5090                let (aq, ad) = e.gelu_tanh_mul_q8_1(&gate, &up_all, &mut act, n_ff, t)?;
5091                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5092            } else {
5093                e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5094                e.matmul(ffn_down, &act, t)?
5095            };
5096            if defer_post_norm { return Ok((f0, attn_out)); }
5097            let mut sn = e.uninit(t * n_embd)?;
5098            e.rms_norm(&f0, bits.post_ffw_norm.float_data(), &mut sn, n_embd, t, eps)?;
5099            return Ok((sn, attn_out));
5100        };
5101
5102        assert!(pre_norm.is_none(), "pre-norm fold is dense-entry only");
5103        // MoE variant (26B): attn_out = cur + x fused with the three attn_out norms
5104        // (ffn_norm + router-scale + pre_ffw_norm_2): ONE launch, chains verbatim. At small t
5105        // (decode + verify) the zsh and moe_in outputs are EMITTED q8_1 (both consumers are
5106        // quantized matmuls — two quantize launches + two f32 round-trips fold away).
5107        let mut attn_out = e.uninit(t * n_embd)?;
5108        let mut router_in = e.uninit(t * n_embd)?;
5109        let fast_moe = match &layer.ffn {
5110            crate::hybrid::Ffn::Moe(m) => m.dev_exps.as_ref().is_some_and(|d| !d.gu_il)
5111                && expert_dp4a_supported(m.gate_exps.qtype)
5112                && expert_dp4a_supported(m.up_exps.qtype)
5113                && expert_dp4a_supported(m.down_exps.qtype)
5114                && std::env::var("MEMRA_GEMMA_MOE_FAST").as_deref() != Ok("0"),
5115            _ => false,
5116        };
5117        let q8z = t < PRIME_MIN_T && fast_moe;
5118        let (zsh_f32, zsh_q8, moe_q8) = if q8z {
5119            let (z0, m2) = e.add_rms_norm3_q8z(cur, x, bits.ffn_norm.float_data(),
5120                                               &mbits.router_scale_pre,
5121                                               mbits.pre_ffw_norm_2.float_data(),
5122                                               &mut attn_out, &mut router_in, n_embd, t, eps)?;
5123            (None, Some(z0), Some(m2))
5124        } else {
5125            let mut zsh = e.uninit(t * n_embd)?;
5126            let mut moe_in = e.uninit(t * n_embd)?;
5127            e.add_rms_norm3(cur, x, bits.ffn_norm.float_data(), &mbits.router_scale_pre,
5128                            mbits.pre_ffw_norm_2.float_data(), &mut attn_out, &mut zsh,
5129                            &mut router_in, &mut moe_in, n_embd, t, eps)?;
5130            (Some((zsh, moe_in)), None, None)
5131        };
5132        let attn_out2 = attn_out;
5133        #[allow(unused_variables)]
5134        let attn_out = &attn_out2;
5135        let n_ff = mbits.shared_gate.out_features();
5136        let (gate, up) = if let Some((zq, zd)) = zsh_q8.as_ref() {
5137            if t == 1 {
5138                match e.matmul_q4_fused2(&mbits.shared_gate, &mbits.shared_up, zq, zd)? {
5139                    Some(p) => p,
5140                    None => {
5141                        let h0 = e.zeros(0)?;
5142                        (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, 1)?,
5143                         e.matmul_pre(&mbits.shared_up, zq, zd, &h0, 1)?)
5144                    }
5145                }
5146            } else {
5147                // verify t 2..15: the batched mmvq twins consume the pre-quantized pair.
5148                let h0 = e.zeros(0)?;
5149                (e.matmul_pre(&mbits.shared_gate, zq, zd, &h0, t)?,
5150                 e.matmul_pre(&mbits.shared_up, zq, zd, &h0, t)?)
5151            }
5152        } else {
5153            let (zsh, _) = zsh_f32.as_ref().unwrap();
5154            (e.matmul(&mbits.shared_gate, zsh, t)?, e.matmul(&mbits.shared_up, zsh, t)?)
5155        };
5156        let mut act = e.uninit(t * n_ff)?;
5157        e.gelu_tanh_mul(&gate, &up, &mut act, t * n_ff)?;
5158        let mlp0 = e.matmul(&mbits.shared_down, &act, t)?;
5159        let crate::hybrid::Ffn::Moe(m) = &layer.ffn else { panic!("gemma4 layer not MoE") };
5160        let moe0 = match (&moe_q8, &zsh_f32) {
5161            (Some(mq), _) => self.gemma4_moe_q8(e, m, mbits, mq, &router_in, t)?,
5162            (None, Some((_, moe_in))) => self.gemma4_moe(e, m, mbits, moe_in, &router_in, t)?,
5163            _ => unreachable!(),
5164        };
5165        // post_ffw_norm_1(mlp0) + post_ffw_norm_2(moe0): one fused launch, per-row verbatim.
5166        let mut mlp = e.uninit(t * n_embd)?;
5167        let mut moe = e.uninit(t * n_embd)?;
5168        e.rms_norm2x(&mlp0, &moe0, mbits.post_ffw_norm_1.float_data(),
5169                     mbits.post_ffw_norm_2.float_data(), &mut mlp, &mut moe, n_embd, t, eps)?;
5170
5171        // combine: rms_norm(mlp + moe, post_ffw_norm) + attn_out, then the layer output scalar.
5172        // add+norm fused (add_rms_norm == add then rms_norm, kernel-check-pinned identity).
5173        let mut sum = e.uninit(t * n_embd)?;
5174        let mut sn = e.uninit(t * n_embd)?;
5175        e.add_rms_norm(&mlp, &moe, bits.post_ffw_norm.float_data(), &mut sum, &mut sn,
5176                       n_embd, t, eps)?;
5177        Ok((sn, attn_out2))
5178    }
5179
5180    /// tail_add with the next attn_norm emitted PRE-QUANTIZED q8_1 (decode/verify loops).
5181    fn gemma4_layer_tail_add_nq(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5182                                cur: &CudaSlice<f32>, x: &CudaSlice<f32>, t: usize,
5183                                next_norm: Option<&CudaSlice<f32>>)
5184                                -> Result<(CudaSlice<f32>, Option<(CudaSlice<i8>, CudaSlice<f32>)>), Box<dyn std::error::Error>> {
5185        let n_embd = self.cfg.n_embd as usize;
5186        let bits = layer.gemma4.as_ref().unwrap();
5187        let (sn, attn_out) = self.gemma4_layer_tail_core(e, layer, cur, x, t)?;
5188        let mut xn = e.uninit(t * n_embd)?;
5189        match next_norm {
5190            Some(w) => {
5191                let pair = e.add_scale_rms_norm_q8_1(&sn, &attn_out, bits.layer_scale, w, &mut xn,
5192                                                     n_embd, t, self.cfg.rms_eps)?;
5193                Ok((xn, Some(pair)))
5194            }
5195            None => {
5196                e.add_scale(&sn, &attn_out, bits.layer_scale, &mut xn, t * n_embd)?;
5197                Ok((xn, None))
5198            }
5199        }
5200    }
5201
5202    /// gemma4 prefill: `last_only` = forward_last semantics (lm_head on the final row only).
5203    /// R4: final logits softcapped 30*tanh(l/30) on host (monotonic — argmax unaffected).
5204    fn gemma4_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
5205                      -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5206        // E4B routes to its own forward regardless of the caller's entry point (forward /
5207        // forward_last / prime paths all funnel here for gemma4).
5208        if self.is_gemma4_e4b() { return self.gemma4_e4b_forward(e, tokens, last_only); }
5209        let n_embd = self.cfg.n_embd as usize;
5210        let t = tokens.len();
5211        let pos: Vec<i32> = (0..t as i32).collect();
5212        let pos_d = e.htod_i32(&pos)?;
5213
5214        let mut x = self.embed(e, tokens)?;
5215        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5216        // MEMRA_GEMMA_PROBE=1: per-layer trunk stats (host rms + max of the LAST token row) —
5217        // the bring-up bisect vs llama-eval-callback node stats.
5218        let probe = std::env::var("MEMRA_GEMMA_PROBE").is_ok();
5219        let stat = |e: &Engine, x: &CudaSlice<f32>, tag: &str| -> Result<(), Box<dyn std::error::Error>> {
5220            let h = e.dtoh(x)?;
5221            let bad = h.iter().filter(|v| !v.is_finite()).count();
5222            let mx = h.iter().filter(|v| v.is_finite()).fold(0.0f32, |m, v| m.max(v.abs()));
5223            eprintln!("[gemma-probe] {tag}: tok0_first3={:?} bad={bad} max={mx:.3e}", &h[..3]);
5224            Ok(())
5225        };
5226        if probe { stat(e, &x, "embed")?; }
5227        for (il, layer) in self.layers.iter().enumerate() {
5228            x = self.gemma4_layer(e, il, layer, &x, &pos_d, t)?;
5229            if probe { stat(e, &x, &format!("L{il}"))?; }
5230        }
5231        let mut hn = e.zeros(t * n_embd)?;
5232        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, self.cfg.rms_eps)?;
5233        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5234        let n_vocab = self.output.out_features();
5235        let logits = if last_only {
5236            let hv = e.view(&hn, t * n_embd);
5237            let last_row = hv.slice((t - 1) * n_embd..t * n_embd);
5238            let mut hlast = e.zeros(n_embd)?;
5239            e.copy_view_into(&mut hlast, 0, &last_row, n_embd)?;
5240            let mut ld = e.matmul(&self.output, &hlast, 1)?;
5241            e.softcap(&mut ld, cap, n_vocab)?;
5242            self.gemma4_suppress(e, &mut ld, 1)?;
5243            e.dtoh(&ld)?
5244        } else {
5245            let mut ld = e.matmul(&self.output, &hn, t)?;
5246            e.softcap(&mut ld, cap, t * n_vocab)?;
5247            self.gemma4_suppress(e, &mut ld, t)?;
5248            e.dtoh(&ld)?
5249        };
5250        Ok(logits)
5251    }
5252
5253    /// gemma4 BATCHED PROMPT PRIME v0 (fresh cache only): the prefill graph over the whole
5254    /// prompt with each layer's post-rope K / weightless-normed V appended into the quantized
5255    /// KV cache (decode-append row math). Returns (last-row softcapped logits, h_seed = last
5256    /// pre-output_norm hidden, hiddens = full pre-output_norm stack [T, n_embd]).
5257    pub(crate) fn gemma4_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
5258                               -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5259        assert_eq!(cache.pos, 0, "gemma4 prime v0 is fresh-prompt only");
5260        let n_embd = self.cfg.n_embd as usize;
5261        let eps = self.cfg.rms_eps;
5262        let t = tokens.len();
5263        let pos: Vec<i32> = (0..t as i32).collect();
5264        let pos_d = e.htod_i32(&pos)?;
5265        let mut x = self.embed(e, tokens)?;
5266        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
5267        for (il, layer) in self.layers.iter().enumerate() {
5268            let mut h = e.zeros(t * n_embd)?;
5269            e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5270            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer not full-attn") };
5271            let o = self.gemma4_attn_prime(e, fa, il, &h, &pos_d, t, Some(cache))?;
5272            let mut cur = e.zeros(t * n_embd)?;
5273            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
5274            x = self.gemma4_layer_tail_add(e, layer, &cur, &x, t)?;
5275            self.dflash_tap(e, cache, il, &x, t)?;
5276        }
5277        cache.pos += t;
5278        let hiddens = e.clone_dtod(&x)?;
5279        let xv = e.view(&x, t * n_embd);
5280        let last_row = xv.slice((t - 1) * n_embd..t * n_embd);
5281        let mut h_seed = e.zeros(n_embd)?;
5282        e.copy_view_into(&mut h_seed, 0, &last_row, n_embd)?;
5283        let mut hn = e.uninit(n_embd)?;
5284        e.rms_norm(&h_seed, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5285        let mut ld = e.matmul(&self.output, &hn, 1)?;
5286        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
5287        e.softcap(&mut ld, cap, self.output.out_features())?;
5288        self.gemma4_suppress(e, &mut ld, 1)?;
5289        let logits = e.dtoh(&ld)?;
5290        Ok((logits, h_seed, hiddens))
5291    }
5292
5293    /// gemma4 T=1 decode attention: per-layer geometry, quantized-KV append + fa_decode
5294    /// (vec kernels at hd 256, generic scalar at the globals' hd 512), weightless V-norm,
5295    /// dual rope, scale 1.0. Takes the attn-normed input PRE-QUANTIZED (the cross-layer
5296    /// fused norm emits q8 directly — the f32 h never materializes).
5297    fn gemma4_decode_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5298                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
5299                          pos_d: &CudaSlice<i32>, cache: &mut Cache)
5300                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5301        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
5302        let eps = self.cfg.rms_eps;
5303        let aux = self.gemma4_aux.as_ref().unwrap();
5304        let (hq, hdq) = (hq, hdq);
5305        let h0 = e.zeros(0)?;
5306        let h = &h0;
5307        let (q0, k0, v0) = if swa {
5308            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
5309                Some(t3) => t3,
5310                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5311                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
5312                         e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?),
5313            }
5314        } else {
5315            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, &hq, &hdq)? {
5316                Some(p) => p,
5317                None => (e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
5318                         e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?),
5319            };
5320            let v0 = e.clone_dtod(&k0)?;
5321            (q0, k0, v0)
5322        };
5323        let mut q = e.uninit(nh * hd)?;
5324        let mut k = e.uninit(nkv * hd)?;
5325        let mut v = e.uninit(nkv * hd)?;
5326        // E4B wave-3 fold, m=1 completion (2026-07-23): the rows arms took this fold at
5327        // 550fcfa5; the decode-step trio kept 2 launches/layer it doesn't need.
5328        let ff = if swa { None } else {
5329            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5330        };
5331        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5332                            &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
5333                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
5334        let kvl = cache.kv[il].as_mut().unwrap();
5335        e.append_kv_quantized(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len,
5336                              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()))?;
5337        kvl.len += 1;
5338        // R6 decode: SWA layers attend only the last `sliding_window` keys — a token-aligned
5339        // VIEW OFFSET into the quantized cache (keys carry absolute rope; the mask is purely
5340        // positional). Globals attend the full history.
5341        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5342        let mut attn = e.uninit(nh * hd)?;
5343        // global layers: SAME hd512 rows twin as verify with t=1 (parity law).
5344        if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
5345            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5346            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5347            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5348            // device-len: eager syncs the counter (async arg-store), dc keeps it live.
5349            let base = kvl.len as i32;
5350            e.i32_set_k(&mut kvl.len_d, base)?;
5351            e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1, scale,
5352                             kvl.k_tok_bytes, kvl.v_tok_bytes, Some((&kvl.len_d, -1)), false,
5353                             false, None)?;
5354            return Ok(e.matmul(&fa.wo, &attn, 1)?);
5355        }
5356        // windowed regime: SAME rows_w kernel as verify with t=1 (parity law — see verify_attn).
5357        if swa && kvl.len > win && hd == 256
5358            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5359            let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5360            let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5361            let base = kvl.len as i32;
5362            e.i32_set_k(&mut kvl.len_d, base)?;
5363            e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1, 1, scale,
5364                               win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
5365            return Ok(e.matmul(&fa.wo, &attn, 1)?);
5366        }
5367        let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) } else { (0, kvl.len) };
5368        let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
5369                                     (off_tok + t_kv) * kvl.k_tok_bytes);
5370        let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
5371                                     (off_tok + t_kv) * kvl.v_tok_bytes);
5372        e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
5373                    kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
5374        Ok(e.matmul(&fa.wo, &attn, 1)?)
5375    }
5376
5377    /// gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in
5378    /// device counters; ZERO varying host kernel args. `cap_bucket_max` = Some(bucket) for graph
5379    /// capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for
5380    /// the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target =
5381    /// gemma4_decode_step_h's token stream). V1 scope: t_kv <= sliding_window (no window views
5382    /// in-graph; the driver gates).
5383    #[allow(clippy::too_many_arguments)]
5384    pub fn gemma4_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
5385                                 pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5386                                 embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5387                                 n_vocab: usize, cap_bucket_max: Option<(usize, usize)>)
5388                                 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
5389        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
5390        self.gemma4_decode_step_dc_into(e, token_d, pos_d, embd_gpu, embd_qt, embd_rb, cache,
5391                                        n_vocab, cap_bucket_max, &mut tok_out)?;
5392        Ok(tok_out)
5393    }
5394
5395    /// CAPTURE body: argmax lands in the PERSISTENT `tok_out` (same buffer = same address on
5396    /// every replay; pass `token_d` itself for the self-feeding graph loop).
5397    #[allow(clippy::too_many_arguments)]
5398    pub fn gemma4_decode_step_dc_into(&self, e: &Engine, token_d: &CudaSlice<u32>,
5399                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5400                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5401                                      n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5402                                      tok_out: &mut CudaSlice<u32>)
5403                                      -> Result<(), Box<dyn std::error::Error>> {
5404        let n_embd = self.cfg.n_embd as usize;
5405        let eps = self.cfg.rms_eps;
5406        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
5407        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
5408        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5409        let n_layers = self.layers.len();
5410        for (il, layer) in self.layers.iter().enumerate() {
5411            let (hq, hdq) = match h_carry.take() {
5412                Some(p) => p,
5413                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
5414            };
5415            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5416            let o = self.gemma4_decode_attn_dc(e, fa, il, &hq, &hdq, pos_d, cache, cap_bucket_max)?;
5417            let mut cur = e.uninit(n_embd)?;
5418            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
5419            let next_norm = if il + 1 < n_layers {
5420                Some(self.layers[il + 1].attn_norm.float_data())
5421            } else { None };
5422            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
5423            x = xn;
5424            h_carry = hn;
5425        }
5426        let mut hn = e.uninit(n_embd)?;
5427        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
5428        let mut logits = e.matmul(&self.output, &hn, 1)?;
5429        self.gemma4_suppress(e, &mut logits, 1)?;   // cap skipped (monotonic); the mask is not
5430        e.argmax_token_device_into(&logits, tok_out, n_vocab)?;
5431        e.inc_seqlen(pos_d)?;
5432        if cap_bucket_max.is_none() { cache.pos += 1; }
5433        Ok(())
5434    }
5435
5436    /// Persistent transient slots for the ALLOC-FREE captured dc step (the graph door):
5437    /// every buffer the step produces per token lives here, allocated ONCE pre-capture, so
5438    /// the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax,
5439    /// osrt 2026-07-23). Sized for the model's max per-layer shapes.
5440
5441    /// Build the slot set (call OUTSIDE any capture).
5442    pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn std::error::Error>> {
5443        let n_embd = self.cfg.n_embd as usize;
5444        let n_vocab = self.output.out_features();
5445        let n_layers = self.layers.len();
5446        let (mut qmax, mut kvmax, mut ffmax) = (0usize, 0usize, 0usize);
5447        for il in 0..n_layers {
5448            let (hd, nkv, nh, _b, _s, _w) = self.gemma4_geom(il);
5449            qmax = qmax.max(nh * hd);
5450            kvmax = kvmax.max(nkv * hd);
5451            if let crate::hybrid::Ffn::Dense { ffn_gate, .. } = &self.layers[il].ffn {
5452                ffmax = ffmax.max(ffn_gate.out_features());
5453            }
5454        }
5455        Ok(G4DcSlots {
5456            x: e.uninit(n_embd)?, xn: e.uninit(n_embd)?, cur: e.uninit(n_embd)?,
5457            hq: e.alloc_i8_uninit(n_embd)?, hd_: e.uninit(n_embd / 32)?,
5458            q0: e.uninit(qmax)?, k0: e.uninit(kvmax)?, v0: e.uninit(kvmax)?,
5459            q: e.uninit(qmax)?, k: e.uninit(kvmax)?, v: e.uninit(kvmax)?,
5460            attn: e.uninit(qmax)?, o: e.uninit(n_embd)?,
5461            attn_out: e.uninit(n_embd)?, zsh: e.uninit(n_embd)?,
5462            // zq/zd feed the wo matvec (nh*hd rows — 4096 on hd512 globals > n_embd),
5463            // the ffn entry (n_embd) and the lm_head (n_embd): size for the max.
5464            zq: e.alloc_i8_uninit(n_embd.max(qmax))?, zd: e.uninit(n_embd.max(qmax) / 32)?,
5465            gate: e.uninit(ffmax)?, up: e.uninit(ffmax)?,
5466            act: e.uninit(ffmax)?, actq: e.alloc_i8_uninit(ffmax)?, actd: e.uninit(ffmax / 32)?,
5467            f0: e.uninit(n_embd)?, sn: e.uninit(n_embd)?,
5468            hn: e.uninit(n_embd)?, logits: e.uninit(n_vocab)?,
5469        })
5470    }
5471
5472    /// m=1 pre-quantized matvec into a slot — mirrors matmul_pre's m=1 mmvq route exactly
5473    /// (rp4-mirror bytes; mmvq_supports guaranteed for gemma4 q4_0/q6_K).
5474    fn g4_matvec_m1_into(&self, e: &Engine, w: &crate::model::GpuTensor,
5475                         aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y: &mut CudaSlice<f32>)
5476                         -> Result<(), Box<dyn std::error::Error>> {
5477        use crate::model::GpuTensor;
5478        let (bytes, qtype, row_bytes, scale, rp) = match w {
5479            GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
5480                (bytes, *qtype, *row_bytes, *scale, *rp),
5481            _ => return Err("g4_matvec_m1_into: non-quant tensor".into()),
5482        };
5483        let (mbytes, mrp) = match w {
5484            GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5485            _ => (bytes, rp),
5486        };
5487        e.qmatvec_mmvq_into(mbytes, aq, ad, 1, w.in_features(), w.out_features(),
5488                            qtype, row_bytes, scale, mrp, y)
5489    }
5490
5491    /// ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of
5492    /// `gemma4_decode_step_dc_into` at t=1 with every transient slot-fed. Dense gemma4 only
5493    /// (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).
5494    #[allow(clippy::too_many_arguments)]
5495    pub fn gemma4_decode_step_dc_slotted(&self, e: &Engine, token_d: &CudaSlice<u32>,
5496                                         pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
5497                                         embd_qt: i32, embd_rb: usize, cache: &mut Cache,
5498                                         n_vocab: usize, cap_bucket_max: Option<(usize, usize)>,
5499                                         sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>,
5500                                         ring: Option<(&mut CudaSlice<u32>, usize)>)
5501                                         -> Result<(), Box<dyn std::error::Error>> {
5502        let n_embd = self.cfg.n_embd as usize;
5503        let eps = self.cfg.rms_eps;
5504        e.embed_gather_device_into(embd_gpu, token_d, &mut sl.x, n_embd, embd_qt, embd_rb)?;
5505        e.scale_inplace(&mut sl.x, (n_embd as f32).sqrt(), n_embd)?;
5506        let n_layers = self.layers.len();
5507        let mut has_carry = false;
5508        for il in 0..n_layers {
5509            if !has_carry {
5510                e.rms_norm_q8_1_into(&sl.x, self.layers[il].attn_norm.float_data(), n_embd, 1,
5511                                     eps, &mut sl.hq, &mut sl.hd_)?;
5512            }
5513            has_carry = true;
5514            let layer = &self.layers[il];
5515            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
5516            self.gemma4_decode_attn_dc_slotted(e, fa, il, pos_d, cache, cap_bucket_max, sl)?;
5517            e.rms_norm(&sl.o, layer.post_attn_norm.float_data(), &mut sl.cur, n_embd, 1, eps)?;
5518            let next_norm = if il + 1 < n_layers {
5519                Some(self.layers[il + 1].attn_norm.float_data())
5520            } else { None };
5521            self.gemma4_layer_tail_slotted(e, layer, next_norm, sl)?;
5522            std::mem::swap(&mut sl.x, &mut sl.xn);
5523        }
5524        e.rms_norm(&sl.x, self.output_norm.float_data(), &mut sl.hn, n_embd, 1, eps)?;
5525        e.quantize_q8_1_into(&sl.hn, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
5526        // lm_head via the same m=1 mmvq route as matmul (q6_K on the gemma4 family).
5527        {
5528            let (zq, zd) = (&sl.zq, &sl.zd);
5529            let zq = unsafe { &*(zq as *const CudaSlice<i8>) };
5530            let zd = unsafe { &*(zd as *const CudaSlice<f32>) };
5531            self.g4_matvec_m1_into(e, &self.output, zq, zd, &mut sl.logits)?;
5532        }
5533        self.gemma4_suppress(e, &mut sl.logits, 1)?;
5534        e.argmax_token_device_into(&sl.logits, tok_out, n_vocab)?;
5535        if let Some((ring, base)) = ring {
5536            // in-graph token parking: slot = (pos - base) % ring.len(); the door drains
5537            // with ONE sync per chunk instead of a per-token dtoh (the launch-serialization
5538            // fix — llama's 885us/launch overlaps GPU work because nothing waits per token).
5539            e.plain_tok_ring(tok_out, pos_d, base, ring)?;
5540        }
5541        e.inc_seqlen(pos_d)?;
5542        if cap_bucket_max.is_none() { cache.pos += 1; }
5543        Ok(())
5544    }
5545
5546    /// Slot-fed dc attention: mirrors gemma4_decode_attn_dc's CAPTURE arm kernel-for-kernel.
5547    #[allow(clippy::too_many_arguments)]
5548    fn gemma4_decode_attn_dc_slotted(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer,
5549                                     il: usize, pos_d: &CudaSlice<i32>, cache: &mut Cache,
5550                                     cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots)
5551                                     -> Result<(), Box<dyn std::error::Error>> {
5552        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
5553        let eps = self.cfg.rms_eps;
5554        let aux = self.gemma4_aux.as_ref().unwrap();
5555        {
5556            let hq = unsafe { &*(&sl.hq as *const CudaSlice<i8>) };
5557            let hdq = unsafe { &*(&sl.hd_ as *const CudaSlice<f32>) };
5558            if swa {
5559                if !e.matmul_q4_fused3_into(&fa.wq, &fa.wk, &fa.wv, hq, hdq,
5560                                            &mut sl.q0, &mut sl.k0, &mut sl.v0)? {
5561                    return Err("slotted step: fused3 unavailable (non-uniform trunk)".into());
5562                }
5563            } else {
5564                if !e.matmul_q4_fused2_into(&fa.wq, &fa.wk, hq, hdq, &mut sl.q0, &mut sl.k0)? {
5565                    return Err("slotted step: fused2 unavailable".into());
5566                }
5567                let k0r = unsafe { &*(&sl.k0 as *const CudaSlice<f32>) };
5568                e.copy_into(&mut sl.v0, 0, k0r, nkv * hd)?;
5569            }
5570        }
5571        // E4B wave-3 fold, m=1 completion (2026-07-23) — MUST mirror the dc_into arm
5572        // kernel-for-kernel (graph stream-identity gate).
5573        let ff = if swa { None } else {
5574            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5575        };
5576        let kvl = cache.kv[il].as_mut().unwrap();
5577        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
5578        if crate::Engine::qkv_append_on() {
5579            // append fold (2026-07-23): mirrors dc_into.
5580            e.rms_norm_qkv_rope_append_dc(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(),
5581                fa.k_norm.float_data(), &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
5582                pos_d, nh, nkv, base, 1.0, ff, eps,
5583                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
5584        } else {
5585            e.rms_norm_qkv_rope(&sl.q0, &sl.k0, &sl.v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5586                                &aux.ones, &mut sl.q, &mut sl.k, &mut sl.v, hd, nh, nkv,
5587                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
5588            e.append_kv_quantized_dc(&sl.k, &sl.v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
5589                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
5590                                     kv_fp8)?;
5591        }
5592        e.inc_seqlen(&mut kvl.len_d)?;
5593        let (b_swa, b_glob) = cap_bucket_max.expect("slotted step is capture-only");
5594        let k_view = e.view_u8(&kvl.k, kvl.k.len());
5595        let v_view = e.view_u8(&kvl.v, kvl.v.len());
5596        let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
5597        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5598        // combine-q8 emit (wave-5b m=1 port): the rows arms quantize inside the combine —
5599        // the standalone quantize launch runs only on the non-rows fallback. MUST mirror
5600        // the dc_into arm branch-for-branch (stream gate).
5601        let mut fa_q8 = false;
5602        if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
5603            e.fa_decode_rows(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, b_glob - 1,
5604                             1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
5605                             Some((&kvl.len_d, -1)), false, false,
5606                             Some((&mut sl.zq, &mut sl.zd)))?;
5607            fa_q8 = true;
5608        } else if swa && b_swa > win && hd == 256 && rows_on {
5609            e.fa_decode_rows_w(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv,
5610                               &kvl.len_d, -1, 1, scale, win,
5611                               kvl.k_tok_bytes, kvl.v_tok_bytes,
5612                               Some((&mut sl.zq, &mut sl.zd)))?;
5613            fa_q8 = true;
5614        } else {
5615            let b = if swa { b_swa } else { b_glob };
5616            e.fa_decode_dc(&sl.q, &k_view, &v_view, &mut sl.attn, hd, nh, nkv, &kvl.len_d, b,
5617                           scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
5618                           swa && crate::Engine::wkv_on())?;
5619        }
5620        if !fa_q8 {
5621            let aq = unsafe { &*(&sl.attn as *const CudaSlice<f32>) };
5622            e.quantize_q8_1_into(aq, 1, nh * hd, &mut sl.zq, &mut sl.zd)?;
5623        }
5624        {
5625            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
5626            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
5627            self.g4_matvec_m1_into(e, &fa.wo, zq, zd, &mut sl.o)?;
5628        }
5629        Ok(())
5630    }
5631
5632    /// Slot-fed dense layer tail: mirrors gemma4_layer_tail_core (t=1, no pre-norm fold) +
5633    /// tail_add_nq kernel-for-kernel; the next layer's (hq, hd_) carry lands in the slots.
5634    fn gemma4_layer_tail_slotted(&self, e: &Engine, layer: &crate::hybrid::HybridLayer,
5635                                 next_norm: Option<&CudaSlice<f32>>, sl: &mut G4DcSlots)
5636                                 -> Result<(), Box<dyn std::error::Error>> {
5637        let n_embd = self.cfg.n_embd as usize;
5638        let eps = self.cfg.rms_eps;
5639        let bits = layer.gemma4.as_ref().unwrap();
5640        let crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } = &layer.ffn
5641        else { return Err("slotted tail: dense ffn only".into()) };
5642        e.add_rms_norm(&sl.cur, &sl.x, bits.ffn_norm.float_data(), &mut sl.attn_out,
5643                       &mut sl.zsh, n_embd, 1, eps)?;
5644        let n_ff = ffn_gate.out_features();
5645        {
5646            let zshr = unsafe { &*(&sl.zsh as *const CudaSlice<f32>) };
5647            e.quantize_q8_1_into(zshr, 1, n_embd, &mut sl.zq, &mut sl.zd)?;
5648        }
5649        {
5650            let zq = unsafe { &*(&sl.zq as *const CudaSlice<i8>) };
5651            let zd = unsafe { &*(&sl.zd as *const CudaSlice<f32>) };
5652            if !e.matmul_q4_fused2_into(ffn_gate, ffn_up, zq, zd, &mut sl.gate, &mut sl.up)? {
5653                return Err("slotted tail: ffn fused2 unavailable".into());
5654            }
5655        }
5656        debug_assert!(e.uses_q8_1_fast(ffn_down));
5657        {
5658            let upr = unsafe { &*(&sl.up as *const CudaSlice<f32>) };
5659            let upv = e.view(upr, n_ff);
5660            let up_all = upv.slice(0..n_ff);
5661            let gr = unsafe { &*(&sl.gate as *const CudaSlice<f32>) };
5662            e.gelu_tanh_mul_q8_1_into(gr, &up_all, &mut sl.act, n_ff, 1,
5663                                      &mut sl.actq, &mut sl.actd)?;
5664        }
5665        {
5666            let aq = unsafe { &*(&sl.actq as *const CudaSlice<i8>) };
5667            let ad = unsafe { &*(&sl.actd as *const CudaSlice<f32>) };
5668            self.g4_matvec_m1_into(e, ffn_down, aq, ad, &mut sl.f0)?;
5669        }
5670        e.rms_norm(&sl.f0, bits.post_ffw_norm.float_data(), &mut sl.sn, n_embd, 1, eps)?;
5671        match next_norm {
5672            Some(w) => {
5673                e.add_scale_rms_norm_q8_1_into(&sl.sn, &sl.attn_out, bits.layer_scale, w,
5674                                               &mut sl.xn, n_embd, 1, eps,
5675                                               &mut sl.hq, &mut sl.hd_)?;
5676            }
5677            None => {
5678                e.add_scale(&sl.sn, &sl.attn_out, bits.layer_scale, &mut sl.xn, n_embd)?;
5679            }
5680        }
5681        Ok(())
5682    }
5683
5684    /// dc attention: same math as gemma4_decode_attn, KV slot/lengths from device counters.
5685    #[allow(clippy::too_many_arguments)]
5686    fn gemma4_decode_attn_dc(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
5687                             hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
5688                             pos_d: &CudaSlice<i32>, cache: &mut Cache,
5689                             cap_bucket_max: Option<(usize, usize)>)
5690                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5691        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
5692        let eps = self.cfg.rms_eps;
5693        let aux = self.gemma4_aux.as_ref().unwrap();
5694        let (q0, k0, v0) = if swa {
5695            match e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)? {
5696                Some(t3) => t3,
5697                None => {
5698                    let h0 = e.zeros(0)?;
5699                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
5700                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?,
5701                     e.matmul_pre(&fa.wv, hq, hdq, &h0, 1)?)
5702                }
5703            }
5704        } else {
5705            let (q0, k0) = match e.matmul_q4_fused2(&fa.wq, &fa.wk, hq, hdq)? {
5706                Some(p) => p,
5707                None => {
5708                    let h0 = e.zeros(0)?;
5709                    (e.matmul_pre(&fa.wq, hq, hdq, &h0, 1)?,
5710                     e.matmul_pre(&fa.wk, hq, hdq, &h0, 1)?)
5711                }
5712            };
5713            let v0 = e.clone_dtod(&k0)?;
5714            (q0, k0, v0)
5715        };
5716        let mut q = e.uninit(nh * hd)?;
5717        let mut k = e.uninit(nkv * hd)?;
5718        let mut v = e.uninit(nkv * hd)?;
5719        // E4B wave-3 fold, m=1 completion (2026-07-23): mirrored by the slotted arm.
5720        let ff = if swa { None } else {
5721            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
5722        };
5723        let kvl = cache.kv[il].as_mut().unwrap();
5724        let kv_fp8 = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
5725        if crate::Engine::qkv_append_on() {
5726            // append fold (2026-07-23): norm+rope+cache-append in ONE launch.
5727            e.rms_norm_qkv_rope_append_dc(&q0, &k0, &v0, fa.q_norm.float_data(),
5728                fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
5729                pos_d, nh, nkv, base, 1.0, ff, eps,
5730                &mut kvl.k, &mut kvl.v, &kvl.len_d, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
5731        } else {
5732            e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
5733                                &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
5734                                pos_d, nh, nkv, base, 1.0, ff, eps)?;
5735            e.append_kv_quantized_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d,
5736                                     kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes, kv_fp8)?;
5737        }
5738        e.inc_seqlen(&mut kvl.len_d)?;
5739        let mut attn = e.uninit(nh * hd)?;
5740        // combine-q8 emit carrier (wave-5b m=1 port): rows arms fill this; the tail then
5741        // rides g4_matvec_m1_into instead of matmul's internal quantize.
5742        let mut fa_q8: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
5743        // Weight prefetch NOT wired here (26B/31B/qwen probes 2026-07-13: 26B flat
5744        // 196.6/196.0 vs 196.4/196.1 — MoE ffn dilutes wo; 31B −0.2% — dense decode sits
5745        // at the DRAM wall, no idle window to front-load into). E4B keeps the arm
5746        // (gemma4_e4b_attn, +0.65% valid window).
5747        match cap_bucket_max {
5748            None => {
5749                // dc-EAGER: host knows the length — R6 window views EXACTLY like the eager
5750                // decode (SWA layers attend the last `sliding_window` keys); the device
5751                // counters carry only the append slot + the graph seam.
5752                kvl.len += 1;
5753                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5754                if !swa && hd == 512 && kvl.len >= crate::fa512_min_tkv()
5755                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5756                    // global layers: SAME hd512 rows twin as verify, t=1 (parity law).
5757                    // dc: len_d is live (inc_seqlen) — device-len rides it, plus=-1.
5758                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5759                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5760                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
5761                    e.fa_decode_rows(&q, &kp, &vp, &mut attn, hd, nh, nkv, kvl.len - 1, 1,
5762                                     scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
5763                                     Some((&kvl.len_d, -1)), false, false,
5764                                     Some((&mut aq8, &mut ad8)))?;
5765                    fa_q8 = Some((aq8, ad8));
5766                } else if swa && kvl.len > win && hd == 256
5767                    && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
5768                    // windowed regime: SAME rows_w kernel as verify, t=1 (parity law).
5769                    let kp = e.view_u8(&kvl.k, kvl.len * kvl.k_tok_bytes);
5770                    let vp = e.view_u8(&kvl.v, kvl.len * kvl.v_tok_bytes);
5771                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
5772                    e.fa_decode_rows_w(&q, &kp, &vp, &mut attn, hd, nh, nkv, &kvl.len_d, -1,
5773                                       1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes,
5774                                       Some((&mut aq8, &mut ad8)))?;
5775                    fa_q8 = Some((aq8, ad8));
5776                } else {
5777                    let (off_tok, t_kv) = if swa && kvl.len > win { (kvl.len - win, win) }
5778                                          else { (0, kvl.len) };
5779                    let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
5780                                                 (off_tok + t_kv) * kvl.k_tok_bytes);
5781                    let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
5782                                                 (off_tok + t_kv) * kvl.v_tok_bytes);
5783                    e.fa_decode_kvmod(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t_kv, scale,
5784                                kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
5785                }
5786            }
5787            Some((b_swa, b_glob)) => {
5788                // capture: full-buffer views + device length. PARITY LAW port (graph arc step
5789                // 3): the SAME regime branches as dc-eager, pinned per arm — b_swa is the
5790                // exact-key t_kv for the dc family (n_splits constant per bucket), b_glob is
5791                // the RUNG max for the rows family (kernels derive per-replay splits from
5792                // kvl.len_d; grid/partials sized for the rung end, excess splits exit).
5793                let k_view = e.view_u8(&kvl.k, kvl.k.len());
5794                let v_view = e.view_u8(&kvl.v, kvl.v.len());
5795                let rows_on = std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0");
5796                let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5797                if !swa && hd == 512 && b_glob >= crate::fa512_min_tkv() && rows_on {
5798                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
5799                    e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, b_glob - 1,
5800                                     1, scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
5801                                     Some((&kvl.len_d, -1)), false, false,
5802                                     Some((&mut aq8, &mut ad8)))?;
5803                    fa_q8 = Some((aq8, ad8));
5804                } else if swa && b_swa > win && hd == 256 && rows_on {
5805                    let (mut aq8, mut ad8) = e.uninit_q8_pair(nh * hd)?;
5806                    e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
5807                                       &kvl.len_d, -1, 1, scale, win,
5808                                       kvl.k_tok_bytes, kvl.v_tok_bytes,
5809                                       Some((&mut aq8, &mut ad8)))?;
5810                    fa_q8 = Some((aq8, ad8));
5811                } else {
5812                    let b = if swa { b_swa } else { b_glob };
5813                    e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, b,
5814                                   scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
5815                                   swa && crate::Engine::wkv_on())?;
5816                }
5817            }
5818        }
5819        // combine-q8 emit (wave-5b m=1 port): the rows arms produced the wo activation pair
5820        // in-combine — ride the slotted arm's exact matvec route (parity by construction).
5821        if let Some((aq8, ad8)) = fa_q8 {
5822            let mut y = e.uninit(fa.wo.out_features())?;
5823            self.g4_matvec_m1_into(e, &fa.wo, &aq8, &ad8, &mut y)?;
5824            return Ok(y);
5825        }
5826        Ok(e.matmul(&fa.wo, &attn, 1)?)
5827    }
5828
5829    /// gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full
5830    /// dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch
5831    /// per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window
5832    /// views in-graph); caller gates and falls back to the dc-eager loop.
5833    pub fn gemma4_generate_graph(&self, e: &Engine, prompt_pos: usize, first_token: u32,
5834                                 cache: &mut Cache, max_new: usize, eos: &[u32],
5835                                 mut on_token: impl FnMut(u32) -> bool)
5836                                 -> Result<(Vec<u32>, crate::decode::StopReason), Box<dyn std::error::Error>> {
5837        if self.is_gemma4_e4b() {
5838            return Err("E4B graph serving is unwired (HANDOVER-E4B.md) — dc-eager is the serving arm".into());
5839        }
5840        use crate::decode::StopReason;
5841        let n_vocab = self.output.out_features();
5842        let n_embd = self.cfg.n_embd as usize;
5843        let embd_gpu = self.embd_gpu.get_or_init(|| {
5844            e.upload_u8(&self.embd.raw).expect("embed table upload")
5845        });
5846        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
5847        for kvl in cache.kv.iter_mut().flatten() {
5848            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5849        }
5850        let mut token_d = e.stream().clone_htod(&[first_token])?;
5851        let mut pos_d = e.htod_i32(&[prompt_pos as i32])?;
5852        let g4 = self.cfg.gemma4.as_ref().unwrap();
5853        let (hd_s, hd_g) = (g4.key_length_swa as usize, g4.key_length_global as usize);
5854        // per-layer nkv: swa vs global counts from the pattern (uniform within class).
5855        let nkv_s = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
5856            .find(|p| *p.1).map(|p| *p.0 as usize).unwrap_or(8);
5857        let nkv_g = g4.head_count_kv.iter().zip(g4.swa_pattern.iter())
5858            .find(|p| !*p.1).map(|p| *p.0 as usize).unwrap_or(2);
5859        let mut graphs: std::collections::HashMap<((bool, usize), (bool, usize), bool, bool),
5860                                                  (cudarc::driver::CudaGraph,
5861                                                   Vec<Box<dyn std::any::Any + Send>>)> = Default::default();
5862        // ALLOC-FREE capture: persistent transient slots — the captured graph carries zero
5863        // mem nodes (the 226us/launch tax). Slots must outlive every cached graph.
5864        let mut slots = self.g4_dc_slots(e)?;
5865        // Chunked replay ring: tokens park on-device; ONE drain sync per chunk. ring_base is
5866        // baked at the door entry (the modulo keeps every capture valid indefinitely).
5867        const RING: usize = 64;
5868        // DRAIN pinned at 1 (2026-07-23): relaunching the SAME graph exec before its prior
5869        // launch completes is ILLEGAL (chunk=4 -> ILLEGAL_ADDRESS; chunk=1 clean). Pipelining
5870        // needs alternating exec instances — and the measured payoff was ~0 (the ~200us
5871        // cuGraphLaunch host cost already overlaps its own launch's GPU work; llama pays
5872        // 885us/launch the same way). MEMRA_GRAPH_DRAIN raises it only for experiments.
5873        const DRAIN: usize = 1;
5874        let mut ring = e.stream().alloc_zeros::<u32>(RING)?;
5875        let ring_base = prompt_pos;
5876        let mut out = Vec::with_capacity(max_new);
5877        let mut reason = StopReason::MaxNew;
5878        let mut next = first_token;
5879        let mut captures = 0usize;
5880        for _ in 0..max_new {
5881            out.push(next);
5882            if eos.contains(&next) { reason = StopReason::Eos; break; }
5883            if !on_token(next) { reason = StopReason::Callback; break; }
5884            let t_kv = cache.pos + 1;
5885            // Bucket key per ARM (graph arc step 3):
5886            //  - swa component: dc family under the window (exact (fa_vec, n_splits) key —
5887            //    n_splits must be constant per bucket); rows_w above it (window-constant, so
5888            //    the component collapses to a single marker).
5889            //  - global component: dc family under the fa512 floor (exact key); rows_dpl16
5890            //    at/above it — the kernel derives splits from len_d per replay, so buckets
5891            //    are power-of-2 RUNGS (one capture per doubling; grid sized for the rung end).
5892            let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
5893            let f512 = crate::fa512_min_tkv();
5894            let key_s = if t_kv > win { (true, usize::MAX) }
5895                        else { e.fa_bucket_key(t_kv, hd_s, nkv_s, crate::Engine::wkv_on()) };
5896            let (key_g, rung_end) = if t_kv >= f512 {
5897                // strict upper bound: the rung must ROLL at exact powers (t_kv==1024 starts
5898                // the [1024,2048) bucket) — sizing covers every replayed T_kv < end.
5899                let end = (t_kv + 1).next_power_of_two().max(f512 * 2);
5900                ((true, end), end)
5901            } else { (e.fa_bucket_key(t_kv, hd_g, nkv_g, false), t_kv) };
5902            let key = (key_s, key_g, t_kv >= f512, t_kv > win);
5903            if !graphs.contains_key(&key) {
5904                let bucket_max = (t_kv, rung_end);
5905                // snapshot device+host state (the 3 capture-warmup runs leave no residue).
5906                let snap = cache.snapshot(e)?;
5907                let pos_save = e.dtoh_i32_one(&pos_d)?;
5908                let len_save: Vec<Option<i32>> = cache.kv.iter()
5909                    .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
5910                let tok_save = e.dtoh_u32_one(&token_d)?;
5911                // RETAINED capture (2026-07-23): the plain capture_graph left pool-transient
5912                // clones as dead COPY NODES replayed every launch — the E4B 0.74ms/token
5913                // regression class, and this door's measured -8.8%. The keeper pins warmup
5914                // transients so the captured graph holds kernel nodes only.
5915                let graph = {
5916                    let tok_ref = &mut token_d;
5917                    let pos_ref = &mut pos_d;
5918                    let cache_ref = &mut *cache;
5919                    let slots_ref = &mut slots;
5920                    let ring_ref = &mut ring;
5921                    e.capture_graph_retained_flags(
5922                        cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY,
5923                        |e| {
5924                        // self-feeding: the argmax writes token_d itself.
5925                        let tok_in = unsafe { &*(tok_ref as *const CudaSlice<u32>) };
5926                        let sl = unsafe { &mut *(slots_ref as *mut G4DcSlots) };
5927                        let rg = unsafe { &mut *(ring_ref as *mut CudaSlice<u32>) };
5928                        self.gemma4_decode_step_dc_slotted(e, tok_in, pos_ref, embd_gpu, qt, rb,
5929                                                           cache_ref, n_vocab, Some(bucket_max),
5930                                                           sl, tok_ref, Some((rg, ring_base)))
5931                    })?
5932                };
5933                cache.rollback(e, &snap, 0)?;
5934                e.set_i32_one(&mut pos_d, pos_save)?;
5935                for (il, ls) in len_save.iter().enumerate() {
5936                    if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
5937                        e.set_i32_one(&mut kvl.len_d, *v)?;
5938                    }
5939                }
5940                e.set_u32_one(&mut token_d, tok_save)?;
5941                if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
5942                    if let Ok(c) = crate::graph_update::node_census(&graph.0) {
5943                        eprintln!("[graph-census] {c:?}");
5944                    }
5945                }
5946                graphs.insert(key, graph);
5947                captures += 1;
5948            }
5949            // CHUNKED REPLAY: enqueue up to DRAIN launches back-to-back (the ~200us host
5950            // cost of each cuGraphLaunch overlaps the PREVIOUS launch's ~11ms GPU work),
5951            // then ONE sync + ring drain. The chunk must stay inside this bucket key and
5952            // the budget; capture warmups already emitted their tokens through the ring.
5953            let mut chunk = 1usize;
5954            let drain_cap: usize = std::env::var("MEMRA_GRAPH_DRAIN").ok()
5955                .and_then(|v| v.parse().ok()).unwrap_or(DRAIN);
5956            while chunk < drain_cap && out.len() + chunk < max_new {
5957                let t_next = cache.pos + 1 + chunk;
5958                let key_s2 = if t_next > win { (true, usize::MAX) }
5959                             else { e.fa_bucket_key(t_next, hd_s, nkv_s, crate::Engine::wkv_on()) };
5960                let key_g2 = if t_next >= f512 {
5961                    (true, (t_next + 1).next_power_of_two().max(f512 * 2))
5962                } else { e.fa_bucket_key(t_next, hd_g, nkv_g, false) };
5963                if (key_s2, key_g2, t_next >= f512, t_next > win) != key { break; }
5964                chunk += 1;
5965            }
5966            let g = &graphs.get(&key).unwrap().0;
5967            for _ in 0..chunk { g.launch()?; }
5968            e.stream().synchronize()?;
5969            let ringh = e.dtoh_u32(&ring)?;
5970            for j in 0..chunk {
5971                let pos_j = cache.pos + j;
5972                let tok_j = ringh[(pos_j - ring_base) % RING];
5973                cache.pos += 0; // advanced below in one shot
5974                if j + 1 == chunk { next = tok_j; }
5975                else {
5976                    out.push(tok_j);
5977                    if eos.contains(&tok_j) || !on_token(tok_j) {
5978                        reason = if eos.contains(&tok_j) { StopReason::Eos }
5979                                 else { StopReason::Callback };
5980                        // roll device/host state back to the stop point.
5981                        let keep = cache.pos + j + 1;
5982                        e.set_i32_one(&mut pos_d, keep as i32)?;
5983                        for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
5984                            e.set_i32_one(&mut kvl.len_d, keep as i32)?;
5985                            kvl.len = keep;
5986                        }
5987                        cache.pos = keep;
5988                        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
5989                            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
5990                        }
5991                        return Ok((out, reason));
5992                    }
5993                }
5994            }
5995            cache.pos += chunk;
5996            for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) { kvl.len += chunk; }
5997        }
5998        if std::env::var("MEMRA_GRAPH_STATS").is_ok() {
5999            eprintln!("[gemma-graph] captures={captures} buckets={}", graphs.len());
6000        }
6001        Ok((out, reason))
6002    }
6003
6004    /// gemma4 VERIFY step (spec decode): t tokens batched through the trunk at positions
6005    /// pos0..pos0+t-1, K/V rows appended to the quantized cache (caller rolls back rejected
6006    /// rows via kvl.len), per-token causal windowed attend (fa_decode per token — the same
6007    /// kernel family as T=1 decode at each token's t_kv). Returns [t, n_vocab] softcapped
6008    /// logits (host) + advances cache.pos by t.
6009    pub(crate) fn gemma4_decode_step_t(&self, e: &Engine, tokens: &[u32], pos0: usize,
6010                                       cache: &mut Cache)
6011                                       -> Result<Vec<f32>, Box<dyn std::error::Error>> {
6012        Ok(self.gemma4_decode_step_t_h(e, tokens, pos0, cache)?.0)
6013    }
6014
6015    /// GREEDY verify: per-row DEVICE argmax (t x 4B host traffic instead of the t x 1MB logits
6016    /// stack — softcap skipped: tanh is monotonic, per-row argmax unaffected). Returns
6017    /// (argmax ids [t], post-output_norm hidden stack [t, n_embd]).
6018    pub(crate) fn gemma4_decode_step_t_am(&self, e: &Engine, tokens: &[u32], pos0: usize,
6019                                          cache: &mut Cache)
6020                                          -> Result<(Vec<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6021        let (ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6022        let t = tokens.len();
6023        let n_vocab = self.output.out_features();
6024        let mut toks = e.stream().alloc_zeros::<u32>(t)?;
6025        for i in 0..t {
6026            e.argmax_token_device_col(&ld, i, n_vocab, &mut toks, i)?;
6027        }
6028        Ok((e.dtoh_u32(&toks)?, hn))
6029    }
6030
6031    /// Device-token verify (async spec round): tokens live in tok_d[0..t]; per-row argmax
6032    /// lands in a DEVICE buffer (no host logits). Returns (vam_d [t] u32 device, hn stack).
6033    pub(crate) fn gemma4_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6034                                              pos0: usize, cache: &mut Cache)
6035                                              -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6036        let (ld, hn) = self.gemma4_verify_trunk(e, &vec![0u32; t], pos0, cache, Some(tok_d))?;
6037        let n_vocab = self.output.out_features();
6038        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6039        for i in 0..t {
6040            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6041        }
6042        Ok((vam, hn))
6043    }
6044
6045    /// gemma4 verify + the POST-output_norm hidden stack [t, n_embd] (the drafter's h input —
6046    /// llama's h_nextn convention).
6047    pub(crate) fn gemma4_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
6048                                         cache: &mut Cache)
6049                                         -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6050        let (mut ld, hn) = self.gemma4_verify_trunk(e, tokens, pos0, cache, None)?;
6051        let t = tokens.len();
6052        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6053        e.softcap(&mut ld, cap, t * self.output.out_features())?;
6054        Ok((e.dtoh(&ld)?, hn))
6055    }
6056
6057    /// Persistent device scratch for the BURST verify stream (one per generation): rope pos
6058    /// rows [cap] + per-row t_kv counters (device-filled from the round counter, zero H2D).
6059    pub(crate) fn verify_stream_scratch(&self, e: &Engine, cap: usize)
6060                                        -> Result<VerifyStreamScratch, Box<dyn std::error::Error>> {
6061        Ok(VerifyStreamScratch {
6062            pos_d: e.htod_i32(&vec![0i32; cap])?,
6063            row_ctrs: (0..cap).map(|_| e.htod_i32(&[0])).collect::<Result<_, _>>()?,
6064        })
6065    }
6066
6067    /// BURST verify trunk (device-slot twin): tokens from a device buffer, rope positions
6068    /// iota'd from `ctr`, appends/attention at the layers' len_d counters (verify_attn_stream),
6069    /// NO host cache.pos/len advance. Returns (per-row device argmaxes [t], post-norm hidden
6070    /// stack [t, n_embd]) — the burst's accept/seed inputs, zero host readbacks.
6071    /// `scr` = PERSISTENT per-gen scratch (pos rows + per-row t_kv counters): the first cut
6072    /// htod_i32-allocated these per call — 9 pageable H2D copies per round, each a stream
6073    /// sync, exactly the turnaround the burst exists to remove.
6074    pub(crate) fn gemma4_verify_t_am_stream(&self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize,
6075                                            ctr: &CudaSlice<i32>, hint: usize,
6076                                            cache: &mut Cache,
6077                                            scr: &mut VerifyStreamScratch)
6078                                            -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6079        let n_embd = self.cfg.n_embd as usize;
6080        let eps = self.cfg.rms_eps;
6081        assert!(t <= scr.row_ctrs.len() && t <= 64);
6082        e.i32_iota_from(ctr, &mut scr.pos_d, t)?;
6083        for i in 0..t {
6084            e.i32_copy_add(ctr, &mut scr.row_ctrs[i], (i + 1) as i32)?;
6085        }
6086        let (pos_d, row_ctrs) = (&scr.pos_d, &scr.row_ctrs);
6087        let embd_gpu = self.embd_gpu.get_or_init(|| {
6088            e.upload_u8(&self.embd.raw).expect("embed table upload")
6089        });
6090        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6091        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
6092        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6093        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6094        let n_layers = self.layers.len();
6095        for (il, layer) in self.layers.iter().enumerate() {
6096            let (hq, hdq) = match h_carry.take() {
6097                Some(p) => p,
6098                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6099            };
6100            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6101            let o = self.gemma4_verify_attn_stream(e, fa, il, &hq, &hdq, pos_d, t, cache,
6102                                                    hint, row_ctrs)?;
6103            let mut cur = e.uninit(t * n_embd)?;
6104            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6105            let next_norm = if il + 1 < n_layers {
6106                Some(self.layers[il + 1].attn_norm.float_data())
6107            } else { None };
6108            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6109            x = xn;
6110            h_carry = hn;
6111            self.dflash_tap(e, cache, il, &x, t)?;
6112        }
6113        let mut hn = e.uninit(t * n_embd)?;
6114        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6115        let ld = e.matmul(&self.output, &hn, t)?;
6116        let n_vocab = self.output.out_features();
6117        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
6118        for i in 0..t {
6119            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
6120        }
6121        Ok((vam, hn))
6122    }
6123
6124    /// Verify trunk core: returns (UN-softcapped logits device [t, n_vocab], post-output_norm
6125    /// hidden stack [t, n_embd]); appends KV rows + advances cache.pos.
6126    /// DFlash tap write (dflash lane): copy the post-layer residual rows of `x` into the
6127    /// armed sink at the tap slot for `il` (row-major [t, n_taps*hidden]). Per-row D2D
6128    /// copies — t <= block_size on verify; prime pays t*n_taps once per prompt (dedicated
6129    /// kernel later if it shows in the profile).
6130    fn dflash_tap(&self, e: &Engine, cache: &mut Cache, il: usize, x: &CudaSlice<f32>, t: usize)
6131                  -> Result<(), Box<dyn std::error::Error>> {
6132        let Some(taps) = cache.dflash_taps.as_mut() else { return Ok(()) };
6133        let Some(slot) = taps.layer_ids.iter().position(|&l| l == il) else { return Ok(()) };
6134        let h = taps.hidden;
6135        let n_taps = taps.layer_ids.len();
6136        debug_assert_eq!(taps.t, t);
6137        let xv = e.view(x, t * h);
6138        for r in 0..t {
6139            let row = xv.slice(r * h..(r + 1) * h);
6140            e.copy_view_into(&mut taps.buf, r * n_taps * h + slot * h, &row, h)?;
6141        }
6142        Ok(())
6143    }
6144
6145    fn gemma4_verify_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
6146                           tok_dev: Option<&CudaSlice<u32>>)
6147                           -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6148        let n_embd = self.cfg.n_embd as usize;
6149        let eps = self.cfg.rms_eps;
6150        let t = tokens.len();
6151        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6152        let pos_d = e.htod_i32(&pos)?;
6153        let mut x = match tok_dev {
6154            Some(td) => {
6155                let embd_gpu = self.embd_gpu.get_or_init(|| {
6156                    e.upload_u8(&self.embd.raw).expect("embed table upload")
6157                });
6158                let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
6159                e.embed_gather_device_td(embd_gpu, td, t, n_embd, qt, rb)?
6160            }
6161            None => e.htod(&self.embd.gather(n_embd, tokens))?,
6162        };
6163        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6164        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6165        let n_layers = self.layers.len();
6166        for (il, layer) in self.layers.iter().enumerate() {
6167            let (hq, hdq) = match h_carry.take() {
6168                Some(p) => p,
6169                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, t, eps)?,
6170            };
6171            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6172            let o = self.gemma4_verify_attn(e, fa, il, &hq, &hdq, &pos_d, t, cache)?;
6173            let mut cur = e.uninit(t * n_embd)?;
6174            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, t, eps)?;
6175            let next_norm = if il + 1 < n_layers {
6176                Some(self.layers[il + 1].attn_norm.float_data())
6177            } else { None };
6178            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, t, next_norm)?;
6179            x = xn;
6180            h_carry = hn;
6181            self.dflash_tap(e, cache, il, &x, t)?;
6182        }
6183        let mut hn = e.uninit(t * n_embd)?;
6184        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6185        let mut ld = e.matmul(&self.output, &hn, t)?;
6186        self.gemma4_suppress(e, &mut ld, t)?;   // before the per-row argmax consumers
6187        cache.pos += t;
6188        Ok((ld, hn))
6189    }
6190
6191    /// Verify attention: project/norm/rope t rows, append them to the cache, then attend each
6192    /// row causally over [win_off_i .. base+i] via the SAME fa_decode dispatch as T=1 decode.
6193    /// BURST verify attention (device-slot twin of `gemma4_verify_attn`): the append lands
6194    /// at the layer's len_d counter (rows_dc), attention bases ride the counter (rows /
6195    /// rows_w with base_dev), and NO host len is read or bumped — `hint` is a host UPPER
6196    /// bound on base_len used only for split sizing and arm gating (the burst loop passes
6197    /// pos0 + burst slack). Host len mirrors re-sync at the burst drain.
6198    #[allow(clippy::too_many_arguments)]
6199    fn gemma4_verify_attn_stream(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6200                                 hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6201                                 pos_d: &CudaSlice<i32>, t: usize,
6202                                 cache: &mut Cache, hint: usize,
6203                                 row_ctrs: &[CudaSlice<i32>])
6204                                 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6205        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6206        let eps = self.cfg.rms_eps;
6207        let aux = self.gemma4_aux.as_ref().unwrap();
6208        let h0 = e.zeros(0)?;
6209        let h = &h0;
6210        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
6211        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
6212        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6213        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6214        let fused_qkv = if f2b {
6215            if swa {
6216                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6217                    .map(|(a, b, c)| (a, b, Some(c)))
6218            } else {
6219                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6220                    .map(|(a, b)| (a, b, None))
6221            }
6222        } else { None };
6223        let (q0, k0, v0) = match fused_qkv {
6224            Some((a, b, cv)) => {
6225                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6226                (a, b, v)
6227            }
6228            None => {
6229                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6230                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6231                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6232                         else { e.clone_dtod(&k0)? };
6233                (q0, k0, v0)
6234            }
6235        };
6236        let mut q = e.uninit(t * nh * hd)?;
6237        let mut k = e.uninit(t * nkv * hd)?;
6238        let mut v = e.uninit(t * nkv * hd)?;
6239        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
6240        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
6241        let ff = if swa { None } else {
6242            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6243        };
6244        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6245                            &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6246                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
6247        let kvl = cache.kv[il].as_mut().unwrap();
6248        // append at the DEVICE slot; the counter advances by t on-device.
6249        e.append_kv_quantized_rows_dc(&k, &v, &mut kvl.k, &mut kvl.v, &kvl.len_d, t,
6250                                      kvl.kv_dim_k, kvl.kv_dim_v,
6251                                      kvl.k_tok_bytes, kvl.v_tok_bytes,
6252                                      (!swa && crate::Engine::gkv_on())
6253                                          || (swa && crate::Engine::wkv_on()))?;
6254        // len_d is NOT advanced here: the burst's device rollback (spec_rollback_stream) is
6255        // the sole len writer after this round's attention (base stays = old len, plus = 0).
6256        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6257        let mut attn = e.uninit(t * nh * hd)?;
6258        let k_view = e.view_u8(&kvl.k, kvl.k.len());
6259        let v_view = e.view_u8(&kvl.v, kvl.v.len());
6260        // arm gating on the host HINT (upper bound; burst entry guards hint >= the vec floor
6261        // and a stable window regime — the same rung/regime keys as the draft graph).
6262        if swa && hint + 1 >= win {
6263            // fully-windowed rows: per-row window geometry from the counter (plus = 0: len_d
6264            // still holds the pre-append len; row r's T_kv = ctr + r + 1).
6265            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6266                               &kvl.len_d, 0, t, scale, win,
6267                               kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6268        } else if hd == 512 && hint + t < crate::fa512_min_tkv() {
6269            // globals UNDER the fa512 crossover: eager runs the per-row fa_decode_kvmod
6270            // fallback there (parity with t=1 decode) — mirror it with per-row fa_decode_dc
6271            // (bit-correct for any t_kv <= bucket; the drafter dc arc proved the pairing).
6272            // Burst entry gates the horizon onto one side of the crossover, so hint decides
6273            // for every row.
6274            // NO .max(512): fa_decode_dc gates its hd512 dpl16-vs-scalar pick on bucket_max
6275            // (mirroring eager's fa512 floor) — forcing 512 here flipped the arm to dpl16
6276            // while eager ran scalar (il=5 KV drift, the burst's 4/128). SAME LAW capped
6277            // from above (2026-07-13): a regime-pinned hint near the floor (round-graph
6278            // captures use f512-1) pow2-rounds PAST it — clamp under the floor, or the arm
6279            // re-flips to dpl16 (the round-graph 4/64). The scalar unified self-splits, so
6280            // any bucket >= the live length is exact.
6281            let bucket = (hint + t + 2).next_power_of_two()
6282                .min(crate::fa512_min_tkv().saturating_sub(1));
6283            let qv = e.view(&q, t * nh * hd);
6284            for i in 0..t {
6285                let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
6286                let mut q_one = e.uninit(nh * hd)?;
6287                e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6288                let mut a_one = e.uninit(nh * hd)?;
6289                e.fa_decode_dc(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv,
6290                               &row_ctrs[i], bucket, scale,
6291                               kvl.k_tok_bytes, kvl.v_tok_bytes, false)?;
6292                e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6293            }
6294        } else if hd == 512 {
6295            // globals past the crossover: dpl16 dc twin via the shared rows wrapper (hint
6296            // sizes splits — upper bound; splits beyond the device len exit in-kernel).
6297            e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, hint, t, scale,
6298                             kvl.k_tok_bytes, kvl.v_tok_bytes,
6299                             Some((&kvl.len_d, 0)), false, false, None)?;
6300        } else {
6301            // hd256 under-window: v4 device-len rows twin.
6302            e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6303                                &kvl.len_d, hint + t, t, scale,
6304                                kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6305                                swa && crate::Engine::wkv_on())?;
6306        }
6307        Ok(e.matmul(&fa.wo, &attn, t)?)
6308    }
6309
6310    fn gemma4_verify_attn(&self, e: &Engine, fa: &crate::hybrid::FullAttnLayer, il: usize,
6311                          hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6312                          pos_d: &CudaSlice<i32>, t: usize,
6313                          cache: &mut Cache)
6314                          -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6315        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
6316        let eps = self.cfg.rms_eps;
6317        let aux = self.gemma4_aux.as_ref().unwrap();
6318        let n_embd = self.cfg.n_embd as usize;
6319        let _ = n_embd;
6320
6321        let h0 = e.zeros(0)?;
6322        let h = &h0;
6323        // BATCHED FUSED qkv (MEMRA_F2B=1, megakernel microcosm): swa layers fuse all three,
6324        // globals fuse q,k (v := k clone). Bit-identical per row; segments tail-fill.
6325        static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6326        let f2b = *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0"));
6327        let fused_qkv = if f2b {
6328            if swa {
6329                e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6330                    .map(|(a, b, c)| (a, b, Some(c)))
6331            } else {
6332                e.matmul_q4_fused2_batched(&fa.wq, &fa.wk, hq, hdq, t)?
6333                    .map(|(a, b)| (a, b, None))
6334            }
6335        } else { None };
6336        let (q0, k0, v0) = match fused_qkv {
6337            Some((a, b, cv)) => {
6338                let v = match cv { Some(c) => c, None => e.clone_dtod(&b)? };
6339                (a, b, v)
6340            }
6341            None => {
6342                let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6343                let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, t)?;
6344                let v0 = if swa { e.matmul_pre(&fa.wv, hq, hdq, h, t)? }
6345                         else { e.clone_dtod(&k0)? };
6346                (q0, k0, v0)
6347            }
6348        };
6349        let mut q = e.uninit(t * nh * hd)?;
6350        let mut k = e.uninit(t * nkv * hd)?;
6351        let mut v = e.uninit(t * nkv * hd)?;
6352        // E4B wave-3 fold (2026-07-23 decode-dust port): q/k/v norms + q/k rope in ONE
6353        // launch — rope math verbatim on the normed rows; V ones-rms, never roped.
6354        let ff = if swa { None } else {
6355            Some(aux.rope_freqs.as_ref().expect("gemma4 global rope needs rope_freqs.weight"))
6356        };
6357        e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6358                            &aux.ones, &mut q, &mut k, &mut v, hd, nh * t, nkv * t,
6359                            pos_d, nh, nkv, base, 1.0, ff, eps)?;
6360        let kvl = cache.kv[il].as_mut().unwrap();
6361        let base_len = kvl.len;
6362        e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, base_len, t,
6363                                   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()))?;
6364        kvl.len += t;
6365        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6366        let mut attn = e.uninit(t * nh * hd)?;
6367        // ROWS twin when no window offset is in play (globals are hd512-ineligible; SWA rows
6368        // under the window need no offset): ONE launch, per-row causal == the per-token loop.
6369        let rows_ok = (hd == 256 && base_len + 1 >= crate::fa_vec_min_tkv())
6370            // gemma globals: hd512 rows twin in the dpl16 vec regime (row 0 gates the batch);
6371            // decode rides the SAME symbol at t=1 (parity law).
6372            || (hd == 512 && !swa && base_len + 1 >= crate::fa512_min_tkv());
6373        if rows_ok && (!swa || base_len + t <= win) {
6374            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6375            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6376            if hd == 512 {
6377                // device-len twin: sync the counter to the verify base (async arg-store).
6378                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6379                e.fa_decode_rows(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, base_len, t,
6380                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6381                                 Some((&kvl.len_d, 0)), false,
6382                                 swa && crate::Engine::wkv_on(), None)?;
6383            } else {
6384                // hd256: the SAME v4_dc symbol the burst verify launches (PARITY LAW — the
6385                // host-base rows_v4 twin compiles apart from v4_dc and the two-symbol split
6386                // drifted the burst's persisted KV at il=8; one symbol, counter arg-stored).
6387                e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6388                e.fa_decode_rows_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6389                                    &kvl.len_d, base_len + t, t, scale,
6390                                    kvl.k_tok_bytes, kvl.v_tok_bytes, 0,
6391                                    swa && crate::Engine::wkv_on())?;
6392            }
6393            return Ok(e.matmul(&fa.wo, &attn, t)?);
6394        }
6395        // WINDOWED rows twin (deep ctx, every row fully windowed): one launch, ABSOLUTE-index
6396        // per-row geometry over the prefix view. PARITY LAW (2026-07-10 root-cause): textually
6397        // identical kernels do NOT compile bit-identically (nvcc unrolls fa_decode_vec_q 2x vs
6398        // its rows_w clone — SASS-proven; the unpinned score `+=` chain then rounds apart), so
6399        // decode-vs-verify parity comes from BOTH sides launching THIS SAME rows_w kernel
6400        // (decode passes t=1) — bitwise equal per position by symbol identity, any lane.
6401        // MEMRA_GEMMA_ROWS_W=0 -> per-token loop (decode falls back to fa_decode views too).
6402        if hd == 256 && swa && base_len + 1 >= win
6403            && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6404            let k_view = e.view_u8(&kvl.k, (base_len + t) * kvl.k_tok_bytes);
6405            let v_view = e.view_u8(&kvl.v, (base_len + t) * kvl.v_tok_bytes);
6406            e.i32_set_k(&mut kvl.len_d, base_len as i32)?;
6407            e.fa_decode_rows_w(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, &kvl.len_d, 0,
6408                               t, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6409            return Ok(e.matmul(&fa.wo, &attn, t)?);
6410        }
6411        for i in 0..t {
6412            let avail = base_len + i + 1;
6413            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
6414            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6415                                         (off_tok + t_kv) * kvl.k_tok_bytes);
6416            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6417                                         (off_tok + t_kv) * kvl.v_tok_bytes);
6418            let qi = e.view(&q, t * nh * hd);
6419            let q_row = qi.slice(i * nh * hd..(i + 1) * nh * hd);
6420            let mut q_one = e.uninit(nh * hd)?;
6421            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6422            let mut a_one = e.uninit(nh * hd)?;
6423            // straddle rounds: windowed rows must ride the SAME rows_w kernel decode uses
6424            // (parity law above); global hd512 rows past the fa512 floor likewise ride the
6425            // rows_dpl16 twin; remaining rows keep the gated fa_decode pair.
6426            if swa && avail > win && hd == 256
6427                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6428                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6429                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6430                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6431                e.fa_decode_rows_w(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, &kvl.len_d, 0,
6432                                   1, scale, win, kvl.k_tok_bytes, kvl.v_tok_bytes, None)?;
6433            } else if !swa && hd == 512 && avail >= crate::fa512_min_tkv()
6434                && std::env::var("MEMRA_GEMMA_ROWS_W").as_deref() != Ok("0") {
6435                let kp = e.view_u8(&kvl.k, avail * kvl.k_tok_bytes);
6436                let vp = e.view_u8(&kvl.v, avail * kvl.v_tok_bytes);
6437                e.i32_set_k(&mut kvl.len_d, (avail - 1) as i32)?;
6438                e.fa_decode_rows(&q_one, &kp, &vp, &mut a_one, hd, nh, nkv, avail - 1, 1,
6439                                 scale, kvl.k_tok_bytes, kvl.v_tok_bytes,
6440                                 Some((&kvl.len_d, 0)), false, false, None)?;
6441            } else {
6442                e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
6443                            kvl.k_tok_bytes, kvl.v_tok_bytes, swa && crate::Engine::wkv_on())?;
6444            }
6445            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6446        }
6447        Ok(e.matmul(&fa.wo, &attn, t)?)
6448    }
6449
6450    /// gemma4 T=1 decode step: R8 layer graph over the cache; returns (softcapped logits host,
6451    /// h_seed = pre-output_norm hidden). Advances cache.pos.
6452    pub(crate) fn gemma4_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
6453                                       -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6454        // M1-PP2 door (crate::pp): 2-stage split with an explicit activation handoff.
6455        // Default OFF — unset env means this branch is never taken. The gemma4 arm stays
6456        // 2-stage in M2 (the N-stage gate model is the generic-arm 9B); N>2 warns + runs
6457        // unsplit rather than guessing a fence.
6458        if let Some(split) = crate::pp::pp2_split(self.layers.len()) {
6459            return self.gemma4_decode_step_h_pp2(e, token, cache, split);
6460        }
6461        if crate::pp::pp_cuts(self.layers.len()).is_some() {
6462            crate::pp::warn_unwired_once("gemma4 eager decode (N>2)");
6463        }
6464        let n_embd = self.cfg.n_embd as usize;
6465        let eps = self.cfg.rms_eps;
6466        let pos_d = e.htod_i32(&[cache.pos as i32])?;
6467        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
6468        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
6469        // cross-layer fusion: each tail's closing add+scale also EMITS the next layer's
6470        // attn-normed input pre-quantized q8_1 (the mixer consumes only quantized matmuls).
6471        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6472        let n_layers = self.layers.len();
6473        for (il, layer) in self.layers.iter().enumerate() {
6474            let (hq, hdq) = match h_carry.take() {
6475                Some(p) => p,
6476                None => e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, 1, eps)?,
6477            };
6478            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6479            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, &pos_d, cache)?;
6480            let mut cur = e.uninit(n_embd)?;
6481            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6482            let next_norm = if il + 1 < n_layers {
6483                Some(self.layers[il + 1].attn_norm.float_data())
6484            } else { None };
6485            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6486            x = xn;
6487            h_carry = hn;
6488        }
6489        let mut hn = e.uninit(n_embd)?;
6490        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
6491        let h_seed = e.clone_dtod(&x)?;
6492        let mut ld = e.matmul(&self.output, &hn, 1)?;
6493        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6494        e.softcap(&mut ld, cap, self.output.out_features())?;   // R4 on device (262k host tanh ~ms/step)
6495        self.gemma4_suppress(e, &mut ld, 1)?;
6496        let logits = e.dtoh(&ld)?;
6497        cache.pos += 1;
6498        Ok((logits, h_seed))
6499    }
6500
6501    /// M1-PP2 stage subgraph (gemma4 arm): layers [lo, hi) of the gemma4 T=1 walk with the
6502    /// pre-quantized next-norm carry LOCAL to the range. Enters with a materialized residual
6503    /// `x` (the range head runs its own `rms_norm_q8_1` against ITS layer's attn_norm — for
6504    /// lo == 0 that is exactly the unsplit loop's il==0 arm); exits with the residual
6505    /// materialized (range tail passes next_norm = None, exactly the unsplit last-layer arm).
6506    /// Bit-identity of the cut relies on the kernel-check-pinned `add_scale_rms_norm_q8_1 ==
6507    /// add_scale then rms_norm_q8_1` identity (`pp2-gate` verifies end-to-end).
6508    fn gemma4_decode_layers(&self, e: &Engine, mut x: CudaSlice<f32>, lo: usize, hi: usize,
6509                            pos_d: &CudaSlice<i32>, cache: &mut Cache)
6510                            -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6511        let n_embd = self.cfg.n_embd as usize;
6512        let eps = self.cfg.rms_eps;
6513        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6514        for il in lo..hi {
6515            let layer = &self.layers[il];
6516            let (hq, hdq) = match h_carry.take() {
6517                Some(p) => p,
6518                // range head: il == lo — norm against THIS layer's attn_norm.
6519                None => e.rms_norm_q8_1(&x, self.layers[il].attn_norm.float_data(), n_embd, 1, eps)?,
6520            };
6521            let Mixer::Full(fa) = &layer.mixer else { panic!("gemma4 layer {il} not full-attn") };
6522            let o = self.gemma4_decode_attn(e, fa, il, &hq, &hdq, pos_d, cache)?;
6523            let mut cur = e.uninit(n_embd)?;
6524            e.rms_norm(&o, layer.post_attn_norm.float_data(), &mut cur, n_embd, 1, eps)?;
6525            let next_norm = if il + 1 < hi {
6526                Some(self.layers[il + 1].attn_norm.float_data())
6527            } else { None };
6528            let (xn, hn) = self.gemma4_layer_tail_add_nq(e, layer, &cur, &x, 1, next_norm)?;
6529            x = xn;
6530            h_carry = hn;
6531        }
6532        Ok(x)
6533    }
6534
6535    /// M1-PP2 (increment 2): `gemma4_decode_step_h` as TWO stage subgraphs, each on its
6536    /// own stream (and device, under MEMRA_PP_DEVICES) with the transport-selected
6537    /// boundary handoff — same choreography as the generic arm (decode.rs), same
6538    /// ownership contract (crate::pp). Stage 0 = embed+scale + layers [0, split);
6539    /// stage 1 = layers [split, n) + output_norm + softcapped head.
6540    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam. Gate: `pp2-gate`.
6541    fn gemma4_decode_step_h_pp2(&self, e: &Engine, token: u32, cache: &mut Cache, split: usize)
6542                                -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6543        if crate::pp::pp2_streams_off() {
6544            return self.gemma4_decode_step_h_pp2_samestream(e, token, cache, split);
6545        }
6546        let rt = crate::pp::Pp2Rt::get(e)?;
6547        let e0 = rt.engine(0, e);
6548        let e1 = rt.engine(1, e);
6549        let n_embd = self.cfg.n_embd as usize;
6550        let eps = self.cfg.rms_eps;
6551
6552        // ---- STAGE 0 (its own stream): embed + sqrt(n_embd) scale + layers [0, split) ----
6553        let (pos_d, slot) = {
6554            let _st0 = rt.enter(0);
6555            let pos_d = e0.htod_i32(&[cache.pos as i32])?;
6556            let mut x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
6557            e0.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
6558            let x = self.gemma4_decode_layers(e0, x, 0, split, &pos_d, cache)?;
6559            let slot = rt.tx(0, &x, n_embd)?;
6560            (pos_d, slot)
6561        };
6562
6563        // ---- STAGE 1 (its own stream): RX + layers [split, n) + softcapped head ----
6564        let _st1 = rt.enter(1);
6565        let x = rt.rx(0, slot, n_embd)?;
6566        let x = self.gemma4_decode_layers(e1, x, split, self.layers.len(), &pos_d, cache)?;
6567
6568        let mut hn = e1.uninit(n_embd)?;
6569        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
6570        let h_seed = e1.clone_dtod(&x)?;
6571        let mut ld = e1.matmul(&self.output, &hn, 1)?;
6572        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6573        e1.softcap(&mut ld, cap, self.output.out_features())?;
6574        self.gemma4_suppress(e1, &mut ld, 1)?;
6575        let logits = e1.dtoh(&ld)?;
6576        cache.pos += 1;
6577        Ok((logits, h_seed))
6578    }
6579
6580    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 gemma4 pp2 body verbatim — both
6581    /// stage subgraphs on the ambient compute stream, boundary = two plain dtod copies.
6582    fn gemma4_decode_step_h_pp2_samestream(&self, e: &Engine, token: u32, cache: &mut Cache,
6583                                           split: usize)
6584                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6585        let n_embd = self.cfg.n_embd as usize;
6586        let eps = self.cfg.rms_eps;
6587        let pos_d = e.htod_i32(&[cache.pos as i32])?;
6588
6589        // ---- STAGE 0: embed + sqrt(n_embd) scale + layers [0, split) ----
6590        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
6591        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
6592        let x = self.gemma4_decode_layers(e, x, 0, split, &pos_d, cache)?;
6593
6594        // ---- STAGE BOUNDARY: explicit [n_embd] activation handoff (TX copy, RX copy) ----
6595        let boundary_tx = e.clone_dtod(&x)?;
6596        let boundary_rx = e.clone_dtod(&boundary_tx)?;
6597
6598        // ---- STAGE 1: layers [split, n) + output_norm + softcapped head ----
6599        let x = self.gemma4_decode_layers(e, boundary_rx, split, self.layers.len(), &pos_d, cache)?;
6600
6601        let mut hn = e.uninit(n_embd)?;
6602        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
6603        let h_seed = e.clone_dtod(&x)?;
6604        let mut ld = e.matmul(&self.output, &hn, 1)?;
6605        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
6606        e.softcap(&mut ld, cap, self.output.out_features())?;
6607        self.gemma4_suppress(e, &mut ld, 1)?;
6608        let logits = e.dtoh(&ld)?;
6609        cache.pos += 1;
6610        Ok((logits, h_seed))
6611    }
6612}
6613
6614// ===================================================================================== //
6615//  gemma-4 E4B (per-layer embeddings + KV-sharing) — FIRST-LIGHT forward.               //
6616//  Dedicated simple path (Stage-B matmuls, per-row causal attention over the quantized  //
6617//  cache) so the tuned 26B/31B paths stay untouched. Arch: research/gemma4-bringup/     //
6618//  e4b-arch-map.md; llama reference: src/models/gemma4.cpp (E4B arms). Wired: forward   //
6619//  (prefill logits), prime (tokenwise-equivalent batched trunk), eager decode_step_h.   //
6620//  NOT wired: dc/graph serving arms, verify/spec, chunked prime (HANDOVER-E4B.md).      //
6621// ===================================================================================== //
6622impl HybridModel {
6623    pub fn is_gemma4_e4b(&self) -> bool {
6624        self.gemma4_aux.as_ref().is_some_and(|a| a.e4b.is_some())
6625    }
6626
6627    /// E4B per-layer geometry: nh/nkv derive from the layer's OWN tensor shapes (the E4B GGUF
6628    /// ships a scalar head_count_kv; swa q 8x256 / kv 2x256, global q 4x512 / kv 1x512).
6629    /// KV-shared layers report the SHARE TARGET's kv count (their wk IS the target's tensor).
6630    fn gemma4_e4b_geom(&self, il: usize) -> (usize, usize, usize, f32, f32, bool) {
6631        let g = self.cfg.gemma4.as_ref().unwrap();
6632        let swa = g.swa_pattern[il];
6633        let hd = if swa { g.key_length_swa } else { g.key_length_global } as usize;
6634        let Mixer::Full(fa) = &self.layers[il].mixer else { panic!("e4b layer {il} not full-attn") };
6635        let nh = fa.wq.out_features() / hd;
6636        let nkv = fa.wk.out_features() / hd;
6637        (hd, nkv, nh, if swa { g.rope_base_swa } else { g.rope_base_global }, 1.0, swa)
6638    }
6639
6640    /// E4B KV-share target: Some(target) for the trailing shared layers, None = own cache.
6641    fn gemma4_e4b_kv_target(&self, il: usize) -> Option<usize> {
6642        self.layers[il].gemma4.as_ref()
6643            .and_then(|b| b.e4b.as_ref())
6644            .and_then(|e4| e4.kv_share.map(|t| t as usize))
6645    }
6646
6647    /// E4B prologue: inp_pl [t][n_layer][n_epl] =
6648    ///   ( gather(per_layer_tok_embd, tok)*sqrt(n_epl)
6649    ///   + rms_norm(model_proj . x_scaled * 1/sqrt(n_embd), proj_norm) ) * 1/sqrt(2)
6650    /// (llama gemma4.cpp build_inp_per_layer + project_per_layer_inputs, exact order).
6651    fn gemma4_e4b_inp_pl(&self, e: &Engine, tokens: &[u32], x_scaled: &CudaSlice<f32>, t: usize)
6652                         -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6653        let tok_d = e.stream().clone_htod(&tokens.to_vec())?;
6654        self.gemma4_e4b_inp_pl_dev(e, &tok_d, x_scaled, t)
6655    }
6656
6657    /// Device-token prologue core (dc arm shares it: token ids never touch the host).
6658    fn gemma4_e4b_inp_pl_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
6659                             x_scaled: &CudaSlice<f32>, t: usize)
6660                             -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6661        let aux = self.gemma4_aux.as_ref().unwrap();
6662        let m = aux.e4b.as_ref().unwrap();
6663        let n_embd = self.cfg.n_embd as usize;
6664        let n_layer = self.layers.len();
6665        let width = m.n_epl * n_layer;
6666        let tbl = m.tok_tbl_gpu.get_or_init(|| {
6667            e.upload_u8(&m.tok_embd_bytes).expect("e4b per-layer token table upload")
6668        });
6669        let mut a = e.embed_gather_device_td(tbl, tok_d, t, width, m.tok_embd_qt,
6670                                             m.tok_embd_row_bytes)?;
6671        e.scale_inplace(&mut a, (m.n_epl as f32).sqrt(), t * width)?;
6672        let mut p = e.matmul(&m.model_proj, x_scaled, t)?;
6673        e.scale_inplace(&mut p, 1.0 / (n_embd as f32).sqrt(), t * width)?;
6674        let mut pn = e.uninit(t * width)?;
6675        e.rms_norm(&p, m.proj_norm.float_data(), &mut pn, m.n_epl, t * n_layer,
6676                   self.cfg.rms_eps)?;
6677        let mut out = e.uninit(t * width)?;
6678        e.add_scale(&a, &pn, 1.0 / 2f32.sqrt(), &mut out, t * width)?;
6679        Ok(out)
6680    }
6681
6682    /// E4B attention (t-wide, causal, per-row fa over the QUANTIZED cache — first-light
6683    /// correctness path; the fa rows arms come later). Own-KV layers project+norm+rope k/v
6684    /// and append t rows; KV-shared layers are Q-only over the target layer's cache (which
6685    /// already holds this forward's rows — the target runs earlier in the stack).
6686    #[allow(clippy::too_many_arguments)]
6687    fn gemma4_e4b_attn(&self, e: &Engine, il: usize,
6688                       hq: &CudaSlice<i8>, hdq: &CudaSlice<f32>,
6689                       pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
6690                       dc_bucket: Option<usize>)
6691                       -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6692        let (hd, nkv, nh, base, scale, swa) = self.gemma4_e4b_geom(il);
6693        let eps = self.cfg.rms_eps;
6694        let aux = self.gemma4_aux.as_ref().unwrap();
6695        let Mixer::Full(fa) = &self.layers[il].mixer else { unreachable!() };
6696        // pre-quantized layer input (E4B fusion port, 2026-07-12): ONE norm+quant feeds
6697        // wq/wk/wv via matmul_pre — the first-light path quantized the same h three times
6698        // (the profile's 342 quantize_q8_1/token; glue = 26% of the 6.1ms token).
6699        let h0 = e.zeros(0)?;
6700        let h = &h0;
6701
6702        let ff = if swa { None } else {
6703            Some(aux.rope_freqs.as_ref().expect("e4b global rope needs rope_freqs.weight"))
6704        };
6705        let share = self.gemma4_e4b_kv_target(il);
6706        // Own-KV arms keep the f32 (k, v) alive for the prime fa arm below.
6707        let mut kv_f32: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
6708        let mut q;
6709        if let Some(_tgt) = share {
6710            let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, t)?;
6711            q = e.uninit(t * nh * hd)?;
6712            // Q-only through the same fused norm+rope kernel (rk = 0: the k/v segments are
6713            // empty; q0 stands in for the unused k/v pointers).
6714            let mut kdummy = e.uninit(1)?;
6715            let mut vdummy = e.uninit(1)?;
6716            e.rms_norm_qkv_rope(&q0, &q0, &q0, fa.q_norm.float_data(),
6717                                fa.q_norm.float_data(), &aux.ones,
6718                                &mut q, &mut kdummy, &mut vdummy, hd, nh * t, 0,
6719                                pos_d, nh, 1, base, 1.0, ff, eps)?;
6720        } else {
6721            // wave-4b: ONE concat matvec (wq|wk|wv) at t == 1 when the cat tensor exists;
6722            // else the fused3 grid launch; else per-matvec. The cat output is contiguous
6723            // q|k|v rows — the cat norm+rope twin consumes it directly.
6724            let e4bits = self.layers[il].gemma4.as_ref().and_then(|g| g.e4b.as_ref());
6725            let cat = e4bits.and_then(|e4| e4.qkv_cat.as_ref());
6726            q = e.uninit(t * nh * hd)?;
6727            let mut k = e.uninit(t * nkv * hd)?;
6728            let mut v = e.uninit(t * nkv * hd)?;
6729            if t == 1 && cat.is_some() {
6730                let qkv0 = e.matmul_pre(cat.unwrap(), hq, hdq, h, 1)?;
6731                e.rms_norm_qkv_rope_cat(&qkv0, fa.q_norm.float_data(), fa.k_norm.float_data(),
6732                                        &aux.ones, &mut q, &mut k, &mut v, hd, nh, nkv,
6733                                        pos_d, nh, nkv, base, 1.0, ff, eps)?;
6734            } else {
6735                let (q0, k0, v0) = match if t == 1 {
6736                    e.matmul_q4_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hdq)?
6737                } else {
6738                    // E4B verify f3 port (2026-07-14): the 31B segmented-grid batched qkv
6739                    // on E4B's real-V triple — same MEMRA_F2B seam, bit-identical per row.
6740                    static F2B_QKV: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6741                    if *F2B_QKV.get_or_init(|| std::env::var("MEMRA_F2B").as_deref() != Ok("0")) {
6742                        e.matmul_q4_fused3_batched(&fa.wq, &fa.wk, &fa.wv, hq, hdq, t)?
6743                    } else { None }
6744                } {
6745                    Some(triple) => triple,
6746                    None => (e.matmul_pre(&fa.wq, hq, hdq, h, t)?,
6747                             e.matmul_pre(&fa.wk, hq, hdq, h, t)?,
6748                             e.matmul_pre(&fa.wv, hq, hdq, h, t)?),   // E4B: real v (K != V)
6749                };
6750                // wave-3 fold: q/k/v norms + q/k rope in ONE launch (rope math verbatim on
6751                // the normed rows; V ones-rms, never roped).
6752                e.rms_norm_qkv_rope(&q0, &k0, &v0, fa.q_norm.float_data(),
6753                                    fa.k_norm.float_data(), &aux.ones, &mut q, &mut k, &mut v,
6754                                    hd, nh * t, nkv * t, pos_d, nh, nkv, base, 1.0, ff, eps)?;
6755            }
6756            let kvl = cache.kv[il].as_mut().unwrap();
6757            // class flag must match the cache dims (the g-threading sweep hardcoded `false`
6758            // here and the wkv default corrupted E4B: q8_0 bytes into an e4m3 cache — the
6759            // degenerate tok-0 stream, 2026-07-12).
6760            let cls = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6761            if dc_bucket.is_some() {
6762                // DC arm (graph serving): append at the len_d slot, advance the counter
6763                // in-stream — replay-correct, no host len in the launch args. Host mirrors
6764                // are NOT touched here (the replay loop owns them; a bump at capture-record
6765                // time would double-count the capture iteration).
6766                debug_assert!(t == 1);
6767                // wave 5c: append + len_d inc fused (one launch; single-block ordering).
6768                e.append_kv_quantized_row_dc_inc(&k, &v, &mut kvl.k, &mut kvl.v,
6769                                                 &mut kvl.len_d, kvl.kv_dim_k, kvl.kv_dim_v,
6770                                                 kvl.k_tok_bytes, kvl.v_tok_bytes, cls)?;
6771            } else {
6772                e.append_kv_quantized_rows(&k, &v, &mut kvl.k, &mut kvl.v, kvl.len, t,
6773                                           kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes,
6774                                           kvl.v_tok_bytes, cls)?;
6775                kvl.len += t;
6776            }
6777            kv_f32 = Some((k, v));
6778        }
6779        // attention: per-row causal fa over the (own or target) quantized cache. The cache
6780        // already contains this forward's rows in both arms; row i attends [.., base+i].
6781        let kvl_idx = share.unwrap_or(il);
6782        let kvl = cache.kv[kvl_idx].as_ref().unwrap();
6783        let base_len = kvl.len - t;   // pre-append length (target appended this forward too)
6784        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
6785        let mut attn = e.uninit(t * nh * hd)?;
6786        // PRIME FA ARMS (perf lane 3, completed 2026-07-31 — the H100 board pinned E4B
6787        // prefill at 482 tok/s: the per-row loop ran the DECODE kernel once per token per
6788        // layer). Fresh-prompt (base_len == 0) batched attention, one launch per layer:
6789        //   - own-KV hd256 under the window: f32 fa_prefill (full causal exact — 26B pattern)
6790        //   - own-KV hd256 above the window (swa): fa_prefill_w windowed twin (12B pattern)
6791        //   - own-KV hd512 globals: fa_prefill_hd512 (12B/31B globals pattern)
6792        //   - KV-shared layers (no f32 k/v): fa_prefill_view over the target's QUANTIZED
6793        //     rows (the T=K verify kernel; the target appended this forward's rows already).
6794        //     swa-shared above the window keeps the per-row loop (no windowed view twin).
6795        // Same numeric class split as the 26B prime (f32/batched prefill vs per-row
6796        // quantized decode); run-gen argmax + chat gates arbitrate. MEMRA_NOFA=1 reverts all.
6797        if t > 1 && base_len == 0 && std::env::var("MEMRA_NOFA").is_err() {
6798            if let Some((kf, vf)) = &kv_f32 {
6799                if hd == 256 && t <= win {
6800                    e.fa_prefill(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true)?;
6801                    return Ok(e.matmul(&fa.wo, &attn, t)?);
6802                }
6803                if hd == 256 && swa && t > win {
6804                    e.fa_prefill_w(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale, true,
6805                                   win)?;
6806                    return Ok(e.matmul(&fa.wo, &attn, t)?);
6807                }
6808                if hd == 512 && !swa {
6809                    e.fa_prefill_hd512(&q, kf, vf, &mut attn, hd, nh, nkv, t, t, scale,
6810                                       true)?;
6811                    return Ok(e.matmul(&fa.wo, &attn, t)?);
6812                }
6813            } else if share.is_some() {
6814                let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6815                let k_view = e.view_u8(&kvl.k, kvl.k.len());
6816                let v_view = e.view_u8(&kvl.v, kvl.v.len());
6817                if hd == 256 && (!swa || t <= win) {
6818                    // inline-dequant quantized-view prefill (fa_prefill_q stamps 256/128)
6819                    e.fa_prefill_view(&q, &k_view, &v_view, &mut attn, hd, nh, nkv, t, t,
6820                                      scale, true, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
6821                    return Ok(e.matmul(&fa.wo, &attn, t)?);
6822                }
6823                // remaining shared classes (swa above the window; hd512 globals): dequant
6824                // the target's t rows ONCE to f32, then the same f32 twins as own-KV.
6825                let kv_dim = nkv * hd;
6826                let mut kf = e.uninit(t * kv_dim)?;
6827                let mut vf = e.uninit(t * kv_dim)?;
6828                e.fa_dequant_kv_view_f32(&k_view, &v_view, &mut kf, &mut vf, kv_dim, kv_dim,
6829                                         t, kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
6830                if hd == 512 {
6831                    e.fa_prefill_hd512(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale,
6832                                       true)?;
6833                } else {
6834                    e.fa_prefill_w(&q, &kf, &vf, &mut attn, hd, nh, nkv, t, t, scale, true,
6835                                   win)?;
6836                }
6837                return Ok(e.matmul(&fa.wo, &attn, t)?);
6838            }
6839        }
6840        if let Some(bucket) = dc_bucket {
6841            // DC arm (t == 1, under-window regime — the generate gate enforces it): ONE
6842            // fa_decode_dc over the live counter. len_d already advanced past this token
6843            // (t_kv = len_d[0], the cache.rs contract). KV-shared layers read the target's
6844            // counter (advanced when the target ran earlier in the stack).
6845            assert!(t == 1);
6846            // hd512 globals: eager (kvmod) picks the SCALAR unified below the fa512 floor,
6847            // and under the window every live t_kv sits below it — cap the capture bucket
6848            // under the floor so fa_decode_dc bakes the same scalar symbol, or the graph
6849            // rides the dpl16 twin and its numeric class diverges from the dc-eager stream
6850            // (E4B-GRAPH-GATE 2/64, 2026-07-12).
6851            let bucket = if hd == 512 && win <= crate::fa512_min_tkv() {
6852                bucket.min(crate::fa512_min_tkv().saturating_sub(1))
6853            } else { bucket };
6854            let k_view = e.view_u8(&kvl.k, kvl.k.len());
6855            let v_view = e.view_u8(&kvl.v, kvl.v.len());
6856            let g = (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on());
6857            // Weight prefetch (SOTA item 3, 2026-07-13): wo's decode plane prefetched into
6858            // L2 across the fa window — fa reads KV only, the weight DRAM lanes are idle
6859            // there (E4B valid window +0.65%: 196.8 vs 195.6). Value-free scheduling op,
6860            // captured into the dc graph like any other launch. Extending the cascade to
6861            // the ffn gate/up planes measured NEGATIVE (193.9 vs 195.8 — 29MB/layer floods
6862            // the fill path and evicts still-hot lines); wo-only is the shipped shape.
6863            // MEMRA_WPF=0 rollback seam.
6864            if crate::Engine::wpf_level() >= 1 {
6865                e.prefetch_weight_l2(&fa.wo)?;
6866            }
6867            // wave 5b: the combine emits the wo input q8 pair directly — the standalone
6868            // quantize launch + the f32 attn round-trip fold away (t=1 fast path only).
6869            if e.uses_q8_1_fast(&fa.wo) {
6870                let mut oq = e.alloc_i8_uninit(nh * hd)?;
6871                let mut od = e.zeros(nh * hd / 32)?;
6872                e.fa_decode_dc_q8(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6873                                  &kvl.len_d, bucket, scale,
6874                                  kvl.k_tok_bytes, kvl.v_tok_bytes, g,
6875                                  Some((&mut oq, &mut od)))?;
6876                return Ok(e.matmul_pre(&fa.wo, &oq, &od, &attn, t)?);
6877            }
6878            e.fa_decode_dc(&q, &k_view, &v_view, &mut attn, hd, nh, nkv,
6879                           &kvl.len_d, bucket, scale,
6880                           kvl.k_tok_bytes, kvl.v_tok_bytes, g)?;
6881            return Ok(e.matmul(&fa.wo, &attn, t)?);
6882        }
6883        for i in 0..t {
6884            let avail = base_len + i + 1;
6885            let (off_tok, t_kv) = if swa && avail > win { (avail - win, win) } else { (0, avail) };
6886            let k_view = e.view_u8_range(&kvl.k, off_tok * kvl.k_tok_bytes,
6887                                         (off_tok + t_kv) * kvl.k_tok_bytes);
6888            let v_view = e.view_u8_range(&kvl.v, off_tok * kvl.v_tok_bytes,
6889                                         (off_tok + t_kv) * kvl.v_tok_bytes);
6890            let qv = e.view(&q, t * nh * hd);
6891            let q_row = qv.slice(i * nh * hd..(i + 1) * nh * hd);
6892            let mut q_one = e.uninit(nh * hd)?;
6893            e.copy_view_into(&mut q_one, 0, &q_row, nh * hd)?;
6894            let mut a_one = e.uninit(nh * hd)?;
6895            // read class MUST match the append class (globals are e4m3 under gkv): the
6896            // swa-only flag decoded global-layer e4m3 bytes as q8_0/q5_1 — attention over
6897            // garbage on EVERY E4B global layer, the cross-mode maxdiff-30 root (2026-07-12).
6898            e.fa_decode_kvmod(&q_one, &k_view, &v_view, &mut a_one, hd, nh, nkv, t_kv, scale,
6899                        kvl.k_tok_bytes, kvl.v_tok_bytes,
6900                        (!swa && crate::Engine::gkv_on())
6901                            || (swa && crate::Engine::wkv_on()))?;
6902            e.copy_into(&mut attn, i * nh * hd, &a_one, nh * hd)?;
6903        }
6904        Ok(e.matmul(&fa.wo, &attn, t)?)
6905    }
6906
6907    /// E4B trunk: embed -> prologue -> layers (attn + dense ffn via the 31B tail_core + the
6908    /// per-layer-embedding tail + layer scale) -> output_norm. Returns (softcapped logits
6909    /// device [t, n_vocab], pre-output_norm hidden [t, n_embd]). Appends t rows per own-KV
6910    /// layer; does NOT advance cache.pos (caller owns pos).
6911    fn gemma4_e4b_trunk(&self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache,
6912                        head_last: bool)
6913                        -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6914        let n_embd = self.cfg.n_embd as usize;
6915        let t = tokens.len();
6916        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6917        let pos_d = e.htod_i32(&pos)?;
6918        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
6919        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
6920        let inp_pl = self.gemma4_e4b_inp_pl(e, tokens, &x, t)?;
6921        self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true, head_last)
6922    }
6923
6924    /// Layer stack + head over prebuilt (x_scaled, inp_pl, device pos) — everything below
6925    /// here is device-driven, so the dc arm shares it verbatim (stream identity with the
6926    /// eager chain by construction: SAME functions, not twins).
6927    fn gemma4_e4b_trunk_core(&self, e: &Engine, x_in: CudaSlice<f32>, inp_pl: CudaSlice<f32>,
6928                             pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache,
6929                             dc_bucket: Option<usize>, cap_logits: bool, head_last: bool)
6930                             -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6931        let n_embd = self.cfg.n_embd as usize;
6932        let eps = self.cfg.rms_eps;
6933        let n_layer = self.layers.len();
6934        let mut x = x_in;
6935        let aux_e4b = self.gemma4_aux.as_ref().unwrap().e4b.as_ref().unwrap();
6936        let n_epl = aux_e4b.n_epl;
6937
6938        // cross-layer fusion (2026-07-12 port of the 26B/31B trunk structure): each layer's
6939        // closing add_scale also EMITS the next layer's attn-normed input pre-quantized
6940        // q8_1 (add_scale_rms_norm_q8_1); the LAST layer emits through output_norm so the
6941        // head rides matmul_pre too. First layer's pair comes from a standalone fused
6942        // norm+quant.
6943        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
6944        for il in 0..n_layer {
6945            let layer = &self.layers[il];
6946            let (hq, hdq) = match h_carry.take() {
6947                Some(p) => p,
6948                None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
6949            };
6950            let o = self.gemma4_e4b_attn(e, il, &hq, &hdq, pos_d, t, cache, dc_bucket)?;
6951            // dense ffn tail with the post-attn norm FOLDED into its entry (one launch for
6952            // rms(o, post_attn_norm) + residual add + ffn_norm — glue-fusion lane).
6953            let bits = layer.gemma4.as_ref().unwrap();
6954            let e4b = bits.e4b.as_ref().expect("e4b layer bits");
6955            // glue wave 5: the tail DEFERS its post_ffw norm — the FFN exit fuses
6956            // rms(f0, post_ffw) + residual add + q8 emit into ONE launch (rms_pre_add_q8_1),
6957            // killing the rms_norm + add_q8_1 pair per layer. T-GENERIC since 2026-07-13:
6958            // the fused single-phase reduction is NOT FP-order-identical to the unfused
6959            // rms_norm+add pair (the original "bit-identical chain" claim was FALSE — the
6960            // t==1 gate left the batched VERIFY on the unfused chain and split E4B verify
6961            // from decode by logit maxdiff ~0.45, the greedy tie-flip at depth: gate 135/256.
6962            // NOFUSE bisect: disabling ONLY this fusion -> verify maxdiff 0.000e0). With the
6963            // gate dropped, decode AND verify ride the same fused chain — parity by
6964            // construction, VERIFY-GATE 0.000e0.
6965            let fuse_exit = e.uses_q8_1_fast(&e4b.inp_gate);
6966            let (sn, attn_out) = self.gemma4_layer_tail_core_pn(
6967                e, layer, &o, &x, t, Some(layer.post_attn_norm.float_data()), fuse_exit)?;
6968            let mut resid = e.uninit(t * n_embd)?;
6969            // per-layer-embedding tail: resid += rms(proj . (gelu(inp_gate . resid) * inp_pl[il]))
6970            // wave-2: the residual add emits q8_1 alongside — inp_gate rides matmul_pre.
6971            // (PLE one-block mega-fusion PROBED NEGATIVE 2026-07-13: argmax-correct but
6972            // 126 vs 189 tok/s — one SM pulling 0.74MB of weights loses to the multi-block
6973            // launch chain it replaced; jsonl row. Kernel deleted per doctrine.)
6974            let g = if fuse_exit {
6975                // sn here = RAW f0 (post_ffw deferred).
6976                let (rq, rd) = e.rms_pre_add_q8_1(&sn, bits.post_ffw_norm.float_data(),
6977                                                  &attn_out, &mut resid, n_embd, t,
6978                                                  self.cfg.rms_eps)?;
6979                e.matmul_pre(&e4b.inp_gate, &rq, &rd, &resid, t)?
6980            } else {
6981                e.add(&sn, &attn_out, &mut resid, t * n_embd)?;
6982                e.matmul(&e4b.inp_gate, &resid, t)?
6983            };
6984            let mut act = e.uninit(t * n_epl)?;
6985            let y = if t == 1 && e.uses_q8_1_fast(&e4b.proj) {
6986                let ipv = e.view(&inp_pl, n_epl * n_layer);
6987                let row = ipv.slice(il * n_epl..(il + 1) * n_epl);
6988                let (aq, ad) = e.gelu_tanh_mul_q8_1(&g, &row, &mut act, n_epl, 1)?;
6989                e.matmul_pre(&e4b.proj, &aq, &ad, &act, t)?
6990            } else {
6991                let mut inp_this = e.uninit(t * n_epl)?;
6992                e.copy_rows_strided(&inp_pl, &mut inp_this, n_epl, t, n_epl * n_layer,
6993                                    il * n_epl)?;
6994                e.gelu_tanh_mul(&g, &inp_this, &mut act, t * n_epl)?;
6995                e.matmul(&e4b.proj, &act, t)?
6996            };
6997            // rms(y, post_norm) + (yn + resid)*layer_scale + next-layer norm+quant emit,
6998            // ONE launch (glue-fusion lane; last layer emits through output_norm).
6999            let next_norm = if il + 1 < n_layer {
7000                self.layers[il + 1].attn_norm.float_data()
7001            } else {
7002                self.output_norm.float_data()
7003            };
7004            let mut xn = e.uninit(t * n_embd)?;
7005            let pair = e.rms_pre_add_scale_rms_norm_q8_1(&y, e4b.post_norm.float_data(),
7006                                                         &resid, bits.layer_scale, next_norm,
7007                                                         &mut xn, n_embd, t, eps)?;
7008            h_carry = Some(pair);
7009            x = xn;
7010        }
7011        // the head consumes the last layer's fused (output_norm) emit. head_last callers
7012        // (prime, last_only forward) need only the final row's logits — the all-T head is
7013        // t*n_vocab of discarded work (E4B prime: ~134ms GEMM + a 2.26GB dtoh kept 1 row).
7014        let (oq, odq) = h_carry.take().unwrap();
7015        let h0 = e.zeros(0)?;
7016        let hm = if head_last { 1 } else { t };
7017        let (hq, hd) = if head_last && t > 1 {
7018            let mut q1 = e.uninit_i8(n_embd)?;
7019            e.dtod_copy_view_i8(&oq.slice((t - 1) * n_embd..t * n_embd), &mut q1)?;
7020            let nb = n_embd / 32;
7021            let mut d1 = e.uninit(nb)?;
7022            e.dtod_copy_view(&odq.slice((t - 1) * nb..t * nb), &mut d1)?;
7023            (q1, d1)
7024        } else {
7025            (oq, odq)
7026        };
7027        let mut ld = e.matmul_pre(&self.output, &hq, &hd, &h0, hm)?;
7028        // softcap is strictly monotonic — greedy (argmax-only) consumers skip it, matching
7029        // the 26B/31B dc precedent (their dc head goes matmul -> argmax with no cap).
7030        // Logit-returning callers (host logits / spec prime) keep the capped emit.
7031        if cap_logits {
7032            let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
7033            e.softcap(&mut ld, cap, hm * self.output.out_features())?;
7034        }
7035        self.gemma4_suppress(e, &mut ld, hm)?;  // mask both capped and argmax-only consumers
7036        Ok((ld, x))
7037    }
7038
7039    /// E4B batched VERIFY (device tokens, the spec round's t=K+1 step): t rows through the
7040    /// e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared
7041    /// layers ride their targets), per-row device argmax + the POST-output_norm hidden
7042    /// stack (the drafter's h convention). Advances cache.pos/kvl.len by t — the spec
7043    /// round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind
7044    /// covers exactly the layers that appended).
7045    pub fn gemma4_e4b_decode_step_t_am_dev(&self, e: &Engine, tok_d: &CudaSlice<u32>,
7046                                                  t: usize, pos0: usize, cache: &mut Cache)
7047                                                  -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7048        let n_embd = self.cfg.n_embd as usize;
7049        let eps = self.cfg.rms_eps;
7050        let pos: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7051        let pos_d = e.htod_i32(&pos)?;
7052        let embd_gpu = self.embd_gpu.get_or_init(|| {
7053            e.upload_u8(&self.embd.raw).expect("embed table upload")
7054        });
7055        let (qt, rb) = self.embd.qt_and_row_bytes(n_embd);
7056        let mut x = e.embed_gather_device_td(embd_gpu, tok_d, t, n_embd, qt, rb)?;
7057        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), t * n_embd)?;
7058        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, tok_d, &x, t)?;
7059        let (ld, xp) = self.gemma4_e4b_trunk_core(e, x, inp_pl, &pos_d, t, cache, None, true,
7060                                                  false)?;
7061        // softcap is monotonic — the per-row argmax is invariant to it (the trunk's head
7062        // emit is already capped, matching the eager chain bit-for-bit).
7063        let n_vocab = self.output.out_features();
7064        let mut vam = e.stream().alloc_zeros::<u32>(t)?;
7065        for i in 0..t {
7066            e.argmax_token_device_col(&ld, i, n_vocab, &mut vam, i)?;
7067        }
7068        let mut hn = e.uninit(t * n_embd)?;
7069        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7070        cache.pos += t;
7071        Ok((vam, hn))
7072    }
7073
7074    /// E4B verify + host logits + POST-output_norm hidden stack (the short-prompt spec
7075    /// prime path — mirror of `gemma4_decode_step_t_h`).
7076    pub(crate) fn gemma4_e4b_decode_step_t_h(&self, e: &Engine, tokens: &[u32], pos0: usize,
7077                                             cache: &mut Cache)
7078                                             -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7079        let n_embd = self.cfg.n_embd as usize;
7080        let eps = self.cfg.rms_eps;
7081        let t = tokens.len();
7082        let (ld, xp) = self.gemma4_e4b_trunk(e, tokens, pos0, cache, false)?;
7083        let mut hn = e.uninit(t * n_embd)?;
7084        e.rms_norm(&xp, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7085        cache.pos += t;
7086        Ok((e.dtoh(&ld)?, hn))
7087    }
7088
7089    /// E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN
7090    /// PLACE (self-feeding replay) and every launch arg is a device counter — pos from
7091    /// pos_d (inc'd in-stream), KV slots from len_d (advanced in-stream), attention from
7092    /// fa_decode_dc at `bucket`. Host mirrors (cache.pos / kvl.len) advance in the caller's
7093    /// replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).
7094    pub fn gemma4_e4b_decode_step_dcg(&self, e: &Engine, token_d: &mut CudaSlice<u32>,
7095                                      pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7096                                      embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7097                                      n_vocab: usize, bucket: usize)
7098                                      -> Result<(), Box<dyn std::error::Error>> {
7099        let n_embd = self.cfg.n_embd as usize;
7100        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7101        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7102        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7103        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, Some(bucket),
7104                                                  false, false)?;
7105        e.argmax_token_device_into(&ld, token_d, n_vocab)?;
7106        e.inc_seqlen(pos_d)?;
7107        Ok(())
7108    }
7109
7110    /// E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides `token_d`, the
7111    /// greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer
7112    /// stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream
7113    /// identity by construction, not by twin-kernel parity). Host KV mirrors advance like the
7114    /// 26B dc-eager arm (window views are host math); len_d stays synced by the caller's
7115    /// entry sync + the appends here don't read it. Graph capture is NOT wired (no
7116    /// cap_bucket_max) — the E4B graph arc comes after the perf gates.
7117    #[allow(clippy::too_many_arguments)]
7118    pub fn gemma4_e4b_decode_step_dc(&self, e: &Engine, token_d: &CudaSlice<u32>,
7119                                     pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>,
7120                                     embd_qt: i32, embd_rb: usize, cache: &mut Cache,
7121                                     n_vocab: usize)
7122                                     -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7123        let n_embd = self.cfg.n_embd as usize;
7124        let eps = self.cfg.rms_eps;
7125        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
7126        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), n_embd)?;
7127        let inp_pl = self.gemma4_e4b_inp_pl_dev(e, token_d, &x, 1)?;
7128        let (ld, _x) = self.gemma4_e4b_trunk_core(e, x, inp_pl, pos_d, 1, cache, None, false,
7129                                                  false)?;
7130        let mut tok_out = e.stream().alloc_zeros::<u32>(1)?;
7131        e.argmax_token_device_into(&ld, &mut tok_out, n_vocab)?;
7132        e.inc_seqlen(pos_d)?;
7133        cache.pos += 1;
7134        let _ = eps;
7135        Ok(tok_out)
7136    }
7137
7138    /// E4B eager T=1 decode step (decode_step_h contract): returns (softcapped logits host,
7139    /// pre-output_norm hidden). Advances cache.pos.
7140    pub(crate) fn gemma4_e4b_decode_step_h(&self, e: &Engine, token: u32, cache: &mut Cache)
7141                                           -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7142        let (ld, x) = self.gemma4_e4b_trunk(e, &[token], cache.pos, cache, false)?;
7143        let logits = e.dtoh(&ld)?;
7144        cache.pos += 1;
7145        Ok((logits, x))
7146    }
7147
7148    /// E4B batched prime (prime_cache contract: (last-row logits host, h_seed pre-norm last
7149    /// row, hidden stack)). First-light: the t-wide trunk (per-row attention) — correct, not
7150    /// fast; the prefill fa arms come later.
7151    pub(crate) fn gemma4_e4b_prime(&self, e: &Engine, tokens: &[u32], cache: &mut Cache)
7152                                   -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7153        assert_eq!(cache.pos, 0, "e4b prime is fresh-prompt only (v0)");
7154        let n_embd = self.cfg.n_embd as usize;
7155        let t = tokens.len();
7156        let (ld, x) = self.gemma4_e4b_trunk(e, tokens, 0, cache, true)?;
7157        cache.pos += t;
7158        let last = e.dtoh(&ld)?;   // head_last: ld is already the final row only
7159        let xv = e.view(&x, t * n_embd);
7160        let row = xv.slice((t - 1) * n_embd..t * n_embd);
7161        let mut h_seed = e.uninit(n_embd)?;
7162        e.copy_view_into(&mut h_seed, 0, &row, n_embd)?;
7163        Ok((last, h_seed, x))
7164    }
7165
7166    /// E4B prefill logits (forward contract — no persistent cache; scratch cache internally).
7167    pub(crate) fn gemma4_e4b_forward(&self, e: &Engine, tokens: &[u32], last_only: bool)
7168                                     -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7169        let mut cache = Cache::new(e, &self.cfg, tokens.len() + 8)?;
7170        let (ld, _x) = self.gemma4_e4b_trunk(e, tokens, 0, &mut cache, last_only)?;
7171        Ok(e.dtoh(&ld)?)   // head_last already reduced to the final row when last_only
7172    }
7173}
7174
7175#[cfg(test)]
7176mod page_prefetch_tests {
7177    use super::{
7178        grouped_worker_prefetch_position, page_prefetch_positions,
7179        page_prefetch_window_from_values, worker_prefetch_positions,
7180    };
7181
7182    #[test]
7183    fn page_prefetch_window_keeps_existing_opt_in_default() {
7184        assert_eq!(page_prefetch_window_from_values(false, None), 0);
7185        assert_eq!(page_prefetch_window_from_values(false, Some("8")), 0);
7186        assert_eq!(page_prefetch_window_from_values(true, None), 1);
7187        assert_eq!(page_prefetch_window_from_values(true, Some("bad")), 1);
7188        assert_eq!(page_prefetch_window_from_values(true, Some("0")), 0);
7189        assert_eq!(page_prefetch_window_from_values(true, Some("8")), 8);
7190    }
7191
7192    #[test]
7193    fn rolling_page_prefetch_advises_each_future_expert_once() {
7194        let advised: Vec<_> = (0..7)
7195            .flat_map(|position| page_prefetch_positions(position, 7, 3))
7196            .collect();
7197        assert_eq!(advised, vec![1, 2, 3, 4, 5, 6]);
7198
7199        let one_ahead: Vec<_> = (0..4)
7200            .flat_map(|position| page_prefetch_positions(position, 4, 1))
7201            .collect();
7202        assert_eq!(one_ahead, vec![1, 2, 3]);
7203        assert!(page_prefetch_positions(0, 4, 0).is_empty());
7204    }
7205
7206    #[test]
7207    fn grouped_worker_prefetch_primes_first_then_each_known_next_once() {
7208        assert_eq!(grouped_worker_prefetch_position(0, None), None);
7209        let positions: Vec<_> = std::iter::once(grouped_worker_prefetch_position(4, None).unwrap())
7210            .chain((0..4).filter_map(|position| {
7211                grouped_worker_prefetch_position(4, Some(position))
7212            }))
7213            .collect();
7214        assert_eq!(positions, vec![0, 1, 2, 3]);
7215        assert_eq!(grouped_worker_prefetch_position(1, Some(0)), None);
7216    }
7217
7218    #[test]
7219    fn rolling_worker_prefetch_primes_current_and_each_future_expert_once() {
7220        let queued: Vec<_> = (0..8)
7221            .flat_map(|position| worker_prefetch_positions(position, 8, 5))
7222            .collect();
7223        assert_eq!(queued, (0..8).collect::<Vec<_>>());
7224
7225        let one_at_a_time: Vec<_> = (0..4)
7226            .flat_map(|position| worker_prefetch_positions(position, 4, 1))
7227            .collect();
7228        assert_eq!(one_at_a_time, vec![0, 1, 2, 3]);
7229        assert!(worker_prefetch_positions(0, 4, 0).is_empty());
7230    }
7231}
7232
7233pub struct G4DcSlots {
7234    x: CudaSlice<f32>, xn: CudaSlice<f32>, cur: CudaSlice<f32>,
7235    hq: CudaSlice<i8>, hd_: CudaSlice<f32>,
7236    q0: CudaSlice<f32>, k0: CudaSlice<f32>, v0: CudaSlice<f32>,
7237    q: CudaSlice<f32>, k: CudaSlice<f32>, v: CudaSlice<f32>,
7238    attn: CudaSlice<f32>, o: CudaSlice<f32>,
7239    attn_out: CudaSlice<f32>, zsh: CudaSlice<f32>,
7240    zq: CudaSlice<i8>, zd: CudaSlice<f32>,
7241    gate: CudaSlice<f32>, up: CudaSlice<f32>,
7242    act: CudaSlice<f32>, actq: CudaSlice<i8>, actd: CudaSlice<f32>,
7243    f0: CudaSlice<f32>, sn: CudaSlice<f32>,
7244    hn: CudaSlice<f32>, logits: CudaSlice<f32>,
7245}
7246